From d2aaa51bff275d2638ba11b5eab7a5e5bc360ece Mon Sep 17 00:00:00 2001
From: LPeter1997 <lenkefi.peti@gmail.com>
Date: Thu, 12 Oct 2023 17:33:31 +0200
Subject: [PATCH] Removed unused regs

---
 .../Internal/Codegen/CilCodegen.cs              | 11 ++++++++++-
 .../Internal/OptimizingIr/Stackifier.cs         | 17 ++++++++++++-----
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/Draco.Compiler/Internal/Codegen/CilCodegen.cs b/src/Draco.Compiler/Internal/Codegen/CilCodegen.cs
index 81c76da2a..e5066afd0 100644
--- a/src/Draco.Compiler/Internal/Codegen/CilCodegen.cs
+++ b/src/Draco.Compiler/Internal/Codegen/CilCodegen.cs
@@ -82,6 +82,7 @@ public CilCodegen(MetadataCodegen metadataCodegen, IProcedure procedure)
     private int? GetRegisterIndex(Register register)
     {
         if (SymbolEqualityComparer.Default.Equals(register.Type, IntrinsicSymbols.Unit)) return null;
+        if (this.stackifier.RegisterUses[register] == 0) return null;
         if (!this.allocatedRegisters.TryGetValue(register, out var allocatedRegister))
         {
             // NOTE: We need to offset by the number of locals
@@ -484,7 +485,15 @@ private void StoreRegister(Register register)
         if (this.treeDepth > 0) return;
 
         var index = this.GetRegisterIndex(register);
-        if (index is null) return;
+        if (index is null)
+        {
+            if (!SymbolEqualityComparer.Default.Equals(register.Type, IntrinsicSymbols.Unit))
+            {
+                // Need to pop
+                this.InstructionEncoder.OpCode(ILOpCode.Pop);
+            }
+            return;
+        }
         this.InstructionEncoder.StoreLocal(index.Value);
     }
 
diff --git a/src/Draco.Compiler/Internal/OptimizingIr/Stackifier.cs b/src/Draco.Compiler/Internal/OptimizingIr/Stackifier.cs
index 28d12bec7..a2623dae2 100644
--- a/src/Draco.Compiler/Internal/OptimizingIr/Stackifier.cs
+++ b/src/Draco.Compiler/Internal/OptimizingIr/Stackifier.cs
@@ -26,18 +26,25 @@ private static ImmutableDictionary<Register, int> CountRegisterUses(IEnumerable<
         return registerUses.ToImmutable();
     }
 
+    /// <summary>
+    /// A dictionary for the number of register uses for each register.
+    /// </summary>
+    public ImmutableDictionary<Register, int> RegisterUses { get; }
+
     private readonly IProcedure procedure;
-    private readonly ImmutableDictionary<Register, int> registerUses;
 
     public Stackifier(IProcedure procedure)
     {
         this.procedure = procedure;
         var instructions = procedure.BasicBlocks.Values.SelectMany(bb => bb.Instructions);
-        // Count the number of register uses
-        this.registerUses = CountRegisterUses(instructions);
+        this.RegisterUses = CountRegisterUses(instructions);
     }
 
-    // TODO: Doc
+    /// <summary>
+    /// Stackifies the given basic block.
+    /// </summary>
+    /// <param name="basicBlock">The basic block to stackify.</param>
+    /// <returns>The new, stackified array of instructions.</returns>
     public ImmutableArray<IInstruction> Stackify(IBasicBlock basicBlock)
     {
         if (!this.procedure.BasicBlocks.Values.Contains(basicBlock))
@@ -80,7 +87,7 @@ public ImmutableArray<IInstruction> Stackify(IBasicBlock basicBlock)
             // If we have a single-use register as a result immediately before this instruction,
             // all good, part of the tree
             if (!stopped
-             && this.registerUses[reg] == 1
+             && this.RegisterUses[reg] == 1
              && instrIterator.Prev is IValueInstruction valueInstr
              && valueInstr.Target == reg)
             {