link: update decl-specific lazy symbols

3 files changed, 37 insertions, 26 deletions

diff --git a/src/link/Coff.zig b/src/link/Coff.zig
index d20d17f2b1..deac8d1fd6 100644
--- a/src/link/Coff.zig
+++ b/src/link/Coff.zig
@@ -1136,7 +1136,11 @@ pub fn lowerUnnamedConst(self: *Coff, tv: TypedValue, decl_index: Module.Decl.In
     return atom.getSymbolIndex().?;
 }
 
-pub fn updateDecl(self: *Coff, module: *Module, decl_index: Module.Decl.Index) !void {
+pub fn updateDecl(
+    self: *Coff,
+    module: *Module,
+    decl_index: Module.Decl.Index,
+) link.File.UpdateDeclError!void {
     if (build_options.skip_non_native and builtin.object_format != .coff) {
         @panic("Attempted to compile for object format that was disabled by build configuration");
     }
@@ -1146,6 +1150,8 @@ pub fn updateDecl(self: *Coff, module: *Module, decl_index: Module.Decl.Index) !
     const tracy = trace(@src());
     defer tracy.end();
 
+    try self.updateLazySymbol(decl_index);
+
     const decl = module.declPtr(decl_index);
 
     if (decl.val.tag() == .extern_fn) {
@@ -1188,7 +1194,8 @@ pub fn updateDecl(self: *Coff, module: *Module, decl_index: Module.Decl.Index) !
     return self.updateDeclExports(module, decl_index, module.getDeclExports(decl_index));
 }
 
-fn updateLazySymbol(self: *Coff, decl: Module.Decl.OptionalIndex, metadata: LazySymbolMetadata) !void {
+fn updateLazySymbol(self: *Coff, decl: ?Module.Decl.Index) !void {
+    const metadata = self.lazy_syms.get(Module.Decl.OptionalIndex.init(decl)) orelse return;
     const mod = self.base.options.module.?;
     if (metadata.text_atom) |atom| try self.updateLazySymbolAtom(
         link.File.LazySymbol.initDecl(.code, decl, mod),
@@ -1402,7 +1409,7 @@ pub fn updateDeclExports(
     module: *Module,
     decl_index: Module.Decl.Index,
     exports: []const *Module.Export,
-) !void {
+) link.File.UpdateDeclExportsError!void {
     if (build_options.skip_non_native and builtin.object_format != .coff) {
         @panic("Attempted to compile for object format that was disabled by build configuration");
     }
@@ -1599,12 +1606,10 @@ pub fn flushModule(self: *Coff, comp: *Compilation, prog_node: *std.Progress.Nod
 
     // Most lazy symbols can be updated when the corresponding decl is,
     // so we only have to worry about the one without an associated decl.
-    if (self.lazy_syms.get(.none)) |metadata| {
-        self.updateLazySymbol(.none, metadata) catch |err| switch (err) {
-            error.CodegenFail => return error.FlushFailure,
-            else => |e| return e,
-        };
-    }
+    self.updateLazySymbol(null) catch |err| switch (err) {
+        error.CodegenFail => return error.FlushFailure,
+        else => |e| return e,
+    };
 
     const gpa = self.base.allocator;
 
diff --git a/src/link/Elf.zig b/src/link/Elf.zig
index b9c113f834..29ae97150e 100644
--- a/src/link/Elf.zig
+++ b/src/link/Elf.zig
@@ -1034,12 +1034,10 @@ pub fn flushModule(self: *Elf, comp: *Compilation, prog_node: *std.Progress.Node
 
     // Most lazy symbols can be updated when the corresponding decl is,
     // so we only have to worry about the one without an associated decl.
-    if (self.lazy_syms.get(.none)) |metadata| {
-        self.updateLazySymbol(.none, metadata) catch |err| switch (err) {
-            error.CodegenFail => return error.FlushFailure,
-            else => |e| return e,
-        };
-    }
+    self.updateLazySymbol(null) catch |err| switch (err) {
+        error.CodegenFail => return error.FlushFailure,
+        else => |e| return e,
+    };
 
     // TODO This linker code currently assumes there is only 1 compilation unit and it
     // corresponds to the Zig source code.
@@ -2579,7 +2577,11 @@ pub fn updateFunc(self: *Elf, module: *Module, func: *Module.Fn, air: Air, liven
     return self.updateDeclExports(module, decl_index, module.getDeclExports(decl_index));
 }
 
-pub fn updateDecl(self: *Elf, module: *Module, decl_index: Module.Decl.Index) !void {
+pub fn updateDecl(
+    self: *Elf,
+    module: *Module,
+    decl_index: Module.Decl.Index,
+) File.UpdateDeclError!void {
     if (build_options.skip_non_native and builtin.object_format != .elf) {
         @panic("Attempted to compile for object format that was disabled by build configuration");
     }
@@ -2590,6 +2592,8 @@ pub fn updateDecl(self: *Elf, module: *Module, decl_index: Module.Decl.Index) !v
     const tracy = trace(@src());
     defer tracy.end();
 
+    try self.updateLazySymbol(decl_index);
+
     const decl = module.declPtr(decl_index);
 
     if (decl.val.tag() == .extern_fn) {
@@ -2656,7 +2660,8 @@ pub fn updateDecl(self: *Elf, module: *Module, decl_index: Module.Decl.Index) !v
     return self.updateDeclExports(module, decl_index, module.getDeclExports(decl_index));
 }
 
-fn updateLazySymbol(self: *Elf, decl: Module.Decl.OptionalIndex, metadata: LazySymbolMetadata) !void {
+fn updateLazySymbol(self: *Elf, decl: ?Module.Decl.Index) !void {
+    const metadata = self.lazy_syms.get(Module.Decl.OptionalIndex.init(decl)) orelse return;
     const mod = self.base.options.module.?;
     if (metadata.text_atom) |atom| try self.updateLazySymbolAtom(
         File.LazySymbol.initDecl(.code, decl, mod),
@@ -2810,7 +2815,7 @@ pub fn updateDeclExports(
     module: *Module,
     decl_index: Module.Decl.Index,
     exports: []const *Module.Export,
-) !void {
+) File.UpdateDeclExportsError!void {
     if (build_options.skip_non_native and builtin.object_format != .elf) {
         @panic("Attempted to compile for object format that was disabled by build configuration");
     }
diff --git a/src/link/MachO.zig b/src/link/MachO.zig
index a57742507d..5136bd84a2 100644
--- a/src/link/MachO.zig
+++ b/src/link/MachO.zig
@@ -495,12 +495,10 @@ pub fn flushModule(self: *MachO, comp: *Compilation, prog_node: *std.Progress.No
 
     // Most lazy symbols can be updated when the corresponding decl is,
     // so we only have to worry about the one without an associated decl.
-    if (self.lazy_syms.get(.none)) |metadata| {
-        self.updateLazySymbol(.none, metadata) catch |err| switch (err) {
-            error.CodegenFail => return error.FlushFailure,
-            else => |e| return e,
-        };
-    }
+    self.updateLazySymbol(null) catch |err| switch (err) {
+        error.CodegenFail => return error.FlushFailure,
+        else => |e| return e,
+    };
 
     const module = self.base.options.module orelse return error.LinkingWithoutZigSourceUnimplemented;
 
@@ -1962,6 +1960,8 @@ pub fn updateDecl(self: *MachO, module: *Module, decl_index: Module.Decl.Index)
     const tracy = trace(@src());
     defer tracy.end();
 
+    try self.updateLazySymbol(decl_index);
+
     const decl = module.declPtr(decl_index);
 
     if (decl.val.tag() == .extern_fn) {
@@ -2036,7 +2036,8 @@ pub fn updateDecl(self: *MachO, module: *Module, decl_index: Module.Decl.Index)
     try self.updateDeclExports(module, decl_index, module.getDeclExports(decl_index));
 }
 
-fn updateLazySymbol(self: *MachO, decl: Module.Decl.OptionalIndex, metadata: LazySymbolMetadata) !void {
+fn updateLazySymbol(self: *MachO, decl: ?Module.Decl.Index) !void {
+    const metadata = self.lazy_syms.get(Module.Decl.OptionalIndex.init(decl)) orelse return;
     const mod = self.base.options.module.?;
     if (metadata.text_atom) |atom| try self.updateLazySymbolAtom(
         File.LazySymbol.initDecl(.code, decl, mod),
@@ -2353,7 +2354,7 @@ pub fn updateDeclExports(
     module: *Module,
     decl_index: Module.Decl.Index,
     exports: []const *Module.Export,
-) !void {
+) File.UpdateDeclExportsError!void {
     if (build_options.skip_non_native and builtin.object_format != .macho) {
         @panic("Attempted to compile for object format that was disabled by build configuration");
     }