wasm linker: handle function data references properly

3 files changed, 52 insertions, 15 deletions

diff --git a/src/codegen.zig b/src/codegen.zig
index e133ee9937..b38844e4e2 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -738,19 +738,32 @@ fn lowerNavRef(
             dev.check(link.File.Tag.wasm.devFeature());
             const wasm = lf.cast(.wasm).?;
             assert(reloc_parent == .none);
-            if (is_obj) {
-                try wasm.out_relocs.append(gpa, .{
-                    .offset = @intCast(code.items.len),
-                    .pointee = .{ .symbol_index = try wasm.navSymbolIndex(nav_index) },
-                    .tag = if (ptr_width_bytes == 4) .memory_addr_i32 else .memory_addr_i64,
-                    .addend = @intCast(offset),
-                });
+            if (is_fn_body) {
+                const gop = try wasm.zcu_indirect_function_set.getOrPut(gpa, nav_index);
+                if (!gop.found_existing) gop.value_ptr.* = {};
+                if (is_obj) {
+                    @panic("TODO add out_reloc for this");
+                } else {
+                    try wasm.func_table_fixups.append(gpa, .{
+                        .table_index = @enumFromInt(gop.index),
+                        .offset = @intCast(code.items.len),
+                    });
+                }
             } else {
-                try wasm.nav_fixups.ensureUnusedCapacity(gpa, 1);
-                wasm.nav_fixups.appendAssumeCapacity(.{
-                    .navs_exe_index = try wasm.refNavExe(nav_index),
-                    .offset = @intCast(code.items.len),
-                });
+                if (is_obj) {
+                    try wasm.out_relocs.append(gpa, .{
+                        .offset = @intCast(code.items.len),
+                        .pointee = .{ .symbol_index = try wasm.navSymbolIndex(nav_index) },
+                        .tag = if (ptr_width_bytes == 4) .memory_addr_i32 else .memory_addr_i64,
+                        .addend = @intCast(offset),
+                    });
+                } else {
+                    try wasm.nav_fixups.ensureUnusedCapacity(gpa, 1);
+                    wasm.nav_fixups.appendAssumeCapacity(.{
+                        .navs_exe_index = try wasm.refNavExe(nav_index),
+                        .offset = @intCast(code.items.len),
+                    });
+                }
             }
             code.appendNTimesAssumeCapacity(0, ptr_width_bytes);
             return;
diff --git a/src/link/Wasm.zig b/src/link/Wasm.zig
index 5dd51c2f12..1ec1064ee0 100644
--- a/src/link/Wasm.zig
+++ b/src/link/Wasm.zig
@@ -151,7 +151,11 @@ uav_fixups: std.ArrayListUnmanaged(UavFixup) = .empty,
 /// List of locations within `string_bytes` that must be patched with the virtual
 /// memory address of a Nav during `flush`.
 /// When emitting an object file, `out_relocs` is used instead.
+/// No functions here only global variables.
 nav_fixups: std.ArrayListUnmanaged(NavFixup) = .empty,
+/// When a nav reference is a function pointer, this tracks the required function
+/// table entry index that needs to overwrite the code in the final output.
+func_table_fixups: std.ArrayListUnmanaged(FuncTableFixup) = .empty,
 /// Symbols to be emitted into an object file. Remains empty when not emitting
 /// an object file.
 symbol_table: std.AutoArrayHashMapUnmanaged(String, void) = .empty,
@@ -307,6 +311,12 @@ pub const NavFixup = extern struct {
     offset: u32,
 };
 
+pub const FuncTableFixup = extern struct {
+    table_index: ZcuIndirectFunctionSetIndex,
+    /// Index into `string_bytes`.
+    offset: u32,
+};
+
 /// Index into `objects`.
 pub const ObjectIndex = enum(u32) {
     _,
@@ -2208,7 +2218,7 @@ pub const FunctionImportId = enum(u32) {
     pub fn pack(unpacked: Unpacked, wasm: *const Wasm) FunctionImportId {
         return switch (unpacked) {
             .object_function_import => |i| @enumFromInt(@intFromEnum(i)),
-            .zcu_import => |i| @enumFromInt(@intFromEnum(i) - wasm.object_function_imports.entries.len),
+            .zcu_import => |i| @enumFromInt(@intFromEnum(i) + wasm.object_function_imports.entries.len),
         };
     }
 
@@ -2295,7 +2305,7 @@ pub const GlobalImportId = enum(u32) {
     pub fn pack(unpacked: Unpacked, wasm: *const Wasm) GlobalImportId {
         return switch (unpacked) {
             .object_global_import => |i| @enumFromInt(@intFromEnum(i)),
-            .zcu_import => |i| @enumFromInt(@intFromEnum(i) - wasm.object_global_imports.entries.len),
+            .zcu_import => |i| @enumFromInt(@intFromEnum(i) + wasm.object_global_imports.entries.len),
         };
     }
 
@@ -2360,7 +2370,7 @@ pub const DataImportId = enum(u32) {
     pub fn pack(unpacked: Unpacked, wasm: *const Wasm) DataImportId {
         return switch (unpacked) {
             .object_data_import => |i| @enumFromInt(@intFromEnum(i)),
-            .zcu_import => |i| @enumFromInt(@intFromEnum(i) - wasm.object_data_imports.entries.len),
+            .zcu_import => |i| @enumFromInt(@intFromEnum(i) + wasm.object_data_imports.entries.len),
         };
     }
 
@@ -3027,6 +3037,7 @@ pub fn deinit(wasm: *Wasm) void {
     wasm.out_relocs.deinit(gpa);
     wasm.uav_fixups.deinit(gpa);
     wasm.nav_fixups.deinit(gpa);
+    wasm.func_table_fixups.deinit(gpa);
 
     wasm.zcu_indirect_function_set.deinit(gpa);
     wasm.object_indirect_function_import_set.deinit(gpa);
diff --git a/src/link/Wasm/Flush.zig b/src/link/Wasm/Flush.zig
index 9e74328c46..802ee6e0e2 100644
--- a/src/link/Wasm/Flush.zig
+++ b/src/link/Wasm/Flush.zig
@@ -54,6 +54,11 @@ const IndirectFunctionTableIndex = enum(u32) {
     fn fromOutputFunctionIndex(f: *const Flush, i: Wasm.OutputFunctionIndex) IndirectFunctionTableIndex {
         return @enumFromInt(f.indirect_function_table.getIndex(i).?);
     }
+
+    fn fromZcuIndirectFunctionSetIndex(i: Wasm.ZcuIndirectFunctionSetIndex) IndirectFunctionTableIndex {
+        // These are the same since those are added to the table first.
+        return @enumFromInt(@intFromEnum(i));
+    }
 };
 
 const DataSegmentGroup = struct {
@@ -755,6 +760,14 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
                 mem.writeInt(u64, wasm.string_bytes.items[nav_fixup.offset..][0..8], vaddr, .little);
             }
         }
+        for (wasm.func_table_fixups.items) |fixup| {
+            const table_index: IndirectFunctionTableIndex = .fromZcuIndirectFunctionSetIndex(fixup.table_index);
+            if (!is64) {
+                mem.writeInt(u32, wasm.string_bytes.items[fixup.offset..][0..4], @intFromEnum(table_index), .little);
+            } else {
+                mem.writeInt(u64, wasm.string_bytes.items[fixup.offset..][0..8], @intFromEnum(table_index), .little);
+            }
+        }
     }
 
     // Data section.