link: cleanup lazy symbols

We now only update one lazy symbol in flushModule.
Updating the rest from updateDecl is TBD.

5 files changed, 158 insertions, 163 deletions

diff --git a/src/Module.zig b/src/Module.zig
index 9dab3edccd..70732369ca 100644
--- a/src/Module.zig
+++ b/src/Module.zig
@@ -555,7 +555,7 @@ pub const Decl = struct {
         _,
 
         pub fn init(oi: ?Index) OptionalIndex {
-            return oi orelse .none;
+            return @intToEnum(OptionalIndex, @enumToInt(oi orelse return .none));
         }
 
         pub fn unwrap(oi: OptionalIndex) ?Index {
diff --git a/src/link.zig b/src/link.zig
index 7ca94f4f6d..45873fec26 100644
--- a/src/link.zig
+++ b/src/link.zig
@@ -1106,23 +1106,21 @@ pub const File = struct {
     };
 
     pub const LazySymbol = struct {
-        kind: enum { code, const_data },
-        ty: Type,
+        pub const Kind = enum { code, const_data };
 
-        pub const Context = struct {
-            mod: *Module,
+        kind: Kind,
+        ty: Type,
 
-            pub fn hash(ctx: @This(), sym: LazySymbol) u32 {
-                var hasher = std.hash.Wyhash.init(0);
-                std.hash.autoHash(&hasher, sym.kind);
-                sym.ty.hashWithHasher(&hasher, ctx.mod);
-                return @truncate(u32, hasher.final());
-            }
+        pub fn initDecl(kind: Kind, decl: Module.Decl.OptionalIndex, mod: *Module) LazySymbol {
+            return .{ .kind = kind, .ty = if (decl.unwrap()) |decl_index|
+                mod.declPtr(decl_index).val.castTag(.ty).?.data
+            else
+                Type.anyerror };
+        }
 
-            pub fn eql(ctx: @This(), lhs: LazySymbol, rhs: LazySymbol, _: usize) bool {
-                return lhs.kind == rhs.kind and lhs.ty.eql(rhs.ty, ctx.mod);
-            }
-        };
+        pub fn getDecl(self: LazySymbol) Module.Decl.OptionalIndex {
+            return Module.Decl.OptionalIndex.init(self.ty.getOwnerDeclOrNull());
+        }
     };
 
     pub const C = @import("link/C.zig");
diff --git a/src/link/Coff.zig b/src/link/Coff.zig
index 3a94feb841..825afff36d 100644
--- a/src/link/Coff.zig
+++ b/src/link/Coff.zig
@@ -145,16 +145,11 @@ const Section = struct {
     free_list: std.ArrayListUnmanaged(Atom.Index) = .{},
 };
 
-const LazySymbolTable = std.ArrayHashMapUnmanaged(
-    link.File.LazySymbol,
-    LazySymbolMetadata,
-    link.File.LazySymbol.Context,
-    true,
-);
+const LazySymbolTable = std.AutoArrayHashMapUnmanaged(Module.Decl.OptionalIndex, LazySymbolMetadata);
 
 const LazySymbolMetadata = struct {
-    atom: Atom.Index,
-    section: u16,
+    text_atom: ?Atom.Index = null,
+    rdata_atom: ?Atom.Index = null,
     alignment: u32,
 };
 
@@ -1176,10 +1171,28 @@ pub fn updateDecl(self: *Coff, module: *Module, decl_index: Module.Decl.Index) !
     return self.updateDeclExports(module, decl_index, module.getDeclExports(decl_index));
 }
 
-fn updateLazySymbol(
+fn updateLazySymbol(self: *Coff, decl: Module.Decl.OptionalIndex, metadata: LazySymbolMetadata) !void {
+    const mod = self.base.options.module.?;
+    if (metadata.text_atom) |atom| try self.updateLazySymbolAtom(
+        link.File.LazySymbol.initDecl(.code, decl, mod),
+        atom,
+        self.text_section_index.?,
+        metadata.alignment,
+    );
+    if (metadata.rdata_atom) |atom| try self.updateLazySymbolAtom(
+        link.File.LazySymbol.initDecl(.const_data, decl, mod),
+        atom,
+        self.rdata_section_index.?,
+        metadata.alignment,
+    );
+}
+
+fn updateLazySymbolAtom(
     self: *Coff,
-    lazy_sym: link.File.LazySymbol,
-    lazy_metadata: LazySymbolMetadata,
+    sym: link.File.LazySymbol,
+    atom_index: Atom.Index,
+    section_index: u16,
+    required_alignment: u32,
 ) !void {
     const gpa = self.base.allocator;
     const mod = self.base.options.module.?;
@@ -1188,16 +1201,15 @@ fn updateLazySymbol(
     defer code_buffer.deinit();
 
     const name = try std.fmt.allocPrint(gpa, "__lazy_{s}_{}", .{
-        @tagName(lazy_sym.kind),
-        lazy_sym.ty.fmt(mod),
+        @tagName(sym.kind),
+        sym.ty.fmt(mod),
     });
     defer gpa.free(name);
 
-    const atom_index = lazy_metadata.atom;
     const atom = self.getAtomPtr(atom_index);
     const local_sym_index = atom.getSymbolIndex().?;
 
-    const src = if (lazy_sym.ty.getOwnerDeclOrNull()) |owner_decl|
+    const src = if (sym.ty.getOwnerDeclOrNull()) |owner_decl|
         mod.declPtr(owner_decl).srcLoc()
     else
         Module.SrcLoc{
@@ -1205,14 +1217,9 @@ fn updateLazySymbol(
             .parent_decl_node = undefined,
             .lazy = .unneeded,
         };
-    const res = try codegen.generateLazySymbol(
-        &self.base,
-        src,
-        lazy_sym,
-        &code_buffer,
-        .none,
-        .{ .parent_atom_index = local_sym_index },
-    );
+    const res = try codegen.generateLazySymbol(&self.base, src, sym, &code_buffer, .none, .{
+        .parent_atom_index = local_sym_index,
+    });
     const code = switch (res) {
         .ok => code_buffer.items,
         .fail => |em| {
@@ -1221,11 +1228,10 @@ fn updateLazySymbol(
         },
     };
 
-    const required_alignment = lazy_metadata.alignment;
     const code_len = @intCast(u32, code.len);
     const symbol = atom.getSymbolPtr(self);
     try self.setSymbolName(symbol, name);
-    symbol.section_number = @intToEnum(coff.SectionNumber, lazy_metadata.section + 1);
+    symbol.section_number = @intToEnum(coff.SectionNumber, section_index + 1);
     symbol.type = .{ .complex_type = .NULL, .base_type = .NULL };
 
     const vaddr = try self.allocateAtom(atom_index, code_len, required_alignment);
@@ -1250,24 +1256,18 @@ fn updateLazySymbol(
 
 pub fn getOrCreateAtomForLazySymbol(
     self: *Coff,
-    lazy_sym: link.File.LazySymbol,
+    sym: link.File.LazySymbol,
     alignment: u32,
 ) !Atom.Index {
-    const gop = try self.lazy_syms.getOrPutContext(self.base.allocator, lazy_sym, .{
-        .mod = self.base.options.module.?,
-    });
+    const gop = try self.lazy_syms.getOrPut(self.base.allocator, sym.getDecl());
     errdefer _ = self.lazy_syms.pop();
-    if (!gop.found_existing) {
-        gop.value_ptr.* = .{
-            .atom = try self.createAtom(),
-            .section = switch (lazy_sym.kind) {
-                .code => self.text_section_index.?,
-                .const_data => self.rdata_section_index.?,
-            },
-            .alignment = alignment,
-        };
-    }
-    return gop.value_ptr.atom;
+    if (!gop.found_existing) gop.value_ptr.* = .{ .alignment = alignment };
+    const atom = switch (sym.kind) {
+        .code => &gop.value_ptr.text_atom,
+        .const_data => &gop.value_ptr.rdata_atom,
+    };
+    if (atom.* == null) atom.* = try self.createAtom();
+    return atom.*.?;
 }
 
 pub fn getOrCreateAtomForDecl(self: *Coff, decl_index: Module.Decl.Index) !Atom.Index {
@@ -1600,17 +1600,13 @@ pub fn flushModule(self: *Coff, comp: *Compilation, prog_node: *std.Progress.Nod
     sub_prog_node.activate();
     defer sub_prog_node.end();
 
-    {
-        var lazy_it = self.lazy_syms.iterator();
-        while (lazy_it.next()) |lazy_entry| {
-            self.updateLazySymbol(
-                lazy_entry.key_ptr.*,
-                lazy_entry.value_ptr.*,
-            ) catch |err| switch (err) {
-                error.CodegenFail => return error.FlushFailure,
-                else => |e| return e,
-            };
-        }
+    // 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,
+        };
     }
 
     const gpa = self.base.allocator;
@@ -2489,6 +2485,7 @@ const Module = @import("../Module.zig");
 const Object = @import("Coff/Object.zig");
 const Relocation = @import("Coff/Relocation.zig");
 const StringTable = @import("strtab.zig").StringTable;
+const Type = @import("../type.zig").Type;
 const TypedValue = @import("../TypedValue.zig");
 
 pub const base_tag: link.File.Tag = .coff;
diff --git a/src/link/Elf.zig b/src/link/Elf.zig
index 8079b09166..ec05f5895a 100644
--- a/src/link/Elf.zig
+++ b/src/link/Elf.zig
@@ -64,8 +64,8 @@ const Section = struct {
 };
 
 const LazySymbolMetadata = struct {
-    atom: Atom.Index,
-    shdr: u16,
+    text_atom: ?Atom.Index = null,
+    rodata_atom: ?Atom.Index = null,
     alignment: u32,
 };
 
@@ -208,7 +208,7 @@ relocs: RelocTable = .{},
 
 const RelocTable = std.AutoHashMapUnmanaged(Atom.Index, std.ArrayListUnmanaged(Atom.Reloc));
 const UnnamedConstTable = std.AutoHashMapUnmanaged(Module.Decl.Index, std.ArrayListUnmanaged(Atom.Index));
-const LazySymbolTable = std.ArrayHashMapUnmanaged(File.LazySymbol, LazySymbolMetadata, File.LazySymbol.Context, true);
+const LazySymbolTable = std.AutoArrayHashMapUnmanaged(Module.Decl.OptionalIndex, LazySymbolMetadata);
 
 /// When allocating, the ideal_capacity is calculated by
 /// actual_capacity + (actual_capacity / ideal_factor)
@@ -1065,17 +1065,13 @@ pub fn flushModule(self: *Elf, comp: *Compilation, prog_node: *std.Progress.Node
     sub_prog_node.activate();
     defer sub_prog_node.end();
 
-    {
-        var lazy_it = self.lazy_syms.iterator();
-        while (lazy_it.next()) |lazy_entry| {
-            self.updateLazySymbol(
-                lazy_entry.key_ptr.*,
-                lazy_entry.value_ptr.*,
-            ) catch |err| switch (err) {
-                error.CodegenFail => return error.FlushFailure,
-                else => |e| return e,
-            };
-        }
+    // 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,
+        };
     }
 
     // TODO This linker code currently assumes there is only 1 compilation unit and it
@@ -2424,22 +2420,16 @@ pub fn freeDecl(self: *Elf, decl_index: Module.Decl.Index) void {
     }
 }
 
-pub fn getOrCreateAtomForLazySymbol(self: *Elf, lazy_sym: File.LazySymbol, alignment: u32) !Atom.Index {
-    const gop = try self.lazy_syms.getOrPutContext(self.base.allocator, lazy_sym, .{
-        .mod = self.base.options.module.?,
-    });
+pub fn getOrCreateAtomForLazySymbol(self: *Elf, sym: File.LazySymbol, alignment: u32) !Atom.Index {
+    const gop = try self.lazy_syms.getOrPut(self.base.allocator, sym.getDecl());
     errdefer _ = self.lazy_syms.pop();
-    if (!gop.found_existing) {
-        gop.value_ptr.* = .{
-            .atom = try self.createAtom(),
-            .shdr = switch (lazy_sym.kind) {
-                .code => self.text_section_index.?,
-                .const_data => self.rodata_section_index.?,
-            },
-            .alignment = alignment,
-        };
-    }
-    return gop.value_ptr.atom;
+    if (!gop.found_existing) gop.value_ptr.* = .{ .alignment = alignment };
+    const atom = switch (sym.kind) {
+        .code => &gop.value_ptr.text_atom,
+        .const_data => &gop.value_ptr.rodata_atom,
+    };
+    if (atom.* == null) atom.* = try self.createAtom();
+    return atom.*.?;
 }
 
 pub fn getOrCreateAtomForDecl(self: *Elf, decl_index: Module.Decl.Index) !Atom.Index {
@@ -2708,7 +2698,29 @@ 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, lazy_sym: File.LazySymbol, lazy_metadata: LazySymbolMetadata) !void {
+fn updateLazySymbol(self: *Elf, decl: Module.Decl.OptionalIndex, metadata: LazySymbolMetadata) !void {
+    const mod = self.base.options.module.?;
+    if (metadata.text_atom) |atom| try self.updateLazySymbolAtom(
+        File.LazySymbol.initDecl(.code, decl, mod),
+        atom,
+        self.text_section_index.?,
+        metadata.alignment,
+    );
+    if (metadata.rodata_atom) |atom| try self.updateLazySymbolAtom(
+        File.LazySymbol.initDecl(.const_data, decl, mod),
+        atom,
+        self.rodata_section_index.?,
+        metadata.alignment,
+    );
+}
+
+fn updateLazySymbolAtom(
+    self: *Elf,
+    sym: File.LazySymbol,
+    atom_index: Atom.Index,
+    shdr_index: u16,
+    required_alignment: u32,
+) !void {
     const gpa = self.base.allocator;
     const mod = self.base.options.module.?;
 
@@ -2717,19 +2729,18 @@ fn updateLazySymbol(self: *Elf, lazy_sym: File.LazySymbol, lazy_metadata: LazySy
 
     const name_str_index = blk: {
         const name = try std.fmt.allocPrint(gpa, "__lazy_{s}_{}", .{
-            @tagName(lazy_sym.kind),
-            lazy_sym.ty.fmt(mod),
+            @tagName(sym.kind),
+            sym.ty.fmt(mod),
         });
         defer gpa.free(name);
         break :blk try self.shstrtab.insert(gpa, name);
     };
     const name = self.shstrtab.get(name_str_index).?;
 
-    const atom_index = lazy_metadata.atom;
     const atom = self.getAtom(atom_index);
     const local_sym_index = atom.getSymbolIndex().?;
 
-    const src = if (lazy_sym.ty.getOwnerDeclOrNull()) |owner_decl|
+    const src = if (sym.ty.getOwnerDeclOrNull()) |owner_decl|
         mod.declPtr(owner_decl).srcLoc()
     else
         Module.SrcLoc{
@@ -2737,14 +2748,9 @@ fn updateLazySymbol(self: *Elf, lazy_sym: File.LazySymbol, lazy_metadata: LazySy
             .parent_decl_node = undefined,
             .lazy = .unneeded,
         };
-    const res = try codegen.generateLazySymbol(
-        &self.base,
-        src,
-        lazy_sym,
-        &code_buffer,
-        .none,
-        .{ .parent_atom_index = local_sym_index },
-    );
+    const res = try codegen.generateLazySymbol(&self.base, src, sym, &code_buffer, .none, .{
+        .parent_atom_index = local_sym_index,
+    });
     const code = switch (res) {
         .ok => code_buffer.items,
         .fail => |em| {
@@ -2753,7 +2759,6 @@ fn updateLazySymbol(self: *Elf, lazy_sym: File.LazySymbol, lazy_metadata: LazySy
         },
     };
 
-    const shdr_index = lazy_metadata.shdr;
     const phdr_index = self.sections.items(.phdr_index)[shdr_index];
     const local_sym = atom.getSymbolPtr(self);
     local_sym.* = .{
@@ -2764,7 +2769,6 @@ fn updateLazySymbol(self: *Elf, lazy_sym: File.LazySymbol, lazy_metadata: LazySy
         .st_value = 0,
         .st_size = 0,
     };
-    const required_alignment = lazy_metadata.alignment;
     const vaddr = try self.allocateAtom(atom_index, code.len, required_alignment);
     errdefer self.freeAtom(atom_index);
     log.debug("allocated text block for {s} at 0x{x}", .{ name, vaddr });
diff --git a/src/link/MachO.zig b/src/link/MachO.zig
index 2e1076cdf6..b0a05e4050 100644
--- a/src/link/MachO.zig
+++ b/src/link/MachO.zig
@@ -232,16 +232,11 @@ const is_hot_update_compatible = switch (builtin.target.os.tag) {
     else => false,
 };
 
-const LazySymbolTable = std.ArrayHashMapUnmanaged(
-    link.File.LazySymbol,
-    LazySymbolMetadata,
-    link.File.LazySymbol.Context,
-    true,
-);
+const LazySymbolTable = std.AutoArrayHashMapUnmanaged(Module.Decl.OptionalIndex, LazySymbolMetadata);
 
 const LazySymbolMetadata = struct {
-    atom: Atom.Index,
-    section: u8,
+    text_atom: ?Atom.Index = null,
+    data_const_atom: ?Atom.Index = null,
     alignment: u32,
 };
 
@@ -513,17 +508,13 @@ pub fn flushModule(self: *MachO, comp: *Compilation, prog_node: *std.Progress.No
     sub_prog_node.activate();
     defer sub_prog_node.end();
 
-    {
-        var lazy_it = self.lazy_syms.iterator();
-        while (lazy_it.next()) |lazy_entry| {
-            self.updateLazySymbol(
-                lazy_entry.key_ptr.*,
-                lazy_entry.value_ptr.*,
-            ) catch |err| switch (err) {
-                error.CodegenFail => return error.FlushFailure,
-                else => |e| return e,
-            };
-        }
+    // 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,
+        };
     }
 
     const module = self.base.options.module orelse return error.LinkingWithoutZigSourceUnimplemented;
@@ -2309,7 +2300,29 @@ 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, lazy_sym: File.LazySymbol, lazy_metadata: LazySymbolMetadata) !void {
+fn updateLazySymbol(self: *MachO, decl: Module.Decl.OptionalIndex, metadata: LazySymbolMetadata) !void {
+    const mod = self.base.options.module.?;
+    if (metadata.text_atom) |atom| try self.updateLazySymbolAtom(
+        File.LazySymbol.initDecl(.code, decl, mod),
+        atom,
+        self.text_section_index.?,
+        metadata.alignment,
+    );
+    if (metadata.data_const_atom) |atom| try self.updateLazySymbolAtom(
+        File.LazySymbol.initDecl(.const_data, decl, mod),
+        atom,
+        self.data_const_section_index.?,
+        metadata.alignment,
+    );
+}
+
+fn updateLazySymbolAtom(
+    self: *MachO,
+    sym: File.LazySymbol,
+    atom_index: Atom.Index,
+    section_index: u8,
+    required_alignment: u32,
+) !void {
     const gpa = self.base.allocator;
     const mod = self.base.options.module.?;
 
@@ -2318,19 +2331,18 @@ fn updateLazySymbol(self: *MachO, lazy_sym: File.LazySymbol, lazy_metadata: Lazy
 
     const name_str_index = blk: {
         const name = try std.fmt.allocPrint(gpa, "___lazy_{s}_{}", .{
-            @tagName(lazy_sym.kind),
-            lazy_sym.ty.fmt(mod),
+            @tagName(sym.kind),
+            sym.ty.fmt(mod),
         });
         defer gpa.free(name);
         break :blk try self.strtab.insert(gpa, name);
     };
     const name = self.strtab.get(name_str_index).?;
 
-    const atom_index = lazy_metadata.atom;
     const atom = self.getAtomPtr(atom_index);
     const local_sym_index = atom.getSymbolIndex().?;
 
-    const src = if (lazy_sym.ty.getOwnerDeclOrNull()) |owner_decl|
+    const src = if (sym.ty.getOwnerDeclOrNull()) |owner_decl|
         mod.declPtr(owner_decl).srcLoc()
     else
         Module.SrcLoc{
@@ -2338,14 +2350,9 @@ fn updateLazySymbol(self: *MachO, lazy_sym: File.LazySymbol, lazy_metadata: Lazy
             .parent_decl_node = undefined,
             .lazy = .unneeded,
         };
-    const res = try codegen.generateLazySymbol(
-        &self.base,
-        src,
-        lazy_sym,
-        &code_buffer,
-        .none,
-        .{ .parent_atom_index = local_sym_index },
-    );
+    const res = try codegen.generateLazySymbol(&self.base, src, sym, &code_buffer, .none, .{
+        .parent_atom_index = local_sym_index,
+    });
     const code = switch (res) {
         .ok => code_buffer.items,
         .fail => |em| {
@@ -2354,11 +2361,10 @@ fn updateLazySymbol(self: *MachO, lazy_sym: File.LazySymbol, lazy_metadata: Lazy
         },
     };
 
-    const required_alignment = lazy_metadata.alignment;
     const symbol = atom.getSymbolPtr(self);
     symbol.n_strx = name_str_index;
     symbol.n_type = macho.N_SECT;
-    symbol.n_sect = lazy_metadata.section + 1;
+    symbol.n_sect = section_index + 1;
     symbol.n_desc = 0;
 
     const vaddr = try self.allocateAtom(atom_index, code.len, required_alignment);
@@ -2381,26 +2387,16 @@ fn updateLazySymbol(self: *MachO, lazy_sym: File.LazySymbol, lazy_metadata: Lazy
     try self.writeAtom(atom_index, code);
 }
 
-pub fn getOrCreateAtomForLazySymbol(
-    self: *MachO,
-    lazy_sym: File.LazySymbol,
-    alignment: u32,
-) !Atom.Index {
-    const gop = try self.lazy_syms.getOrPutContext(self.base.allocator, lazy_sym, .{
-        .mod = self.base.options.module.?,
-    });
+pub fn getOrCreateAtomForLazySymbol(self: *MachO, sym: File.LazySymbol, alignment: u32) !Atom.Index {
+    const gop = try self.lazy_syms.getOrPut(self.base.allocator, sym.getDecl());
     errdefer _ = self.lazy_syms.pop();
-    if (!gop.found_existing) {
-        gop.value_ptr.* = .{
-            .atom = try self.createAtom(),
-            .section = switch (lazy_sym.kind) {
-                .code => self.text_section_index.?,
-                .const_data => self.data_const_section_index.?,
-            },
-            .alignment = alignment,
-        };
-    }
-    return gop.value_ptr.atom;
+    if (!gop.found_existing) gop.value_ptr.* = .{ .alignment = alignment };
+    const atom = switch (sym.kind) {
+        .code => &gop.value_ptr.text_atom,
+        .const_data => &gop.value_ptr.data_const_atom,
+    };
+    if (atom.* == null) atom.* = try self.createAtom();
+    return atom.*.?;
 }
 
 pub fn getOrCreateAtomForDecl(self: *MachO, decl_index: Module.Decl.Index) !Atom.Index {