wasm linker: track overaligned uavs

3 files changed, 38 insertions, 8 deletions

diff --git a/src/arch/wasm/CodeGen.zig b/src/arch/wasm/CodeGen.zig
index e44380d571..91bb00e61e 100644
--- a/src/arch/wasm/CodeGen.zig
+++ b/src/arch/wasm/CodeGen.zig
@@ -161,6 +161,7 @@ const WValue = union(enum) {
     uav_ref: struct {
         ip_index: InternPool.Index,
         offset: i32 = 0,
+        orig_ptr_ty: InternPool.Index = .none,
     },
     /// Offset from the bottom of the virtual stack, with the offset
     /// pointing to where the value lives.
@@ -1062,9 +1063,9 @@ fn emitWValue(cg: *CodeGen, value: WValue) InnerError!void {
                 try cg.addInst(.{
                     .tag = .uav_ref,
                     .data = if (is_obj) .{
-                        .uav_obj = try wasm.refUavObj(uav.ip_index),
+                        .uav_obj = try wasm.refUavObj(uav.ip_index, uav.orig_ptr_ty),
                     } else .{
-                        .uav_exe = try wasm.refUavExe(uav.ip_index),
+                        .uav_exe = try wasm.refUavExe(uav.ip_index, uav.orig_ptr_ty),
                     },
                 });
             } else {
@@ -1072,10 +1073,10 @@ fn emitWValue(cg: *CodeGen, value: WValue) InnerError!void {
                     .tag = .uav_ref_off,
                     .data = .{
                         .payload = if (is_obj) try cg.addExtra(Mir.UavRefOffObj{
-                            .uav_obj = try wasm.refUavObj(uav.ip_index),
+                            .uav_obj = try wasm.refUavObj(uav.ip_index, uav.orig_ptr_ty),
                             .offset = uav.offset,
                         }) else try cg.addExtra(Mir.UavRefOffExe{
-                            .uav_exe = try wasm.refUavExe(uav.ip_index),
+                            .uav_exe = try wasm.refUavExe(uav.ip_index, uav.orig_ptr_ty),
                             .offset = uav.offset,
                         }),
                     },
@@ -3093,7 +3094,7 @@ fn lowerPtr(cg: *CodeGen, ptr_val: InternPool.Index, prev_offset: u64) InnerErro
     const offset: u64 = prev_offset + ptr.byte_offset;
     return switch (ptr.base_addr) {
         .nav => |nav| return .{ .nav_ref = .{ .nav_index = nav, .offset = @intCast(offset) } },
-        .uav => |uav| return .{ .uav_ref = .{ .ip_index = uav.val, .offset = @intCast(offset) } },
+        .uav => |uav| return .{ .uav_ref = .{ .ip_index = uav.val, .offset = @intCast(offset), .orig_ptr_ty = uav.orig_ty } },
         .int => return cg.lowerConstant(try pt.intValue(Type.usize, offset), Type.usize),
         .eu_payload => return cg.fail("Wasm TODO: lower error union payload pointer", .{}),
         .opt_payload => |opt_ptr| return cg.lowerPtr(opt_ptr, offset),
diff --git a/src/codegen.zig b/src/codegen.zig
index f4b524c872..acd7a4965d 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -675,7 +675,7 @@ fn lowerUavRef(
             } else {
                 try wasm.uav_fixups.ensureUnusedCapacity(gpa, 1);
                 wasm.uav_fixups.appendAssumeCapacity(.{
-                    .uavs_exe_index = try wasm.refUavExe(uav.val),
+                    .uavs_exe_index = try wasm.refUavExe(uav.val, uav.orig_ty),
                     .offset = @intCast(code.items.len),
                     .addend = @intCast(offset),
                 });
diff --git a/src/link/Wasm.zig b/src/link/Wasm.zig
index f64f767db1..850339d64d 100644
--- a/src/link/Wasm.zig
+++ b/src/link/Wasm.zig
@@ -189,6 +189,9 @@ navs_exe: std.AutoArrayHashMapUnmanaged(InternPool.Nav.Index, ZcuDataExe) = .emp
 uavs_obj: std.AutoArrayHashMapUnmanaged(InternPool.Index, ZcuDataObj) = .empty,
 /// Tracks ref count to optimize LEB encodings for UAV references.
 uavs_exe: std.AutoArrayHashMapUnmanaged(InternPool.Index, ZcuDataExe) = .empty,
+/// Sparse table of uavs that need to be emitted with greater alignment than
+/// the default for the type.
+overaligned_uavs: std.AutoArrayHashMapUnmanaged(InternPool.Index, Alignment) = .empty,
 /// When the key is an enum type, this represents a `@tagName` function.
 zcu_funcs: std.AutoArrayHashMapUnmanaged(InternPool.Index, ZcuFunc) = .empty,
 nav_exports: std.AutoArrayHashMapUnmanaged(NavExport, Zcu.Export.Index) = .empty,
@@ -1945,6 +1948,7 @@ pub const DataSegmentId = enum(u32) {
                 const zcu = wasm.base.comp.zcu.?;
                 const ip = &zcu.intern_pool;
                 const ip_index = i.key(wasm).*;
+                if (wasm.overaligned_uavs.get(ip_index)) |a| return a;
                 const ty: Zcu.Type = .fromInterned(ip.typeOf(ip_index));
                 const result = ty.abiAlignment(zcu);
                 assert(result != .none);
@@ -3085,6 +3089,7 @@ pub fn deinit(wasm: *Wasm) void {
     wasm.navs_obj.deinit(gpa);
     wasm.uavs_exe.deinit(gpa);
     wasm.uavs_obj.deinit(gpa);
+    wasm.overaligned_uavs.deinit(gpa);
     wasm.zcu_funcs.deinit(gpa);
     wasm.nav_exports.deinit(gpa);
     wasm.uav_exports.deinit(gpa);
@@ -4397,10 +4402,22 @@ pub fn symbolNameIndex(wasm: *Wasm, name: String) Allocator.Error!SymbolTableInd
     return @enumFromInt(gop.index);
 }
 
-pub fn refUavObj(wasm: *Wasm, ip_index: InternPool.Index) !UavsObjIndex {
+pub fn refUavObj(wasm: *Wasm, ip_index: InternPool.Index, orig_ptr_ty: InternPool.Index) !UavsObjIndex {
     const comp = wasm.base.comp;
+    const zcu = comp.zcu.?;
+    const ip = &zcu.intern_pool;
     const gpa = comp.gpa;
     assert(comp.config.output_mode == .Obj);
+
+    if (orig_ptr_ty != .none) {
+        const abi_alignment = Zcu.Type.fromInterned(ip.typeOf(ip_index)).abiAlignment(zcu);
+        const explicit_alignment = ip.indexToKey(orig_ptr_ty).ptr_type.flags.alignment;
+        if (explicit_alignment.compare(.gt, abi_alignment)) {
+            const gop = try wasm.overaligned_uavs.getOrPut(gpa, ip_index);
+            gop.value_ptr.* = if (gop.found_existing) gop.value_ptr.maxStrict(explicit_alignment) else explicit_alignment;
+        }
+    }
+
     const gop = try wasm.uavs_obj.getOrPut(gpa, ip_index);
     if (!gop.found_existing) gop.value_ptr.* = .{
         // Lowering the value is delayed to avoid recursion.
@@ -4410,10 +4427,22 @@ pub fn refUavObj(wasm: *Wasm, ip_index: InternPool.Index) !UavsObjIndex {
     return @enumFromInt(gop.index);
 }
 
-pub fn refUavExe(wasm: *Wasm, ip_index: InternPool.Index) !UavsExeIndex {
+pub fn refUavExe(wasm: *Wasm, ip_index: InternPool.Index, orig_ptr_ty: InternPool.Index) !UavsExeIndex {
     const comp = wasm.base.comp;
+    const zcu = comp.zcu.?;
+    const ip = &zcu.intern_pool;
     const gpa = comp.gpa;
     assert(comp.config.output_mode != .Obj);
+
+    if (orig_ptr_ty != .none) {
+        const abi_alignment = Zcu.Type.fromInterned(ip.typeOf(ip_index)).abiAlignment(zcu);
+        const explicit_alignment = ip.indexToKey(orig_ptr_ty).ptr_type.flags.alignment;
+        if (explicit_alignment.compare(.gt, abi_alignment)) {
+            const gop = try wasm.overaligned_uavs.getOrPut(gpa, ip_index);
+            gop.value_ptr.* = if (gop.found_existing) gop.value_ptr.maxStrict(explicit_alignment) else explicit_alignment;
+        }
+    }
+
     const gop = try wasm.uavs_exe.getOrPut(gpa, ip_index);
     if (gop.found_existing) {
         gop.value_ptr.count += 1;