3 files changed, 89 insertions, 42 deletions

diff --git a/src/shader_recompiler/backend/glasm/emit_glasm.cpp b/src/shader_recompiler/backend/glasm/emit_glasm.cpp
index 8e5d575a9..ad27b8b06 100644
--- a/src/shader_recompiler/backend/glasm/emit_glasm.cpp
+++ b/src/shader_recompiler/backend/glasm/emit_glasm.cpp
@@ -39,14 +39,16 @@ struct Identity {
 };
 
 template <bool scalar>
-struct RegWrapper {
-    RegWrapper(EmitContext& ctx, Value value)
-        : reg_alloc{ctx.reg_alloc}, allocated{value.type != Type::Register} {
-        if (allocated) {
+class RegWrapper {
+public:
+    RegWrapper(EmitContext& ctx, const IR::Value& ir_value) : reg_alloc{ctx.reg_alloc} {
+        const Value value{reg_alloc.Peek(ir_value)};
+        if (value.type == Type::Register) {
+            inst = ir_value.InstRecursive();
+            reg = Register{value};
+        } else {
             const bool is_long{value.type == Type::F64 || value.type == Type::U64};
             reg = is_long ? reg_alloc.AllocLongReg() : reg_alloc.AllocReg();
-        } else {
-            reg = Register{value};
         }
         switch (value.type) {
         case Type::Register:
@@ -68,8 +70,11 @@ struct RegWrapper {
             break;
         }
     }
+
     ~RegWrapper() {
-        if (allocated) {
+        if (inst) {
+            reg_alloc.Unref(*inst);
+        } else {
             reg_alloc.FreeReg(reg);
         }
     }
@@ -78,19 +83,42 @@ struct RegWrapper {
         return std::conditional_t<scalar, ScalarRegister, Register>{Value{reg}};
     }
 
+private:
     RegAlloc& reg_alloc;
+    IR::Inst* inst{};
     Register reg{};
-    bool allocated{};
+};
+
+template <typename ArgType>
+class ValueWrapper {
+public:
+    ValueWrapper(EmitContext& ctx, const IR::Value& ir_value_)
+        : reg_alloc{ctx.reg_alloc}, ir_value{ir_value_}, value{reg_alloc.Peek(ir_value)} {}
+
+    ~ValueWrapper() {
+        if (!ir_value.IsImmediate()) {
+            reg_alloc.Unref(*ir_value.InstRecursive());
+        }
+    }
+
+    ArgType Extract() {
+        return value;
+    }
+
+private:
+    RegAlloc& reg_alloc;
+    const IR::Value& ir_value;
+    ArgType value;
 };
 
 template <typename ArgType>
 auto Arg(EmitContext& ctx, const IR::Value& arg) {
     if constexpr (std::is_same_v<ArgType, Register>) {
-        return RegWrapper<false>{ctx, ctx.reg_alloc.Consume(arg)};
+        return RegWrapper<false>{ctx, arg};
     } else if constexpr (std::is_same_v<ArgType, ScalarRegister>) {
-        return RegWrapper<true>{ctx, ctx.reg_alloc.Consume(arg)};
+        return RegWrapper<true>{ctx, arg};
     } else if constexpr (std::is_base_of_v<Value, ArgType>) {
-        return Identity{ArgType{ctx.reg_alloc.Consume(arg)}};
+        return ValueWrapper<ArgType>{ctx, arg};
     } else if constexpr (std::is_same_v<ArgType, const IR::Value&>) {
         return Identity{arg};
     } else if constexpr (std::is_same_v<ArgType, u32>) {
diff --git a/src/shader_recompiler/backend/glasm/reg_alloc.cpp b/src/shader_recompiler/backend/glasm/reg_alloc.cpp
index 1a65a5e7d..f556f3aee 100644
--- a/src/shader_recompiler/backend/glasm/reg_alloc.cpp
+++ b/src/shader_recompiler/backend/glasm/reg_alloc.cpp
@@ -21,10 +21,40 @@ Register RegAlloc::LongDefine(IR::Inst& inst) {
     return Define(inst, true);
 }
 
+Value RegAlloc::Peek(const IR::Value& value) {
+    return value.IsImmediate() ? MakeImm(value) : PeekInst(*value.InstRecursive());
+}
+
 Value RegAlloc::Consume(const IR::Value& value) {
-    if (!value.IsImmediate()) {
-        return Consume(*value.InstRecursive());
+    return value.IsImmediate() ? MakeImm(value) : ConsumeInst(*value.InstRecursive());
+}
+
+void RegAlloc::Unref(IR::Inst& inst) {
+    inst.DestructiveRemoveUsage();
+    if (!inst.HasUses()) {
+        Free(inst.Definition<Id>());
     }
+}
+
+Register RegAlloc::AllocReg() {
+    Register ret;
+    ret.type = Type::Register;
+    ret.id = Alloc(false);
+    return ret;
+}
+
+Register RegAlloc::AllocLongReg() {
+    Register ret;
+    ret.type = Type::Register;
+    ret.id = Alloc(true);
+    return ret;
+}
+
+void RegAlloc::FreeReg(Register reg) {
+    Free(reg.id);
+}
+
+Value RegAlloc::MakeImm(const IR::Value& value) {
     Value ret;
     switch (value.Type()) {
     case IR::Type::U1:
@@ -53,43 +83,24 @@ Value RegAlloc::Consume(const IR::Value& value) {
     return ret;
 }
 
-Register RegAlloc::AllocReg() {
-    Register ret;
-    ret.type = Type::Register;
-    ret.id = Alloc(false);
-    return ret;
-}
-
-Register RegAlloc::AllocLongReg() {
-    Register ret;
-    ret.type = Type::Register;
-    ret.id = Alloc(true);
-    return ret;
-}
-
-void RegAlloc::FreeReg(Register reg) {
-    Free(reg.id);
+Register RegAlloc::Define(IR::Inst& inst, bool is_long) {
+    inst.SetDefinition<Id>(Alloc(is_long));
+    return Register{PeekInst(inst)};
 }
 
-Register RegAlloc::Define(IR::Inst& inst, bool is_long) {
-    const Id id{Alloc(is_long)};
-    inst.SetDefinition<Id>(id);
-    Register ret;
+Value RegAlloc::PeekInst(IR::Inst& inst) {
+    Value ret;
     ret.type = Type::Register;
-    ret.id = id;
+    ret.id = inst.Definition<Id>();
     return ret;
 }
 
-Value RegAlloc::Consume(IR::Inst& inst) {
-    const Id id{inst.Definition<Id>()};
+Value RegAlloc::ConsumeInst(IR::Inst& inst) {
     inst.DestructiveRemoveUsage();
     if (!inst.HasUses()) {
-        Free(id);
+        Free(inst.Definition<Id>());
     }
-    Value ret;
-    ret.type = Type::Register;
-    ret.id = id;
-    return ret;
+    return PeekInst(inst);
 }
 
 Id RegAlloc::Alloc(bool is_long) {
diff --git a/src/shader_recompiler/backend/glasm/reg_alloc.h b/src/shader_recompiler/backend/glasm/reg_alloc.h
index 8df73ca18..5742436cf 100644
--- a/src/shader_recompiler/backend/glasm/reg_alloc.h
+++ b/src/shader_recompiler/backend/glasm/reg_alloc.h
@@ -99,8 +99,12 @@ public:
 
     Register LongDefine(IR::Inst& inst);
 
+    [[nodiscard]] Value Peek(const IR::Value& value);
+
     Value Consume(const IR::Value& value);
 
+    void Unref(IR::Inst& inst);
+
     [[nodiscard]] Register AllocReg();
 
     [[nodiscard]] Register AllocLongReg();
@@ -123,9 +127,13 @@ private:
     static constexpr size_t NUM_REGS = 4096;
     static constexpr size_t NUM_ELEMENTS = 4;
 
+    Value MakeImm(const IR::Value& value);
+
     Register Define(IR::Inst& inst, bool is_long);
 
-    Value Consume(IR::Inst& inst);
+    Value PeekInst(IR::Inst& inst);
+
+    Value ConsumeInst(IR::Inst& inst);
 
     Id Alloc(bool is_long);