From 576bb3f0a965cd7ff3dad6076567657f18d6675e Mon Sep 17 00:00:00 2001
From: Luuk de Gram <luuk@degram.dev>
Date: Thu, 13 Oct 2022 21:33:26 +0200
Subject: [PATCH] wasm: de -and increment reference count locals

When reusing an operand it increases the reference count, then when
an operand dies it will only decrease the reference count. When
this reaches 0, the local will be virtually freed, meaning it can be
re-used for a new local.
---
 src/arch/wasm/CodeGen.zig | 33 +++++++++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/arch/wasm/CodeGen.zig b/src/arch/wasm/CodeGen.zig
index b2dedd6d56..7543885fe9 100644
--- a/src/arch/wasm/CodeGen.zig
+++ b/src/arch/wasm/CodeGen.zig
@@ -789,8 +789,13 @@ fn iterateBigTomb(self: *Self, inst: Air.Inst.Index, operand_count: usize) !BigT
 fn processDeath(self: *Self, ref: Air.Inst.Ref) void {
     const inst = Air.refToIndex(ref) orelse return;
     if (self.air.instructions.items(.tag)[inst] == .constant) return;
-    var value = self.values.get(ref) orelse return;
-    value.free(self);
+    const value = self.values.getPtr(ref) orelse return;
+    if (value.* != .local) return;
+    std.debug.print("Decreasing reference for ref: %{?d}\n", .{Air.refToIndex(ref)});
+    value.local.references -= 1; // if this panics, a call to `reuseOperand` was forgotten by the developer
+    if (value.local.references == 0) {
+        value.free(self);
+    }
 }
 
 /// Appends a MIR instruction and returns its index within the list of instructions
@@ -909,10 +914,29 @@ fn emitWValue(self: *Self, value: WValue) InnerError!void {
             try self.addInst(.{ .tag = .memory_address, .data = .{ .payload = extra_index } });
         },
         .function_index => |index| try self.addLabel(.function_index, index), // write function index and generate relocation
-        .stack_offset => try self.addLabel(.local_get, self.bottom_stack_value.local), // caller must ensure to address the offset
+        .stack_offset => try self.addLabel(.local_get, self.bottom_stack_value.local.value), // caller must ensure to address the offset
     }
 }
 
+/// If given a local or stack-offset, increases the reference count by 1.
+/// The old `WValue` found at instruction `ref` is then replaced by the
+/// modified `WValue` and returned. When given a non-local or non-stack-offset,
+/// returns the given `operand` itself instead.
+fn reuseOperand(self: *Self, ref: Air.Inst.Ref, operand: WValue) WValue {
+    if (operand != .local and operand != .stack_offset) return operand;
+    var copy = operand;
+    switch (copy) {
+        .local => |*local| local.references += 1,
+        .stack_offset => |*stack_offset| stack_offset.references += 1,
+        else => unreachable,
+    }
+
+    const gop = self.values.getOrPutAssumeCapacity(ref);
+    assert(gop.found_existing);
+    gop.value_ptr.* = copy;
+    return copy;
+}
+
 /// Creates one locals for a given `Type`.
 /// Returns a corresponding `Wvalue` with `local` as active tag
 fn allocLocal(self: *Self, ty: Type) InnerError!WValue {
@@ -2956,7 +2980,8 @@ fn airUnreachable(self: *Self, inst: Air.Inst.Index) InnerError!void {
 fn airBitcast(self: *Self, inst: Air.Inst.Index) InnerError!void {
     const ty_op = self.air.instructions.items(.data)[inst].ty_op;
     const result = if (!self.liveness.isUnused(inst)) result: {
-        break :result try self.resolveInst(ty_op.operand);
+        const operand = try self.resolveInst(ty_op.operand);
+        break :result self.reuseOperand(ty_op.operand, operand);
     } else WValue{ .none = {} };
     self.finishAir(inst, result, &.{});
 }