implement indirect function table for object functions

4 files changed, 91 insertions, 27 deletions

diff --git a/src/arch/wasm/CodeGen.zig b/src/arch/wasm/CodeGen.zig
index 8d70103d1f..4023658742 100644
--- a/src/arch/wasm/CodeGen.zig
+++ b/src/arch/wasm/CodeGen.zig
@@ -1027,7 +1027,7 @@ fn emitWValue(cg: *CodeGen, value: WValue) InnerError!void {
             const ip = &zcu.intern_pool;
             if (ip.getNav(nav_ref.nav_index).isExternOrFn(ip)) {
                 assert(nav_ref.offset == 0);
-                const gop = try wasm.indirect_function_table.getOrPut(comp.gpa, nav_ref.nav_index);
+                const gop = try wasm.zcu_indirect_function_set.getOrPut(comp.gpa, nav_ref.nav_index);
                 if (!gop.found_existing) gop.value_ptr.* = {};
                 try cg.addInst(.{
                     .tag = .func_ref,
@@ -1056,7 +1056,7 @@ fn emitWValue(cg: *CodeGen, value: WValue) InnerError!void {
             if (ip.isFunctionType(ip.typeOf(uav.ip_index))) {
                 assert(uav.offset == 0);
                 const owner_nav = ip.toFunc(uav.ip_index).owner_nav;
-                const gop = try wasm.indirect_function_table.getOrPut(comp.gpa, owner_nav);
+                const gop = try wasm.zcu_indirect_function_set.getOrPut(comp.gpa, owner_nav);
                 if (!gop.found_existing) gop.value_ptr.* = {};
                 try cg.addInst(.{
                     .tag = .func_ref,
diff --git a/src/arch/wasm/Mir.zig b/src/arch/wasm/Mir.zig
index 8d22d99108..3ff6b8072b 100644
--- a/src/arch/wasm/Mir.zig
+++ b/src/arch/wasm/Mir.zig
@@ -615,7 +615,7 @@ pub const Inst = struct {
         intrinsic: Intrinsic,
         uav_obj: Wasm.UavsObjIndex,
         uav_exe: Wasm.UavsExeIndex,
-        indirect_function_table_index: Wasm.IndirectFunctionTableIndex,
+        indirect_function_table_index: Wasm.ZcuIndirectFunctionSetIndex,
 
         comptime {
             switch (builtin.mode) {
diff --git a/src/link/Wasm.zig b/src/link/Wasm.zig
index 5f836955d2..277d092cba 100644
--- a/src/link/Wasm.zig
+++ b/src/link/Wasm.zig
@@ -260,7 +260,9 @@ table_imports: std.AutoArrayHashMapUnmanaged(String, TableImport.Index) = .empty
 
 /// All functions that have had their address taken and therefore might be
 /// called via a `call_indirect` function.
-indirect_function_table: std.AutoArrayHashMapUnmanaged(InternPool.Nav.Index, void) = .empty,
+zcu_indirect_function_set: std.AutoArrayHashMapUnmanaged(InternPool.Nav.Index, void) = .empty,
+object_indirect_function_import_set: std.AutoArrayHashMapUnmanaged(String, void) = .empty,
+object_indirect_function_set: std.AutoArrayHashMapUnmanaged(ObjectFunctionIndex, void) = .empty,
 
 error_name_table_ref_count: u32 = 0,
 
@@ -288,8 +290,8 @@ error_name_bytes: std.ArrayListUnmanaged(u8) = .empty,
 /// is stored. No need to serialize; trivially reconstructed.
 error_name_offs: std.ArrayListUnmanaged(u32) = .empty,
 
-/// Index into `Wasm.indirect_function_table`.
-pub const IndirectFunctionTableIndex = enum(u32) {
+/// Index into `Wasm.zcu_indirect_function_set`.
+pub const ZcuIndirectFunctionSetIndex = enum(u32) {
     _,
 };
 
@@ -1312,8 +1314,8 @@ pub const TableImport = extern struct {
                 .unresolved => unreachable,
                 .__indirect_function_table => .{
                     .flags = .{ .has_max = true, .is_shared = false },
-                    .min = @intCast(wasm.indirect_function_table.entries.len + 1),
-                    .max = @intCast(wasm.indirect_function_table.entries.len + 1),
+                    .min = @intCast(wasm.flush_buffer.indirect_function_table.entries.len + 1),
+                    .max = @intCast(wasm.flush_buffer.indirect_function_table.entries.len + 1),
                 },
                 .object_table => |i| i.ptr(wasm).limits(),
             };
@@ -3025,7 +3027,10 @@ pub fn deinit(wasm: *Wasm) void {
     wasm.out_relocs.deinit(gpa);
     wasm.uav_fixups.deinit(gpa);
     wasm.nav_fixups.deinit(gpa);
-    wasm.indirect_function_table.deinit(gpa);
+
+    wasm.zcu_indirect_function_set.deinit(gpa);
+    wasm.object_indirect_function_import_set.deinit(gpa);
+    wasm.object_indirect_function_set.deinit(gpa);
 
     wasm.string_bytes.deinit(gpa);
     wasm.string_table.deinit(gpa);
@@ -3515,6 +3520,7 @@ fn markDataImport(
 }
 
 fn markRelocations(wasm: *Wasm, relocs: ObjectRelocation.IterableSlice) link.File.FlushError!void {
+    const gpa = wasm.base.comp.gpa;
     for (relocs.slice.tags(wasm), relocs.slice.pointees(wasm), relocs.slice.offsets(wasm)) |tag, pointee, offset| {
         if (offset >= relocs.end) break;
         switch (tag) {
@@ -3522,6 +3528,11 @@ fn markRelocations(wasm: *Wasm, relocs: ObjectRelocation.IterableSlice) link.Fil
             .function_import_index_i32,
             .function_import_offset_i32,
             .function_import_offset_i64,
+            => {
+                const name = pointee.symbol_name;
+                const i: FunctionImport.Index = @enumFromInt(wasm.object_function_imports.getIndex(name).?);
+                try markFunctionImport(wasm, name, i.value(wasm), i);
+            },
             .table_import_index_sleb,
             .table_import_index_i32,
             .table_import_index_sleb64,
@@ -3530,6 +3541,7 @@ fn markRelocations(wasm: *Wasm, relocs: ObjectRelocation.IterableSlice) link.Fil
             .table_import_index_rel_sleb64,
             => {
                 const name = pointee.symbol_name;
+                try wasm.object_indirect_function_import_set.put(gpa, name, {});
                 const i: FunctionImport.Index = @enumFromInt(wasm.object_function_imports.getIndex(name).?);
                 try markFunctionImport(wasm, name, i.value(wasm), i);
             },
@@ -3564,13 +3576,18 @@ fn markRelocations(wasm: *Wasm, relocs: ObjectRelocation.IterableSlice) link.Fil
             .function_index_i32,
             .function_offset_i32,
             .function_offset_i64,
+            => try markFunction(wasm, pointee.function.chaseWeak(wasm)),
             .table_index_sleb,
             .table_index_i32,
             .table_index_sleb64,
             .table_index_i64,
             .table_index_rel_sleb,
             .table_index_rel_sleb64,
-            => try markFunction(wasm, pointee.function.chaseWeak(wasm)),
+            => {
+                const function = pointee.function;
+                try wasm.object_indirect_function_set.put(gpa, function, {});
+                try markFunction(wasm, function.chaseWeak(wasm));
+            },
             .global_index_leb,
             .global_index_i32,
             => try markGlobal(wasm, pointee.global.chaseWeak(wasm)),
diff --git a/src/link/Wasm/Flush.zig b/src/link/Wasm/Flush.zig
index 281f915b31..d53a508f9b 100644
--- a/src/link/Wasm/Flush.zig
+++ b/src/link/Wasm/Flush.zig
@@ -34,9 +34,28 @@ function_imports: std.AutoArrayHashMapUnmanaged(String, Wasm.FunctionImportId) =
 global_imports: std.AutoArrayHashMapUnmanaged(String, Wasm.GlobalImportId) = .empty,
 data_imports: std.AutoArrayHashMapUnmanaged(String, Wasm.DataImportId) = .empty,
 
+indirect_function_table: std.AutoArrayHashMapUnmanaged(Wasm.OutputFunctionIndex, void) = .empty,
+
 /// For debug purposes only.
 memory_layout_finished: bool = false,
 
+/// Index into `indirect_function_table`.
+const IndirectFunctionTableIndex = enum(u32) {
+    _,
+
+    fn fromObjectFunctionHandlingWeak(wasm: *const Wasm, index: Wasm.ObjectFunctionIndex) IndirectFunctionTableIndex {
+        return fromOutputFunctionIndex(&wasm.flush_buffer, .fromObjectFunctionHandlingWeak(wasm, index));
+    }
+
+    fn fromSymbolName(wasm: *const Wasm, name: String) IndirectFunctionTableIndex {
+        return fromOutputFunctionIndex(&wasm.flush_buffer, .fromSymbolName(wasm, name));
+    }
+
+    fn fromOutputFunctionIndex(f: *const Flush, i: Wasm.OutputFunctionIndex) IndirectFunctionTableIndex {
+        return @enumFromInt(f.indirect_function_table.getIndex(i).?);
+    }
+};
+
 const DataSegmentGroup = struct {
     first_segment: Wasm.DataSegmentId,
     end_addr: u32,
@@ -46,6 +65,7 @@ pub fn clear(f: *Flush) void {
     f.data_segments.clearRetainingCapacity();
     f.data_segment_groups.clearRetainingCapacity();
     f.binary_bytes.clearRetainingCapacity();
+    f.indirect_function_table.clearRetainingCapacity();
     f.memory_layout_finished = false;
 }
 
@@ -57,6 +77,7 @@ pub fn deinit(f: *Flush, gpa: Allocator) void {
     f.function_imports.deinit(gpa);
     f.global_imports.deinit(gpa);
     f.data_imports.deinit(gpa);
+    f.indirect_function_table.deinit(gpa);
     f.* = undefined;
 }
 
@@ -156,6 +177,17 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
 
     if (diags.hasErrors()) return error.LinkFailure;
 
+    // Merge indirect function tables.
+    try f.indirect_function_table.ensureUnusedCapacity(gpa, wasm.zcu_indirect_function_set.entries.len +
+        wasm.object_indirect_function_import_set.entries.len + wasm.object_indirect_function_set.entries.len);
+    // This one goes first so the indexes can be stable for MIR lowering.
+    for (wasm.zcu_indirect_function_set.keys()) |nav_index|
+        f.indirect_function_table.putAssumeCapacity(.fromIpNav(wasm, nav_index), {});
+    for (wasm.object_indirect_function_import_set.keys()) |symbol_name|
+        f.indirect_function_table.putAssumeCapacity(.fromSymbolName(wasm, symbol_name), {});
+    for (wasm.object_indirect_function_set.keys()) |object_function_index|
+        f.indirect_function_table.putAssumeCapacity(.fromObjectFunction(wasm, object_function_index), {});
+
     // TODO only include init functions for objects with must_link=true or
     // which have any alive functions inside them.
     if (wasm.object_init_funcs.items.len > 0) {
@@ -213,7 +245,7 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
 
     try wasm.tables.ensureUnusedCapacity(gpa, 1);
 
-    if (wasm.indirect_function_table.entries.len > 0) {
+    if (f.indirect_function_table.entries.len > 0) {
         wasm.tables.putAssumeCapacity(.__indirect_function_table, {});
     }
 
@@ -634,7 +666,7 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
     }
 
     // element section
-    if (wasm.indirect_function_table.entries.len > 0) {
+    if (f.indirect_function_table.entries.len > 0) {
         const header_offset = try reserveVecSectionHeader(gpa, binary_bytes);
 
         // indirect function table elements
@@ -650,9 +682,8 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
         if (flags == 0x02) {
             try leb.writeUleb128(binary_writer, @as(u8, 0)); // represents funcref
         }
-        try leb.writeUleb128(binary_writer, @as(u32, @intCast(wasm.indirect_function_table.entries.len)));
-        for (wasm.indirect_function_table.keys()) |nav_index| {
-            const func_index: Wasm.OutputFunctionIndex = .fromIpNav(wasm, nav_index);
+        try leb.writeUleb128(binary_writer, @as(u32, @intCast(f.indirect_function_table.entries.len)));
+        for (f.indirect_function_table.keys()) |func_index| {
             try leb.writeUleb128(binary_writer, @intFromEnum(func_index));
         }
 
@@ -1459,23 +1490,23 @@ fn applyRelocs(code: []u8, code_offset: u32, relocs: Wasm.ObjectRelocation.Itera
             .function_index_leb => reloc_leb_function(sliced_code, .fromObjectFunctionHandlingWeak(wasm, pointee.function)),
             .function_offset_i32 => @panic("TODO this value is not known yet"),
             .function_offset_i64 => @panic("TODO this value is not known yet"),
-            .table_index_i32 => @panic("TODO indirect function table needs to support object functions too"),
-            .table_index_i64 => @panic("TODO indirect function table needs to support object functions too"),
-            .table_index_rel_sleb => @panic("TODO indirect function table needs to support object functions too"),
-            .table_index_rel_sleb64 => @panic("TODO indirect function table needs to support object functions too"),
-            .table_index_sleb => @panic("TODO indirect function table needs to support object functions too"),
-            .table_index_sleb64 => @panic("TODO indirect function table needs to support object functions too"),
+            .table_index_i32 => reloc_u32_table_index(sliced_code, .fromObjectFunctionHandlingWeak(wasm, pointee.function)),
+            .table_index_i64 => reloc_u64_table_index(sliced_code, .fromObjectFunctionHandlingWeak(wasm, pointee.function)),
+            .table_index_rel_sleb => @panic("TODO what does this reloc tag mean?"),
+            .table_index_rel_sleb64 => @panic("TODO what does this reloc tag mean?"),
+            .table_index_sleb => reloc_sleb_table_index(sliced_code, .fromObjectFunctionHandlingWeak(wasm, pointee.function)),
+            .table_index_sleb64 => reloc_sleb64_table_index(sliced_code, .fromObjectFunctionHandlingWeak(wasm, pointee.function)),
 
             .function_import_index_i32 => reloc_u32_function(sliced_code, .fromSymbolName(wasm, pointee.symbol_name)),
             .function_import_index_leb => reloc_leb_function(sliced_code, .fromSymbolName(wasm, pointee.symbol_name)),
             .function_import_offset_i32 => @panic("TODO this value is not known yet"),
             .function_import_offset_i64 => @panic("TODO this value is not known yet"),
-            .table_import_index_i32 => @panic("TODO indirect function table needs to support object functions too"),
-            .table_import_index_i64 => @panic("TODO indirect function table needs to support object functions too"),
-            .table_import_index_rel_sleb => @panic("TODO indirect function table needs to support object functions too"),
-            .table_import_index_rel_sleb64 => @panic("TODO indirect function table needs to support object functions too"),
-            .table_import_index_sleb => @panic("TODO indirect function table needs to support object functions too"),
-            .table_import_index_sleb64 => @panic("TODO indirect function table needs to support object functions too"),
+            .table_import_index_i32 => reloc_u32_table_index(sliced_code, .fromSymbolName(wasm, pointee.symbol_name)),
+            .table_import_index_i64 => reloc_u64_table_index(sliced_code, .fromSymbolName(wasm, pointee.symbol_name)),
+            .table_import_index_rel_sleb => @panic("TODO what does this reloc tag mean?"),
+            .table_import_index_rel_sleb64 => @panic("TODO what does this reloc tag mean?"),
+            .table_import_index_sleb => reloc_sleb_table_index(sliced_code, .fromSymbolName(wasm, pointee.symbol_name)),
+            .table_import_index_sleb64 => reloc_sleb64_table_index(sliced_code, .fromSymbolName(wasm, pointee.symbol_name)),
 
             .global_index_i32 => reloc_u32_global(sliced_code, .fromObjectGlobalHandlingWeak(wasm, pointee.global)),
             .global_index_leb => reloc_leb_global(sliced_code, .fromObjectGlobalHandlingWeak(wasm, pointee.global)),
@@ -1517,6 +1548,22 @@ fn applyRelocs(code: []u8, code_offset: u32, relocs: Wasm.ObjectRelocation.Itera
     }
 }
 
+fn reloc_u32_table_index(code: []u8, i: IndirectFunctionTableIndex) void {
+    mem.writeInt(u32, code[0..4], @intFromEnum(i), .little);
+}
+
+fn reloc_u64_table_index(code: []u8, i: IndirectFunctionTableIndex) void {
+    mem.writeInt(u64, code[0..8], @intFromEnum(i), .little);
+}
+
+fn reloc_sleb_table_index(code: []u8, i: IndirectFunctionTableIndex) void {
+    leb.writeSignedFixed(5, code[0..5], @intFromEnum(i));
+}
+
+fn reloc_sleb64_table_index(code: []u8, i: IndirectFunctionTableIndex) void {
+    leb.writeSignedFixed(11, code[0..11], @intFromEnum(i));
+}
+
 fn reloc_u32_function(code: []u8, function: Wasm.OutputFunctionIndex) void {
     mem.writeInt(u32, code[0..4], @intFromEnum(function), .little);
 }