From 49fddbf4c11d4ea493c1e5b6176052aeacc1ad10 Mon Sep 17 00:00:00 2001
From: Luuk de Gram <luuk@degram.dev>
Date: Mon, 22 May 2023 19:06:38 +0200
Subject: [PATCH] wasm: `union_init` correctly store the tag

Previously we would only store the payload, but not the actual tag
that was set. This meant miscompilations where it would incorrectly
return the tag value.

This also adds a tiny optimization for payloads which are not `byRef`
by directly storing them based on offset, rather than first calculating
a pointer to an offset.
---
 src/arch/wasm/CodeGen.zig | 40 ++++++++++++++++++++++++++++++++-------
 1 file changed, 33 insertions(+), 7 deletions(-)

diff --git a/src/arch/wasm/CodeGen.zig b/src/arch/wasm/CodeGen.zig
index 7872d89e9b..ff09919492 100644
--- a/src/arch/wasm/CodeGen.zig
+++ b/src/arch/wasm/CodeGen.zig
@@ -4982,26 +4982,52 @@ fn airUnionInit(func: *CodeGen, inst: Air.Inst.Index) InnerError!void {
     const result = result: {
         const union_ty = func.air.typeOfIndex(inst);
         const layout = union_ty.unionGetLayout(func.target);
+        const union_obj = union_ty.cast(Type.Payload.Union).?.data;
+        const field = union_obj.fields.values()[extra.field_index];
+        const field_name = union_obj.fields.keys()[extra.field_index];
+
+        const tag_int = blk: {
+            const tag_ty = union_ty.unionTagTypeHypothetical();
+            const enum_field_index = tag_ty.enumFieldIndex(field_name).?;
+            var tag_val_payload: Value.Payload.U32 = .{
+                .base = .{ .tag = .enum_field_index },
+                .data = @intCast(u32, enum_field_index),
+            };
+            const tag_val = Value.initPayload(&tag_val_payload.base);
+            break :blk try func.lowerConstant(tag_val, tag_ty);
+        };
         if (layout.payload_size == 0) {
             if (layout.tag_size == 0) {
                 break :result WValue{ .none = {} };
             }
             assert(!isByRef(union_ty, func.target));
-            break :result WValue{ .imm32 = extra.field_index };
+            break :result tag_int;
         }
         assert(isByRef(union_ty, func.target));
 
         const result_ptr = try func.allocStack(union_ty);
         const payload = try func.resolveInst(extra.init);
-        const union_obj = union_ty.cast(Type.Payload.Union).?.data;
-        assert(union_obj.haveFieldTypes());
-        const field = union_obj.fields.values()[extra.field_index];
-
         if (layout.tag_align >= layout.payload_align) {
-            const payload_ptr = try func.buildPointerOffset(result_ptr, layout.tag_size, .new);
-            try func.store(payload_ptr, payload, field.ty, 0);
+            if (isByRef(field.ty, func.target)) {
+                const payload_ptr = try func.buildPointerOffset(result_ptr, layout.tag_size, .new);
+                try func.store(payload_ptr, payload, field.ty, 0);
+            } else {
+                try func.store(result_ptr, payload, field.ty, @intCast(u32, layout.tag_size));
+            }
+
+            if (layout.tag_size > 0) {
+                try func.store(result_ptr, tag_int, union_obj.tag_ty, 0);
+            }
         } else {
             try func.store(result_ptr, payload, field.ty, 0);
+            if (layout.tag_size > 0) {
+                try func.store(
+                    result_ptr,
+                    tag_int,
+                    union_obj.tag_ty,
+                    @intCast(u32, layout.payload_size),
+                );
+            }
         }
         break :result result_ptr;
     };