Merge pull request #10638 from Luukdegram/wasm-refactor

Stage2: wasm - Refactor lowering constants

2 files changed, 129 insertions, 34 deletions

diff --git a/src/link/Wasm.zig b/src/link/Wasm.zig
index e9ebb669fe..82c1f1f630 100644
--- a/src/link/Wasm.zig
+++ b/src/link/Wasm.zig
@@ -19,7 +19,7 @@ const trace = @import("../tracy.zig").trace;
 const build_options = @import("build_options");
 const wasi_libc = @import("../wasi_libc.zig");
 const Cache = @import("../Cache.zig");
-const TypedValue = @import("../TypedValue.zig");
+const Type = @import("../type.zig").Type;
 const LlvmObject = @import("../codegen/llvm.zig").Object;
 const Air = @import("../Air.zig");
 const Liveness = @import("../Liveness.zig");
@@ -234,12 +234,12 @@ pub fn updateFunc(self: *Wasm, module: *Module, func: *Module.Fn, air: Air, live
         .locals = .{},
         .target = self.base.options.target,
         .bin_file = self,
-        .global_error_set = self.base.options.module.?.global_error_set,
+        .module = module,
     };
     defer codegen.deinit();
 
     // generate the 'code' section for the function declaration
-    const result = codegen.genFunc() catch |err| switch (err) {
+    codegen.genFunc() catch |err| switch (err) {
         error.CodegenFail => {
             decl.analysis = .codegen_failure;
             try module.failed_decls.put(module.gpa, decl, codegen.err_msg);
@@ -247,7 +247,7 @@ pub fn updateFunc(self: *Wasm, module: *Module, func: *Module.Fn, air: Air, live
         },
         else => |e| return e,
     };
-    return self.finishUpdateDecl(decl, result, &codegen);
+    return self.finishUpdateDecl(decl, codegen.code.items);
 }
 
 // Generate code for the Decl, storing it in memory to be later written to
@@ -264,40 +264,37 @@ pub fn updateDecl(self: *Wasm, module: *Module, decl: *Module.Decl) !void {
 
     decl.link.wasm.clear();
 
-    var codegen: CodeGen = .{
+    var code_writer = std.ArrayList(u8).init(self.base.allocator);
+    defer code_writer.deinit();
+    var decl_gen: CodeGen.DeclGen = .{
         .gpa = self.base.allocator,
-        .air = undefined,
-        .liveness = undefined,
-        .values = .{},
-        .code = std.ArrayList(u8).init(self.base.allocator),
         .decl = decl,
-        .err_msg = undefined,
-        .locals = .{},
-        .target = self.base.options.target,
+        .symbol_index = decl.link.wasm.sym_index,
         .bin_file = self,
-        .global_error_set = self.base.options.module.?.global_error_set,
+        .err_msg = undefined,
+        .code = &code_writer,
+        .module = module,
     };
-    defer codegen.deinit();
 
     // generate the 'code' section for the function declaration
-    const result = codegen.genDecl() catch |err| switch (err) {
+    const result = decl_gen.genDecl() catch |err| switch (err) {
         error.CodegenFail => {
             decl.analysis = .codegen_failure;
-            try module.failed_decls.put(module.gpa, decl, codegen.err_msg);
+            try module.failed_decls.put(module.gpa, decl, decl_gen.err_msg);
             return;
         },
         else => |e| return e,
     };
 
-    return self.finishUpdateDecl(decl, result, &codegen);
-}
-
-fn finishUpdateDecl(self: *Wasm, decl: *Module.Decl, result: CodeGen.Result, codegen: *CodeGen) !void {
-    const code: []const u8 = switch (result) {
-        .appended => @as([]const u8, codegen.code.items),
-        .externally_managed => |payload| payload,
+    const code = switch (result) {
+        .externally_managed => |data| data,
+        .appended => code_writer.items,
     };
 
+    return self.finishUpdateDecl(decl, code);
+}
+
+fn finishUpdateDecl(self: *Wasm, decl: *Module.Decl, code: []const u8) !void {
     if (decl.isExtern()) {
         try self.addOrUpdateImport(decl);
         return;
@@ -306,9 +303,75 @@ fn finishUpdateDecl(self: *Wasm, decl: *Module.Decl, result: CodeGen.Result, cod
     if (code.len == 0) return;
     const atom: *Atom = &decl.link.wasm;
     atom.size = @intCast(u32, code.len);
+    atom.alignment = decl.ty.abiAlignment(self.base.options.target);
+    self.symbols.items[atom.sym_index].name = decl.name;
     try atom.code.appendSlice(self.base.allocator, code);
 }
 
+/// Creates a new local symbol for a given type (and its bytes it's represented by)
+/// and then append it as a 'contained' atom onto the Decl.
+pub fn createLocalSymbol(self: *Wasm, decl: *Module.Decl, ty: Type) !u32 {
+    assert(ty.zigTypeTag() != .Fn); // cannot create local symbols for functions
+    var symbol: Symbol = .{
+        .name = "unnamed_local",
+        .flags = 0,
+        .tag = .data,
+        .index = undefined,
+    };
+    symbol.setFlag(.WASM_SYM_BINDING_LOCAL);
+    symbol.setFlag(.WASM_SYM_VISIBILITY_HIDDEN);
+
+    var atom = Atom.empty;
+    atom.alignment = ty.abiAlignment(self.base.options.target);
+    try self.symbols.ensureUnusedCapacity(self.base.allocator, 1);
+
+    if (self.symbols_free_list.popOrNull()) |index| {
+        atom.sym_index = index;
+        self.symbols.items[index] = symbol;
+    } else {
+        atom.sym_index = @intCast(u32, self.symbols.items.len);
+        self.symbols.appendAssumeCapacity(symbol);
+    }
+
+    try decl.link.wasm.locals.append(self.base.allocator, atom);
+    return atom.sym_index;
+}
+
+pub fn updateLocalSymbolCode(self: *Wasm, decl: *Module.Decl, symbol_index: u32, code: []const u8) !void {
+    const atom = decl.link.wasm.symbolAtom(symbol_index);
+    atom.size = @intCast(u32, code.len);
+    try atom.code.appendSlice(self.base.allocator, code);
+}
+
+/// For a given decl, find the given symbol index's atom, and create a relocation for the type.
+/// Returns the given pointer address
+pub fn getDeclVAddr(self: *Wasm, decl: *Module.Decl, ty: Type, symbol_index: u32, target_symbol_index: u32, offset: u32) !u32 {
+    const atom = decl.link.wasm.symbolAtom(symbol_index);
+    const is_wasm32 = self.base.options.target.cpu.arch == .wasm32;
+    if (ty.zigTypeTag() == .Fn) {
+        // We found a function pointer, so add it to our table,
+        // as function pointers are not allowed to be stored inside the data section.
+        // They are instead stored in a function table which are called by index.
+        try self.addTableFunction(target_symbol_index);
+        try atom.relocs.append(self.base.allocator, .{
+            .index = target_symbol_index,
+            .offset = offset,
+            .relocation_type = if (is_wasm32) .R_WASM_TABLE_INDEX_I32 else .R_WASM_TABLE_INDEX_I64,
+        });
+    } else {
+        try atom.relocs.append(self.base.allocator, .{
+            .index = target_symbol_index,
+            .offset = offset,
+            .relocation_type = if (is_wasm32) .R_WASM_MEMORY_ADDR_I32 else .R_WASM_MEMORY_ADDR_I64,
+        });
+    }
+    // we do not know the final address at this point,
+    // as atom allocation will determine the address and relocations
+    // will calculate and rewrite this. Therefore, we simply return the symbol index
+    // that was targeted.
+    return target_symbol_index;
+}
+
 pub fn updateDeclExports(
     self: *Wasm,
     module: *Module,
@@ -329,9 +392,12 @@ pub fn freeDecl(self: *Wasm, decl: *Module.Decl) void {
     }
     const atom = &decl.link.wasm;
     self.symbols_free_list.append(self.base.allocator, atom.sym_index) catch {};
-    atom.deinit(self.base.allocator);
     _ = self.decls.remove(decl);
     self.symbols.items[atom.sym_index].tag = .dead; // to ensure it does not end in the names section
+    for (atom.locals.items) |local_atom| {
+        self.symbols.items[local_atom.sym_index].tag = .dead; // also for any local symbol
+    }
+    atom.deinit(self.base.allocator);
 
     if (decl.isExtern()) {
         const import = self.imports.fetchRemove(decl.link.wasm.sym_index).?.value;
@@ -377,14 +443,16 @@ fn addOrUpdateImport(self: *Wasm, decl: *Module.Decl) !void {
     }
 }
 
-fn parseDeclIntoAtom(self: *Wasm, decl: *Module.Decl) !void {
-    const atom: *Atom = &decl.link.wasm;
+const Kind = union(enum) {
+    data: void,
+    function: FnData,
+};
+
+/// Parses an Atom and inserts its metadata into the corresponding sections.
+fn parseAtom(self: *Wasm, atom: *Atom, kind: Kind) !void {
     const symbol: *Symbol = &self.symbols.items[atom.sym_index];
-    symbol.name = decl.name;
-    atom.alignment = decl.ty.abiAlignment(self.base.options.target);
-    const final_index: u32 = switch (decl.ty.zigTypeTag()) {
-        .Fn => result: {
-            const fn_data = decl.fn_link.wasm;
+    const final_index: u32 = switch (kind) {
+        .function => |fn_data| result: {
             const type_index = fn_data.type_index;
             const index = @intCast(u32, self.functions.items.len + self.imported_functions_count);
             try self.functions.append(self.base.allocator, .{ .type_index = type_index });
@@ -402,7 +470,7 @@ fn parseDeclIntoAtom(self: *Wasm, decl: *Module.Decl) !void {
 
             break :result self.code_section_index.?;
         },
-        else => result: {
+        .data => result: {
             const gop = try self.data_segments.getOrPut(self.base.allocator, ".rodata");
             const atom_index = if (gop.found_existing) blk: {
                 self.segments.items[gop.value_ptr.*].size += atom.size;
@@ -430,7 +498,6 @@ fn parseDeclIntoAtom(self: *Wasm, decl: *Module.Decl) !void {
             });
             symbol.tag = .data;
             symbol.index = info_index;
-            atom.alignment = decl.ty.abiAlignment(self.base.options.target);
 
             break :result atom_index;
         },
@@ -617,7 +684,17 @@ pub fn flushModule(self: *Wasm, comp: *Compilation) !void {
     var decl_it = self.decls.keyIterator();
     while (decl_it.next()) |decl| {
         if (decl.*.isExtern()) continue;
-        try self.parseDeclIntoAtom(decl.*);
+        const atom = &decl.*.link.wasm;
+        if (decl.*.ty.zigTypeTag() == .Fn) {
+            try self.parseAtom(atom, .{ .function = decl.*.fn_link.wasm });
+        } else {
+            try self.parseAtom(atom, .data);
+        }
+
+        // also parse atoms for a decl's locals
+        for (atom.locals.items) |*local_atom| {
+            try self.parseAtom(local_atom, .data);
+        }
     }
 
     try self.setupMemory();
diff --git a/src/link/Wasm/Atom.zig b/src/link/Wasm/Atom.zig
index cabd9f1571..09fb6735f7 100644
--- a/src/link/Wasm/Atom.zig
+++ b/src/link/Wasm/Atom.zig
@@ -31,6 +31,10 @@ next: ?*Atom,
 /// is null when this atom is the first in its order
 prev: ?*Atom,
 
+/// Contains atoms local to a decl, all managed by this `Atom`.
+/// When the parent atom is being freed, it will also do so for all local atoms.
+locals: std.ArrayListUnmanaged(Atom) = .{},
+
 /// Represents a default empty wasm `Atom`
 pub const empty: Atom = .{
     .alignment = 0,
@@ -45,6 +49,11 @@ pub const empty: Atom = .{
 pub fn deinit(self: *Atom, gpa: Allocator) void {
     self.relocs.deinit(gpa);
     self.code.deinit(gpa);
+
+    while (self.locals.popOrNull()) |*local| {
+        local.deinit(gpa);
+    }
+    self.locals.deinit(gpa);
 }
 
 /// Sets the length of relocations and code to '0',
@@ -72,6 +81,15 @@ pub fn getFirst(self: *Atom) *Atom {
     return tmp;
 }
 
+/// Returns the atom for the given `symbol_index`.
+/// This can be either the `Atom` itself, or one of its locals.
+pub fn symbolAtom(self: *Atom, symbol_index: u32) *Atom {
+    if (self.sym_index == symbol_index) return self;
+    return for (self.locals.items) |*local_atom| {
+        if (local_atom.sym_index == symbol_index) break local_atom;
+    } else unreachable; // Used a symbol index not present in this atom or its children.
+}
+
 /// Resolves the relocations within the atom, writing the new value
 /// at the calculated offset.
 pub fn resolveRelocs(self: *Atom, wasm_bin: *const Wasm) !void {