wasm linker: ability to get data and functions from objects

2 files changed, 106 insertions, 42 deletions

diff --git a/src/link/Wasm.zig b/src/link/Wasm.zig
index f27aacb138..0339122796 100644
--- a/src/link/Wasm.zig
+++ b/src/link/Wasm.zig
@@ -180,8 +180,10 @@ dump_argv_list: std.ArrayListUnmanaged([]const u8),
 preloaded_strings: PreloadedStrings,
 
 /// This field is used when emitting an object; `navs_exe` used otherwise.
+/// Does not include externs since that data lives elsewhere.
 navs_obj: std.AutoArrayHashMapUnmanaged(InternPool.Nav.Index, ZcuDataObj) = .empty,
 /// This field is unused when emitting an object; `navs_obj` used otherwise.
+/// Does not include externs since that data lives elsewhere.
 navs_exe: std.AutoArrayHashMapUnmanaged(InternPool.Nav.Index, ZcuDataExe) = .empty,
 /// Tracks all InternPool values referenced by codegen. Needed for outputting
 /// the data segment. This one does not track ref count because object files
@@ -221,6 +223,9 @@ functions: std.AutoArrayHashMapUnmanaged(FunctionImport.Resolution, void) = .emp
 /// Tracks the value at the end of prelink, at which point `functions`
 /// contains only object file functions, and nothing from the Zcu yet.
 functions_end_prelink: u32 = 0,
+
+function_imports_len_prelink: u32 = 0,
+data_imports_len_prelink: u32 = 0,
 /// At the end of prelink, this is populated with needed functions from
 /// objects.
 ///
@@ -3447,6 +3452,9 @@ pub fn prelink(wasm: *Wasm, prog_node: std.Progress.Node) link.File.FlushError!v
         wasm.memories.limits.max = @max(wasm.memories.limits.max, memory_import.limits_max);
         wasm.memories.limits.flags.has_max = wasm.memories.limits.flags.has_max or memory_import.limits_has_max;
     }
+
+    wasm.function_imports_len_prelink = @intCast(wasm.function_imports.entries.len);
+    wasm.data_imports_len_prelink = @intCast(wasm.data_imports.entries.len);
 }
 
 pub fn markFunctionImport(
@@ -3608,7 +3616,7 @@ fn markDataSegment(wasm: *Wasm, segment_index: ObjectDataSegment.Index) link.Fil
     try wasm.markRelocations(segment.relocations(wasm));
 }
 
-fn markDataImport(
+pub fn markDataImport(
     wasm: *Wasm,
     name: String,
     import: *ObjectDataImport,
@@ -4499,11 +4507,40 @@ pub fn navAddr(wasm: *Wasm, nav_index: InternPool.Nav.Index) u32 {
     assert(wasm.flush_buffer.memory_layout_finished);
     const comp = wasm.base.comp;
     assert(comp.config.output_mode != .Obj);
-    // If there is no entry it means the value is zero bits so any address will do.
-    const navs_exe_index: NavsExeIndex = @enumFromInt(wasm.navs_exe.getIndex(nav_index) orelse return 0);
-    log.debug("navAddr {s} {}", .{ navs_exe_index.name(wasm), nav_index });
-    const ds_id: DataSegmentId = .pack(wasm, .{ .nav_exe = navs_exe_index });
-    return wasm.flush_buffer.data_segments.get(ds_id).?;
+    if (wasm.navs_exe.getIndex(nav_index)) |i| {
+        const navs_exe_index: NavsExeIndex = @enumFromInt(i);
+        log.debug("navAddr {s} {}", .{ navs_exe_index.name(wasm), nav_index });
+        const ds_id: DataSegmentId = .pack(wasm, .{ .nav_exe = navs_exe_index });
+        return wasm.flush_buffer.data_segments.get(ds_id).?;
+    }
+    const zcu = comp.zcu.?;
+    const ip = &zcu.intern_pool;
+    const nav = ip.getNav(nav_index);
+    if (nav.getResolvedExtern(ip)) |ext| {
+        if (wasm.getExistingString(ext.name.toSlice(ip))) |symbol_name| {
+            if (wasm.object_data_imports.getPtr(symbol_name)) |import| {
+                switch (import.resolution.unpack(wasm)) {
+                    .unresolved => unreachable,
+                    .object => |object_data_index| {
+                        const object_data = object_data_index.ptr(wasm);
+                        const ds_id: DataSegmentId = .fromObjectDataSegment(wasm, object_data.segment);
+                        const segment_base_addr = wasm.flush_buffer.data_segments.get(ds_id).?;
+                        return segment_base_addr + object_data.offset;
+                    },
+                    .__zig_error_names => @panic("TODO"),
+                    .__zig_error_name_table => @panic("TODO"),
+                    .__heap_base => @panic("TODO"),
+                    .__heap_end => @panic("TODO"),
+                    .uav_exe => @panic("TODO"),
+                    .uav_obj => @panic("TODO"),
+                    .nav_exe => @panic("TODO"),
+                    .nav_obj => @panic("TODO"),
+                }
+            }
+        }
+    }
+    // Otherwise it's a zero bit type; any address will do.
+    return 0;
 }
 
 /// Asserts it is called after `Flush.data_segments` is fully populated and sorted.
diff --git a/src/link/Wasm/Flush.zig b/src/link/Wasm/Flush.zig
index 4e88ad7913..878b911248 100644
--- a/src/link/Wasm/Flush.zig
+++ b/src/link/Wasm/Flush.zig
@@ -122,45 +122,71 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
 
     const entry_name = if (wasm.entry_resolution.isNavOrUnresolved(wasm)) wasm.entry_name else .none;
 
-    // Detect any intrinsics that were called; they need to have dependencies on the symbols marked.
-    // Likewise detect `@tagName` calls so those functions can be included in the output and synthesized.
-    for (wasm.mir_instructions.items(.tag), wasm.mir_instructions.items(.data)) |tag, *data| switch (tag) {
-        .call_intrinsic => {
-            const symbol_name = try wasm.internString(@tagName(data.intrinsic));
-            const i: Wasm.FunctionImport.Index = @enumFromInt(wasm.object_function_imports.getIndex(symbol_name) orelse {
-                return diags.fail("missing compiler runtime intrinsic '{s}' (undefined linker symbol)", .{
-                    @tagName(data.intrinsic),
+    if (comp.zcu) |zcu| {
+        const ip: *const InternPool = &zcu.intern_pool; // No mutations allowed!
+
+        // Detect any intrinsics that were called; they need to have dependencies on the symbols marked.
+        // Likewise detect `@tagName` calls so those functions can be included in the output and synthesized.
+        for (wasm.mir_instructions.items(.tag), wasm.mir_instructions.items(.data)) |tag, *data| switch (tag) {
+            .call_intrinsic => {
+                const symbol_name = try wasm.internString(@tagName(data.intrinsic));
+                const i: Wasm.FunctionImport.Index = @enumFromInt(wasm.object_function_imports.getIndex(symbol_name) orelse {
+                    return diags.fail("missing compiler runtime intrinsic '{s}' (undefined linker symbol)", .{
+                        @tagName(data.intrinsic),
+                    });
                 });
-            });
-            try wasm.markFunctionImport(symbol_name, i.value(wasm), i);
-        },
-        .call_tag_name => {
-            const zcu = comp.zcu.?;
-            const ip = &zcu.intern_pool;
-            assert(ip.indexToKey(data.ip_index) == .enum_type);
-            const gop = try wasm.zcu_funcs.getOrPut(gpa, data.ip_index);
-            if (!gop.found_existing) {
-                wasm.tag_name_table_ref_count += 1;
-                const int_tag_ty = Zcu.Type.fromInterned(data.ip_index).intTagType(zcu);
-                gop.value_ptr.* = .{ .tag_name = .{
-                    .symbol_name = try wasm.internStringFmt("__zig_tag_name_{d}", .{@intFromEnum(data.ip_index)}),
-                    .type_index = try wasm.internFunctionType(.Unspecified, &.{int_tag_ty.ip_index}, .slice_const_u8_sentinel_0, target),
-                    .table_index = @intCast(wasm.tag_name_offs.items.len),
-                } };
-                try wasm.functions.put(gpa, .fromZcuFunc(wasm, @enumFromInt(gop.index)), {});
-                const tag_names = ip.loadEnumType(data.ip_index).names;
-                for (tag_names.get(ip)) |tag_name| {
-                    const slice = tag_name.toSlice(ip);
-                    try wasm.tag_name_offs.append(gpa, @intCast(wasm.tag_name_bytes.items.len));
-                    try wasm.tag_name_bytes.appendSlice(gpa, slice[0 .. slice.len + 1]);
+                try wasm.markFunctionImport(symbol_name, i.value(wasm), i);
+            },
+            .call_tag_name => {
+                assert(ip.indexToKey(data.ip_index) == .enum_type);
+                const gop = try wasm.zcu_funcs.getOrPut(gpa, data.ip_index);
+                if (!gop.found_existing) {
+                    wasm.tag_name_table_ref_count += 1;
+                    const int_tag_ty = Zcu.Type.fromInterned(data.ip_index).intTagType(zcu);
+                    gop.value_ptr.* = .{ .tag_name = .{
+                        .symbol_name = try wasm.internStringFmt("__zig_tag_name_{d}", .{@intFromEnum(data.ip_index)}),
+                        .type_index = try wasm.internFunctionType(.Unspecified, &.{int_tag_ty.ip_index}, .slice_const_u8_sentinel_0, target),
+                        .table_index = @intCast(wasm.tag_name_offs.items.len),
+                    } };
+                    try wasm.functions.put(gpa, .fromZcuFunc(wasm, @enumFromInt(gop.index)), {});
+                    const tag_names = ip.loadEnumType(data.ip_index).names;
+                    for (tag_names.get(ip)) |tag_name| {
+                        const slice = tag_name.toSlice(ip);
+                        try wasm.tag_name_offs.append(gpa, @intCast(wasm.tag_name_bytes.items.len));
+                        try wasm.tag_name_bytes.appendSlice(gpa, slice[0 .. slice.len + 1]);
+                    }
                 }
+            },
+            else => continue,
+        };
+
+        {
+            var i = wasm.function_imports_len_prelink;
+            while (i < f.function_imports.entries.len) {
+                const symbol_name = f.function_imports.keys()[i];
+                if (wasm.object_function_imports.getIndex(symbol_name)) |import_index_usize| {
+                    const import_index: Wasm.FunctionImport.Index = @enumFromInt(import_index_usize);
+                    try wasm.markFunctionImport(symbol_name, import_index.value(wasm), import_index);
+                    f.function_imports.swapRemoveAt(i);
+                    continue;
+                }
+                i += 1;
             }
-        },
-        else => continue,
-    };
+        }
 
-    if (comp.zcu) |zcu| {
-        const ip: *const InternPool = &zcu.intern_pool; // No mutations allowed!
+        {
+            var i = wasm.data_imports_len_prelink;
+            while (i < f.data_imports.entries.len) {
+                const symbol_name = f.data_imports.keys()[i];
+                if (wasm.object_data_imports.getIndex(symbol_name)) |import_index_usize| {
+                    const import_index: Wasm.ObjectDataImport.Index = @enumFromInt(import_index_usize);
+                    try wasm.markDataImport(symbol_name, import_index.value(wasm), import_index);
+                    f.data_imports.swapRemoveAt(i);
+                    continue;
+                }
+                i += 1;
+            }
+        }
 
         if (wasm.error_name_table_ref_count > 0) {
             // Ensure Zcu error name structures are populated.
@@ -437,7 +463,7 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
                 break :b i >= 1 and !wantSegmentMerge(wasm, segment_ids[i - 1], segment_id, category);
             };
             if (want_new_segment) {
-                log.debug("new segment at 0x{x} {} {s} {}", .{ start_addr, segment_id, segment_id.name(wasm), category });
+                log.debug("new segment group at 0x{x} {} {s} {}", .{ start_addr, segment_id, segment_id.name(wasm), category });
                 try f.data_segment_groups.append(gpa, .{
                     .end_addr = @intCast(memory_ptr),
                     .first_segment = first_segment,
@@ -447,6 +473,7 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
 
             const size = segment_id.size(wasm);
             segment_vaddr.* = @intCast(start_addr);
+            log.debug("0x{x} {d} {s}", .{ start_addr, @intFromEnum(segment_id), segment_id.name(wasm) });
             memory_ptr = start_addr + size;
         }
         if (category != .zero) try f.data_segment_groups.append(gpa, .{