Misc. CPU improvements (#519)

* Fix and simplify TranslatorCache

* Fix some assignment alignments, remove some unused usings

* Changes to ILEmitter, separate it from ILEmitterCtx

* Rename ILEmitter to ILMethodBuilder

* Rename LdrLit and *_Fix opcodes

* Revert TranslatorCache impl to the more performant one, fix a few issues with it

* Allow EmitOpCode to be called even after everything has been emitted

* Make Emit and AdvanceOpCode private, simplify it a bit now that it starts emiting from the entry point

* Remove unneeded temp use

* Add missing exit call on TestExclusive

* Use better hash

* Implement the == and != operators

17 files changed, 357 insertions, 272 deletions

diff --git a/ChocolArm64/Translation/IILEmit.cs b/ChocolArm64/Translation/IILEmit.cs
index 3c3925ee..320520eb 100644
--- a/ChocolArm64/Translation/IILEmit.cs
+++ b/ChocolArm64/Translation/IILEmit.cs
@@ -2,6 +2,6 @@ namespace ChocolArm64.Translation
 {
     interface IILEmit
     {
-        void Emit(ILEmitter context);
+        void Emit(ILMethodBuilder context);
     }
 }
\ No newline at end of file
diff --git a/ChocolArm64/Translation/ILBarrier.cs b/ChocolArm64/Translation/ILBarrier.cs
index f931e992..37f7558a 100644
--- a/ChocolArm64/Translation/ILBarrier.cs
+++ b/ChocolArm64/Translation/ILBarrier.cs
@@ -2,6 +2,6 @@ namespace ChocolArm64.Translation
 {
     struct ILBarrier : IILEmit
     {
-        public void Emit(ILEmitter context) { }
+        public void Emit(ILMethodBuilder context) { }
     }
 }
\ No newline at end of file
diff --git a/ChocolArm64/Translation/ILBlock.cs b/ChocolArm64/Translation/ILBlock.cs
index d51e8d9e..13657901 100644
--- a/ChocolArm64/Translation/ILBlock.cs
+++ b/ChocolArm64/Translation/ILBlock.cs
@@ -14,19 +14,21 @@ namespace ChocolArm64.Translation
 
         public bool HasStateStore { get; private set; }
 
-        public List<IILEmit> IlEmitters { get; private set; }
+        private List<IILEmit> _emitters;
+
+        public int Count => _emitters.Count;
 
         public ILBlock Next   { get; set; }
         public ILBlock Branch { get; set; }
 
         public ILBlock()
         {
-            IlEmitters = new List<IILEmit>();
+            _emitters = new List<IILEmit>();
         }
 
-        public void Add(IILEmit ilEmitter)
+        public void Add(IILEmit emitter)
         {
-            if (ilEmitter is ILBarrier)
+            if (emitter is ILBarrier)
             {
                 //Those barriers are used to separate the groups of CIL
                 //opcodes emitted by each ARM instruction.
@@ -35,7 +37,7 @@ namespace ChocolArm64.Translation
                 IntAwOutputs = IntOutputs;
                 VecAwOutputs = VecOutputs;
             }
-            else if (ilEmitter is IlOpCodeLoad ld && ILEmitter.IsRegIndex(ld.Index))
+            else if (emitter is ILOpCodeLoad ld && ILMethodBuilder.IsRegIndex(ld.Index))
             {
                 switch (ld.IoType)
                 {
@@ -44,30 +46,26 @@ namespace ChocolArm64.Translation
                     case IoType.Vector: VecInputs |=  (1L << ld.Index)        & ~VecAwOutputs; break;
                 }
             }
-            else if (ilEmitter is IlOpCodeStore st)
+            else if (emitter is ILOpCodeStore st && ILMethodBuilder.IsRegIndex(st.Index))
             {
-                if (ILEmitter.IsRegIndex(st.Index))
-                {
-                    switch (st.IoType)
-                    {
-                        case IoType.Flag:   IntOutputs |= (1L << st.Index) << 32; break;
-                        case IoType.Int:    IntOutputs |=  1L << st.Index;        break;
-                        case IoType.Vector: VecOutputs |=  1L << st.Index;        break;
-                    }
-                }
-
-                if (st.IoType == IoType.Fields)
+                switch (st.IoType)
                 {
-                    HasStateStore = true;
+                    case IoType.Flag:   IntOutputs |= (1L << st.Index) << 32; break;
+                    case IoType.Int:    IntOutputs |=  1L << st.Index;        break;
+                    case IoType.Vector: VecOutputs |=  1L << st.Index;        break;
                 }
             }
+            else if (emitter is ILOpCodeStoreState)
+            {
+                HasStateStore = true;
+            }
 
-            IlEmitters.Add(ilEmitter);
+            _emitters.Add(emitter);
         }
 
-        public void Emit(ILEmitter context)
+        public void Emit(ILMethodBuilder context)
         {
-            foreach (IILEmit ilEmitter in IlEmitters)
+            foreach (IILEmit ilEmitter in _emitters)
             {
                 ilEmitter.Emit(context);
             }
diff --git a/ChocolArm64/Translation/ILEmitterCtx.cs b/ChocolArm64/Translation/ILEmitterCtx.cs
index c1d1be36..778b29ad 100644
--- a/ChocolArm64/Translation/ILEmitterCtx.cs
+++ b/ChocolArm64/Translation/ILEmitterCtx.cs
@@ -14,15 +14,20 @@ namespace ChocolArm64.Translation
 
         private Dictionary<long, ILLabel> _labels;
 
-        private int _blkIndex;
+        private long _subPosition;
+
         private int _opcIndex;
 
-        private Block[]   _graph;
-        private Block     _root;
-        public  Block     CurrBlock => _graph[_blkIndex];
-        public  OpCode64  CurrOp    => _graph[_blkIndex].OpCodes[_opcIndex];
+        private Block _currBlock;
+
+        public Block    CurrBlock => _currBlock;
+        public OpCode64 CurrOp    => _currBlock?.OpCodes[_opcIndex];
+
+        private Dictionary<Block, ILBlock> _visitedBlocks;
 
-        private ILEmitter _emitter;
+        private Queue<Block> _branchTargets;
+
+        private List<ILBlock> _ilBlocks;
 
         private ILBlock _ilBlock;
 
@@ -33,69 +38,61 @@ namespace ChocolArm64.Translation
         //values needed by some functions, since IL doesn't have a swap instruction.
         //You can use any value here as long it doesn't conflict with the indices
         //for the other registers. Any value >= 64 or < 0 will do.
-        private const int Tmp1Index = -1;
-        private const int Tmp2Index = -2;
-        private const int Tmp3Index = -3;
-        private const int Tmp4Index = -4;
-        private const int Tmp5Index = -5;
-        private const int Tmp6Index = -6;
-
-        public ILEmitterCtx(
-            TranslatorCache cache,
-            Block[]         graph,
-            Block           root,
-            string          subName)
-        {
-            _cache = cache ?? throw new ArgumentNullException(nameof(cache));
-            _graph = graph ?? throw new ArgumentNullException(nameof(graph));
-            _root  = root  ?? throw new ArgumentNullException(nameof(root));
+        private const int IntTmpIndex     = -1;
+        private const int RorTmpIndex     = -2;
+        private const int CmpOptTmp1Index = -3;
+        private const int CmpOptTmp2Index = -4;
+        private const int VecTmp1Index    = -5;
+        private const int VecTmp2Index    = -6;
+
+        public ILEmitterCtx(TranslatorCache cache, Block graph)
+        {
+            _cache     = cache ?? throw new ArgumentNullException(nameof(cache));
+            _currBlock = graph ?? throw new ArgumentNullException(nameof(graph));
 
             _labels = new Dictionary<long, ILLabel>();
 
-            _emitter = new ILEmitter(graph, root, subName);
+            _visitedBlocks = new Dictionary<Block, ILBlock>();
 
-            _ilBlock = _emitter.GetIlBlock(0);
+            _visitedBlocks.Add(graph, new ILBlock());
 
-            _opcIndex = -1;
+            _branchTargets = new Queue<Block>();
 
-            if (graph.Length == 0 || !AdvanceOpCode())
-            {
-                throw new ArgumentException(nameof(graph));
-            }
+            _ilBlocks = new List<ILBlock>();
+
+            _subPosition = graph.Position;
+
+            ResetBlockState();
+
+            AdvanceOpCode();
         }
 
-        public TranslatedSub GetSubroutine()
+        public ILBlock[] GetILBlocks()
         {
-            return _emitter.GetSubroutine();
+            EmitAllOpCodes();
+
+            return _ilBlocks.ToArray();
         }
 
-        public bool AdvanceOpCode()
+        private void EmitAllOpCodes()
         {
-            if (_opcIndex + 1 == CurrBlock.OpCodes.Count &&
-                _blkIndex + 1 == _graph.Length)
+            do
             {
-                return false;
+                EmitOpCode();
             }
+            while (AdvanceOpCode());
+        }
 
-            while (++_opcIndex >= (CurrBlock?.OpCodes.Count ?? 0))
+        private void EmitOpCode()
+        {
+            if (_currBlock == null)
             {
-                _blkIndex++;
-                _opcIndex = -1;
-
-                _optOpLastFlagSet = null;
-                _optOpLastCompare = null;
-
-                _ilBlock = _emitter.GetIlBlock(_blkIndex);
+                return;
             }
 
-            return true;
-        }
-
-        public void EmitOpCode()
-        {
             if (_opcIndex == 0)
             {
-                MarkLabel(GetLabel(CurrBlock.Position));
+                MarkLabel(GetLabel(_currBlock.Position));
 
                 EmitSynchronization();
             }
@@ -109,7 +106,7 @@ namespace ChocolArm64.Translation
         {
             EmitLdarg(TranslatedSub.StateArgIdx);
 
-            EmitLdc_I4(CurrBlock.OpCodes.Count);
+            EmitLdc_I4(_currBlock.OpCodes.Count);
 
             EmitPrivateCall(typeof(CpuThreadState), nameof(CpuThreadState.Synchronize));
 
@@ -126,9 +123,86 @@ namespace ChocolArm64.Translation
             MarkLabel(lblContinue);
         }
 
+        private bool AdvanceOpCode()
+        {
+            if (_currBlock == null)
+            {
+                return false;
+            }
+
+            while (++_opcIndex >= _currBlock.OpCodes.Count)
+            {
+                if (!AdvanceBlock())
+                {
+                    return false;
+                }
+
+                ResetBlockState();
+            }
+
+            return true;
+        }
+
+        private bool AdvanceBlock()
+        {
+            if (_currBlock.Branch != null)
+            {
+                if (_visitedBlocks.TryAdd(_currBlock.Branch, _ilBlock.Branch))
+                {
+                    _branchTargets.Enqueue(_currBlock.Branch);
+                }
+            }
+
+            if (_currBlock.Next != null)
+            {
+                if (_visitedBlocks.TryAdd(_currBlock.Next, _ilBlock.Next))
+                {
+                    _currBlock = _currBlock.Next;
+
+                    return true;
+                }
+                else
+                {
+                    Emit(OpCodes.Br, GetLabel(_currBlock.Next.Position));
+                }
+            }
+
+            return _branchTargets.TryDequeue(out _currBlock);
+        }
+
+        private void ResetBlockState()
+        {
+            _ilBlock = _visitedBlocks[_currBlock];
+
+            _ilBlocks.Add(_ilBlock);
+
+            _ilBlock.Next   = GetOrCreateILBlock(_currBlock.Next);
+            _ilBlock.Branch = GetOrCreateILBlock(_currBlock.Branch);
+
+            _opcIndex = -1;
+
+            _optOpLastFlagSet = null;
+            _optOpLastCompare = null;
+        }
+
+        private ILBlock GetOrCreateILBlock(Block block)
+        {
+            if (block == null)
+            {
+                return null;
+            }
+
+            if (_visitedBlocks.TryGetValue(block, out ILBlock ilBlock))
+            {
+                return ilBlock;
+            }
+
+            return new ILBlock();
+        }
+
         public bool TryOptEmitSubroutineCall()
         {
-            if (CurrBlock.Next == null)
+            if (_currBlock.Next == null)
             {
                 return false;
             }
@@ -148,7 +222,7 @@ namespace ChocolArm64.Translation
                 EmitLdarg(index);
             }
 
-            foreach (Register reg in subroutine.Params)
+            foreach (Register reg in subroutine.SubArgs)
             {
                 switch (reg.Type)
                 {
@@ -160,7 +234,7 @@ namespace ChocolArm64.Translation
 
             EmitCall(subroutine.Method);
 
-            subroutine.AddCaller(_root.Position);
+            subroutine.AddCaller(_subPosition);
 
             return true;
         }
@@ -171,11 +245,11 @@ namespace ChocolArm64.Translation
 
             InstEmitAluHelper.EmitDataLoadOpers(this);
 
-            Stloc(Tmp4Index, IoType.Int);
-            Stloc(Tmp3Index, IoType.Int);
+            Stloc(CmpOptTmp2Index, IoType.Int);
+            Stloc(CmpOptTmp1Index, IoType.Int);
         }
 
-        private Dictionary<Cond, System.Reflection.Emit.OpCode> _branchOps = new Dictionary<Cond, System.Reflection.Emit.OpCode>()
+        private Dictionary<Cond, OpCode> _branchOps = new Dictionary<Cond, OpCode>()
         {
             { Cond.Eq,   OpCodes.Beq    },
             { Cond.Ne,   OpCodes.Bne_Un },
@@ -191,15 +265,15 @@ namespace ChocolArm64.Translation
 
         public void EmitCondBranch(ILLabel target, Cond cond)
         {
-            System.Reflection.Emit.OpCode ilOp;
+            OpCode ilOp;
 
             int intCond = (int)cond;
 
             if (_optOpLastCompare != null &&
                 _optOpLastCompare == _optOpLastFlagSet && _branchOps.ContainsKey(cond))
             {
-                Ldloc(Tmp3Index, IoType.Int, _optOpLastCompare.RegisterSize);
-                Ldloc(Tmp4Index, IoType.Int, _optOpLastCompare.RegisterSize);
+                Ldloc(CmpOptTmp1Index, IoType.Int, _optOpLastCompare.RegisterSize);
+                Ldloc(CmpOptTmp2Index, IoType.Int, _optOpLastCompare.RegisterSize);
 
                 ilOp = _branchOps[cond];
             }
@@ -285,11 +359,11 @@ namespace ChocolArm64.Translation
             }
         }
 
-        public void EmitLsl(int amount) => EmitIlShift(amount, OpCodes.Shl);
-        public void EmitLsr(int amount) => EmitIlShift(amount, OpCodes.Shr_Un);
-        public void EmitAsr(int amount) => EmitIlShift(amount, OpCodes.Shr);
+        public void EmitLsl(int amount) => EmitILShift(amount, OpCodes.Shl);
+        public void EmitLsr(int amount) => EmitILShift(amount, OpCodes.Shr_Un);
+        public void EmitAsr(int amount) => EmitILShift(amount, OpCodes.Shr);
 
-        private void EmitIlShift(int amount, System.Reflection.Emit.OpCode ilOp)
+        private void EmitILShift(int amount, OpCode ilOp)
         {
             if (amount > 0)
             {
@@ -303,14 +377,14 @@ namespace ChocolArm64.Translation
         {
             if (amount > 0)
             {
-                Stloc(Tmp2Index, IoType.Int);
-                Ldloc(Tmp2Index, IoType.Int);
+                Stloc(RorTmpIndex, IoType.Int);
+                Ldloc(RorTmpIndex, IoType.Int);
 
                 EmitLdc_I4(amount);
 
                 Emit(OpCodes.Shr_Un);
 
-                Ldloc(Tmp2Index, IoType.Int);
+                Ldloc(RorTmpIndex, IoType.Int);
 
                 EmitLdc_I4(CurrOp.GetBitsCount() - amount);
 
@@ -336,12 +410,12 @@ namespace ChocolArm64.Translation
             _ilBlock.Add(label);
         }
 
-        public void Emit(System.Reflection.Emit.OpCode ilOp)
+        public void Emit(OpCode ilOp)
         {
             _ilBlock.Add(new ILOpCode(ilOp));
         }
 
-        public void Emit(System.Reflection.Emit.OpCode ilOp, ILLabel label)
+        public void Emit(OpCode ilOp, ILLabel label)
         {
             _ilBlock.Add(new ILOpCodeBranch(ilOp, label));
         }
@@ -353,7 +427,7 @@ namespace ChocolArm64.Translation
 
         public void EmitLdarg(int index)
         {
-            _ilBlock.Add(new IlOpCodeLoad(index, IoType.Arg));
+            _ilBlock.Add(new ILOpCodeLoad(index, IoType.Arg));
         }
 
         public void EmitLdintzr(int index)
@@ -380,24 +454,29 @@ namespace ChocolArm64.Translation
             }
         }
 
-        public void EmitLoadState(Block retBlk)
+        public void EmitLoadState()
         {
-            _ilBlock.Add(new IlOpCodeLoad(Array.IndexOf(_graph, retBlk), IoType.Fields));
+            if (_ilBlock.Next == null)
+            {
+                throw new InvalidOperationException("Can't load state for next block, because there's no next block.");
+            }
+
+            _ilBlock.Add(new ILOpCodeLoadState(_ilBlock.Next));
         }
 
         public void EmitStoreState()
         {
-            _ilBlock.Add(new IlOpCodeStore(Array.IndexOf(_graph, CurrBlock), IoType.Fields));
+            _ilBlock.Add(new ILOpCodeStoreState(_ilBlock));
         }
 
-        public void EmitLdtmp() => EmitLdint(Tmp1Index);
-        public void EmitSttmp() => EmitStint(Tmp1Index);
+        public void EmitLdtmp() => EmitLdint(IntTmpIndex);
+        public void EmitSttmp() => EmitStint(IntTmpIndex);
 
-        public void EmitLdvectmp() => EmitLdvec(Tmp5Index);
-        public void EmitStvectmp() => EmitStvec(Tmp5Index);
+        public void EmitLdvectmp() => EmitLdvec(VecTmp1Index);
+        public void EmitStvectmp() => EmitStvec(VecTmp1Index);
 
-        public void EmitLdvectmp2() => EmitLdvec(Tmp6Index);
-        public void EmitStvectmp2() => EmitStvec(Tmp6Index);
+        public void EmitLdvectmp2() => EmitLdvec(VecTmp2Index);
+        public void EmitStvectmp2() => EmitStvec(VecTmp2Index);
 
         public void EmitLdint(int index) => Ldloc(index, IoType.Int);
         public void EmitStint(int index) => Stloc(index, IoType.Int);
@@ -415,17 +494,17 @@ namespace ChocolArm64.Translation
 
         private void Ldloc(int index, IoType ioType)
         {
-            _ilBlock.Add(new IlOpCodeLoad(index, ioType, CurrOp.RegisterSize));
+            _ilBlock.Add(new ILOpCodeLoad(index, ioType, CurrOp.RegisterSize));
         }
 
         private void Ldloc(int index, IoType ioType, RegisterSize registerSize)
         {
-            _ilBlock.Add(new IlOpCodeLoad(index, ioType, registerSize));
+            _ilBlock.Add(new ILOpCodeLoad(index, ioType, registerSize));
         }
 
         private void Stloc(int index, IoType ioType)
         {
-            _ilBlock.Add(new IlOpCodeStore(index, ioType, CurrOp.RegisterSize));
+            _ilBlock.Add(new ILOpCodeStore(index, ioType, CurrOp.RegisterSize));
         }
 
         public void EmitCallPropGet(Type objType, string propName)
@@ -536,7 +615,7 @@ namespace ChocolArm64.Translation
             EmitZnCheck(OpCodes.Clt, (int)PState.NBit);
         }
 
-        private void EmitZnCheck(System.Reflection.Emit.OpCode ilCmpOp, int flag)
+        private void EmitZnCheck(OpCode ilCmpOp, int flag)
         {
             Emit(OpCodes.Dup);
             Emit(OpCodes.Ldc_I4_0);
diff --git a/ChocolArm64/Translation/ILLabel.cs b/ChocolArm64/Translation/ILLabel.cs
index 4f96edcc..f423a425 100644
--- a/ChocolArm64/Translation/ILLabel.cs
+++ b/ChocolArm64/Translation/ILLabel.cs
@@ -8,12 +8,12 @@ namespace ChocolArm64.Translation
 
         private Label _lbl;
 
-        public void Emit(ILEmitter context)
+        public void Emit(ILMethodBuilder context)
         {
             context.Generator.MarkLabel(GetLabel(context));
         }
 
-        public Label GetLabel(ILEmitter context)
+        public Label GetLabel(ILMethodBuilder context)
         {
             if (!_hasLabel)
             {
diff --git a/ChocolArm64/Translation/ILEmitter.cs b/ChocolArm64/Translation/ILMethodBuilder.cs
index 543528d7..70d9a2db 100644
--- a/ChocolArm64/Translation/ILEmitter.cs
+++ b/ChocolArm64/Translation/ILMethodBuilder.cs
@@ -1,4 +1,3 @@
-using ChocolArm64.Decoders;
 using ChocolArm64.State;
 using System;
 using System.Collections.Generic;
@@ -7,7 +6,7 @@ using System.Runtime.Intrinsics;
 
 namespace ChocolArm64.Translation
 {
-    class ILEmitter
+    class ILMethodBuilder
     {
         public LocalAlloc LocalAlloc { get; private set; }
 
@@ -17,70 +16,23 @@ namespace ChocolArm64.Translation
 
         private ILBlock[] _ilBlocks;
 
-        private ILBlock _root;
-
-        private TranslatedSub _subroutine;
-
         private string _subName;
 
         private int _localsCount;
 
-        public ILEmitter(Block[] graph, Block root, string subName)
+        public ILMethodBuilder(ILBlock[] ilBlocks, string subName)
         {
-            _subName = subName;
-
-            _locals = new Dictionary<Register, int>();
-
-            _ilBlocks = new ILBlock[graph.Length];
-
-            ILBlock GetBlock(int index)
-            {
-                if (index < 0 || index >= _ilBlocks.Length)
-                {
-                    return null;
-                }
-
-                if (_ilBlocks[index] == null)
-                {
-                    _ilBlocks[index] = new ILBlock();
-                }
-
-                return _ilBlocks[index];
-            }
-
-            for (int index = 0; index < _ilBlocks.Length; index++)
-            {
-                ILBlock block = GetBlock(index);
-
-                block.Next   = GetBlock(Array.IndexOf(graph, graph[index].Next));
-                block.Branch = GetBlock(Array.IndexOf(graph, graph[index].Branch));
-            }
-
-            _root = _ilBlocks[Array.IndexOf(graph, root)];
+            _ilBlocks = ilBlocks;
+            _subName  = subName;
         }
 
-        public ILBlock GetIlBlock(int index) => _ilBlocks[index];
-
         public TranslatedSub GetSubroutine()
         {
-            LocalAlloc = new LocalAlloc(_ilBlocks, _root);
+            LocalAlloc = new LocalAlloc(_ilBlocks, _ilBlocks[0]);
 
-            InitSubroutine();
-            InitLocals();
+            List<Register> subArgs = new List<Register>();
 
-            foreach (ILBlock ilBlock in _ilBlocks)
-            {
-                ilBlock.Emit(this);
-            }
-
-            return _subroutine;
-        }
-
-        private void InitSubroutine()
-        {
-            List<Register> Params = new List<Register>();
-
-            void SetParams(long inputs, RegisterType baseType)
+            void SetArgs(long inputs, RegisterType baseType)
             {
                 for (int bit = 0; bit < 64; bit++)
                 {
@@ -88,37 +40,43 @@ namespace ChocolArm64.Translation
 
                     if ((inputs & mask) != 0)
                     {
-                        Params.Add(GetRegFromBit(bit, baseType));
+                        subArgs.Add(GetRegFromBit(bit, baseType));
                     }
                 }
             }
 
-            SetParams(LocalAlloc.GetIntInputs(_root), RegisterType.Int);
-            SetParams(LocalAlloc.GetVecInputs(_root), RegisterType.Vector);
+            SetArgs(LocalAlloc.GetIntInputs(_ilBlocks[0]), RegisterType.Int);
+            SetArgs(LocalAlloc.GetVecInputs(_ilBlocks[0]), RegisterType.Vector);
 
-            DynamicMethod mthd = new DynamicMethod(_subName, typeof(long), GetParamTypes(Params));
+            DynamicMethod method = new DynamicMethod(_subName, typeof(long), GetArgumentTypes(subArgs));
 
-            Generator = mthd.GetILGenerator();
+            Generator = method.GetILGenerator();
 
-            _subroutine = new TranslatedSub(mthd, Params);
-        }
+            TranslatedSub subroutine = new TranslatedSub(method, subArgs);
 
-        private void InitLocals()
-        {
-            int paramsStart = TranslatedSub.FixedArgTypes.Length;
+            int argsStart = TranslatedSub.FixedArgTypes.Length;
 
             _locals = new Dictionary<Register, int>();
 
-            for (int index = 0; index < _subroutine.Params.Count; index++)
+            _localsCount = 0;
+
+            for (int index = 0; index < subroutine.SubArgs.Count; index++)
             {
-                Register reg = _subroutine.Params[index];
+                Register reg = subroutine.SubArgs[index];
 
-                Generator.EmitLdarg(index + paramsStart);
+                Generator.EmitLdarg(index + argsStart);
                 Generator.EmitStloc(GetLocalIndex(reg));
             }
+
+            foreach (ILBlock ilBlock in _ilBlocks)
+            {
+                ilBlock.Emit(this);
+            }
+
+            return subroutine;
         }
 
-        private Type[] GetParamTypes(IList<Register> Params)
+        private Type[] GetArgumentTypes(IList<Register> Params)
         {
             Type[] fixedArgs = TranslatedSub.FixedArgTypes;
 
@@ -140,7 +98,7 @@ namespace ChocolArm64.Translation
         {
             if (!_locals.TryGetValue(reg, out int index))
             {
-                Generator.DeclareLocal(GetLocalType(reg));
+                Generator.DeclareLocal(GetFieldType(reg.Type));
 
                 index = _localsCount++;
 
@@ -150,9 +108,7 @@ namespace ChocolArm64.Translation
             return index;
         }
 
-        public Type GetLocalType(Register reg) => GetFieldType(reg.Type);
-
-        public Type GetFieldType(RegisterType regType)
+        private static Type GetFieldType(RegisterType regType)
         {
             switch (regType)
             {
@@ -182,7 +138,7 @@ namespace ChocolArm64.Translation
 
         public static bool IsRegIndex(int index)
         {
-            return index >= 0 && index < 32;
+            return (uint)index < 32;
         }
     }
 }
\ No newline at end of file
diff --git a/ChocolArm64/Translation/ILOpCode.cs b/ChocolArm64/Translation/ILOpCode.cs
index eb91639e..4021603c 100644
--- a/ChocolArm64/Translation/ILOpCode.cs
+++ b/ChocolArm64/Translation/ILOpCode.cs
@@ -11,7 +11,7 @@ namespace ChocolArm64.Translation
             _ilOp = ilOp;
         }
 
-        public void Emit(ILEmitter context)
+        public void Emit(ILMethodBuilder context)
         {
             context.Generator.Emit(_ilOp);
         }
diff --git a/ChocolArm64/Translation/ILOpCodeBranch.cs b/ChocolArm64/Translation/ILOpCodeBranch.cs
index b0ba2331..22b80b5d 100644
--- a/ChocolArm64/Translation/ILOpCodeBranch.cs
+++ b/ChocolArm64/Translation/ILOpCodeBranch.cs
@@ -13,7 +13,7 @@ namespace ChocolArm64.Translation
             _label = label;
         }
 
-        public void Emit(ILEmitter context)
+        public void Emit(ILMethodBuilder context)
         {
             context.Generator.Emit(_ilOp, _label.GetLabel(context));
         }
diff --git a/ChocolArm64/Translation/ILOpCodeCall.cs b/ChocolArm64/Translation/ILOpCodeCall.cs
index 0b591d46..8486a791 100644
--- a/ChocolArm64/Translation/ILOpCodeCall.cs
+++ b/ChocolArm64/Translation/ILOpCodeCall.cs
@@ -12,7 +12,7 @@ namespace ChocolArm64.Translation
             _mthdInfo = mthdInfo;
         }
 
-        public void Emit(ILEmitter context)
+        public void Emit(ILMethodBuilder context)
         {
             context.Generator.Emit(OpCodes.Call, _mthdInfo);
         }
diff --git a/ChocolArm64/Translation/ILOpCodeConst.cs b/ChocolArm64/Translation/ILOpCodeConst.cs
index 5497eba1..2aaf8676 100644
--- a/ChocolArm64/Translation/ILOpCodeConst.cs
+++ b/ChocolArm64/Translation/ILOpCodeConst.cs
@@ -51,7 +51,7 @@ namespace ChocolArm64.Translation
             _value = new ImmVal { R8 = value };
         }
 
-        public void Emit(ILEmitter context)
+        public void Emit(ILMethodBuilder context)
         {
             switch (_type)
             {
diff --git a/ChocolArm64/Translation/ILOpCodeLoad.cs b/ChocolArm64/Translation/ILOpCodeLoad.cs
index 9dae10cc..c31e06bb 100644
--- a/ChocolArm64/Translation/ILOpCodeLoad.cs
+++ b/ChocolArm64/Translation/ILOpCodeLoad.cs
@@ -3,7 +3,7 @@ using System.Reflection.Emit;
 
 namespace ChocolArm64.Translation
 {
-    struct IlOpCodeLoad : IILEmit
+    struct ILOpCodeLoad : IILEmit
     {
         public int Index { get; private set; }
 
@@ -11,55 +11,26 @@ namespace ChocolArm64.Translation
 
         public RegisterSize RegisterSize { get; private set; }
 
-        public IlOpCodeLoad(int index, IoType ioType, RegisterSize registerSize = 0)
+        public ILOpCodeLoad(int index, IoType ioType, RegisterSize registerSize = 0)
         {
             Index        = index;
             IoType       = ioType;
             RegisterSize = registerSize;
         }
 
-        public void Emit(ILEmitter context)
+        public void Emit(ILMethodBuilder context)
         {
             switch (IoType)
             {
                 case IoType.Arg: context.Generator.EmitLdarg(Index); break;
 
-                case IoType.Fields:
-                {
-                    long intInputs = context.LocalAlloc.GetIntInputs(context.GetIlBlock(Index));
-                    long vecInputs = context.LocalAlloc.GetVecInputs(context.GetIlBlock(Index));
-
-                    LoadLocals(context, intInputs, RegisterType.Int);
-                    LoadLocals(context, vecInputs, RegisterType.Vector);
-
-                    break;
-                }
-
                 case IoType.Flag:   EmitLdloc(context, Index, RegisterType.Flag);   break;
                 case IoType.Int:    EmitLdloc(context, Index, RegisterType.Int);    break;
                 case IoType.Vector: EmitLdloc(context, Index, RegisterType.Vector); break;
             }
         }
 
-        private void LoadLocals(ILEmitter context, long inputs, RegisterType baseType)
-        {
-            for (int bit = 0; bit < 64; bit++)
-            {
-                long mask = 1L << bit;
-
-                if ((inputs & mask) != 0)
-                {
-                    Register reg = ILEmitter.GetRegFromBit(bit, baseType);
-
-                    context.Generator.EmitLdarg(TranslatedSub.StateArgIdx);
-                    context.Generator.Emit(OpCodes.Ldfld, reg.GetField());
-
-                    context.Generator.EmitStloc(context.GetLocalIndex(reg));
-                }
-            }
-        }
-
-        private void EmitLdloc(ILEmitter context, int index, RegisterType registerType)
+        private void EmitLdloc(ILMethodBuilder context, int index, RegisterType registerType)
         {
             Register reg = new Register(index, registerType);
 
diff --git a/ChocolArm64/Translation/ILOpCodeLoadState.cs b/ChocolArm64/Translation/ILOpCodeLoadState.cs
new file mode 100644
index 00000000..ddab6110
--- /dev/null
+++ b/ChocolArm64/Translation/ILOpCodeLoadState.cs
@@ -0,0 +1,42 @@
+using ChocolArm64.State;
+using System.Reflection.Emit;
+
+namespace ChocolArm64.Translation
+{
+    struct ILOpCodeLoadState : IILEmit
+    {
+        private ILBlock _block;
+
+        public ILOpCodeLoadState(ILBlock block)
+        {
+            _block = block;
+        }
+
+        public void Emit(ILMethodBuilder context)
+        {
+            long intInputs = context.LocalAlloc.GetIntInputs(_block);
+            long vecInputs = context.LocalAlloc.GetVecInputs(_block);
+
+            LoadLocals(context, intInputs, RegisterType.Int);
+            LoadLocals(context, vecInputs, RegisterType.Vector);
+        }
+
+        private void LoadLocals(ILMethodBuilder context, long inputs, RegisterType baseType)
+        {
+            for (int bit = 0; bit < 64; bit++)
+            {
+                long mask = 1L << bit;
+
+                if ((inputs & mask) != 0)
+                {
+                    Register reg = ILMethodBuilder.GetRegFromBit(bit, baseType);
+
+                    context.Generator.EmitLdarg(TranslatedSub.StateArgIdx);
+                    context.Generator.Emit(OpCodes.Ldfld, reg.GetField());
+
+                    context.Generator.EmitStloc(context.GetLocalIndex(reg));
+                }
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/ChocolArm64/Translation/ILOpCodeLog.cs b/ChocolArm64/Translation/ILOpCodeLog.cs
index 2c77021b..ebb042b5 100644
--- a/ChocolArm64/Translation/ILOpCodeLog.cs
+++ b/ChocolArm64/Translation/ILOpCodeLog.cs
@@ -9,7 +9,7 @@ namespace ChocolArm64.Translation
             _text = text;
         }
 
-        public void Emit(ILEmitter context)
+        public void Emit(ILMethodBuilder context)
         {
             context.Generator.EmitWriteLine(_text);
         }
diff --git a/ChocolArm64/Translation/ILOpCodeStore.cs b/ChocolArm64/Translation/ILOpCodeStore.cs
index 41326fe1..17a6259c 100644
--- a/ChocolArm64/Translation/ILOpCodeStore.cs
+++ b/ChocolArm64/Translation/ILOpCodeStore.cs
@@ -3,7 +3,7 @@ using System.Reflection.Emit;
 
 namespace ChocolArm64.Translation
 {
-    struct IlOpCodeStore : IILEmit
+    struct ILOpCodeStore : IILEmit
     {
         public int Index { get; private set; }
 
@@ -11,55 +11,26 @@ namespace ChocolArm64.Translation
 
         public RegisterSize RegisterSize { get; private set; }
 
-        public IlOpCodeStore(int index, IoType ioType, RegisterSize registerSize = 0)
+        public ILOpCodeStore(int index, IoType ioType, RegisterSize registerSize = 0)
         {
             Index        = index;
             IoType       = ioType;
             RegisterSize = registerSize;
         }
 
-        public void Emit(ILEmitter context)
+        public void Emit(ILMethodBuilder context)
         {
             switch (IoType)
             {
                 case IoType.Arg: context.Generator.EmitStarg(Index); break;
 
-                case IoType.Fields:
-                {
-                    long intOutputs = context.LocalAlloc.GetIntOutputs(context.GetIlBlock(Index));
-                    long vecOutputs = context.LocalAlloc.GetVecOutputs(context.GetIlBlock(Index));
-
-                    StoreLocals(context, intOutputs, RegisterType.Int);
-                    StoreLocals(context, vecOutputs, RegisterType.Vector);
-                    
-                    break;
-                }
-
                 case IoType.Flag:   EmitStloc(context, Index, RegisterType.Flag);   break;
                 case IoType.Int:    EmitStloc(context, Index, RegisterType.Int);    break;
                 case IoType.Vector: EmitStloc(context, Index, RegisterType.Vector); break;
             }
         }
 
-        private void StoreLocals(ILEmitter context, long outputs, RegisterType baseType)
-        {
-            for (int bit = 0; bit < 64; bit++)
-            {
-                long mask = 1L << bit;
-
-                if ((outputs & mask) != 0)
-                {
-                    Register reg = ILEmitter.GetRegFromBit(bit, baseType);
-
-                    context.Generator.EmitLdarg(TranslatedSub.StateArgIdx);
-                    context.Generator.EmitLdloc(context.GetLocalIndex(reg));
-
-                    context.Generator.Emit(OpCodes.Stfld, reg.GetField());
-                }
-            }
-        }
-
-        private void EmitStloc(ILEmitter context, int index, RegisterType registerType)
+        private void EmitStloc(ILMethodBuilder context, int index, RegisterType registerType)
         {
             Register reg = new Register(index, registerType);
 
diff --git a/ChocolArm64/Translation/ILOpCodeStoreState.cs b/ChocolArm64/Translation/ILOpCodeStoreState.cs
new file mode 100644
index 00000000..458e9eda
--- /dev/null
+++ b/ChocolArm64/Translation/ILOpCodeStoreState.cs
@@ -0,0 +1,42 @@
+using ChocolArm64.State;
+using System.Reflection.Emit;
+
+namespace ChocolArm64.Translation
+{
+    struct ILOpCodeStoreState : IILEmit
+    {
+        private ILBlock _block;
+
+        public ILOpCodeStoreState(ILBlock block)
+        {
+            _block = block;
+        }
+
+        public void Emit(ILMethodBuilder context)
+        {
+            long intOutputs = context.LocalAlloc.GetIntOutputs(_block);
+            long vecOutputs = context.LocalAlloc.GetVecOutputs(_block);
+
+            StoreLocals(context, intOutputs, RegisterType.Int);
+            StoreLocals(context, vecOutputs, RegisterType.Vector);
+        }
+
+        private void StoreLocals(ILMethodBuilder context, long outputs, RegisterType baseType)
+        {
+            for (int bit = 0; bit < 64; bit++)
+            {
+                long mask = 1L << bit;
+
+                if ((outputs & mask) != 0)
+                {
+                    Register reg = ILMethodBuilder.GetRegFromBit(bit, baseType);
+
+                    context.Generator.EmitLdarg(TranslatedSub.StateArgIdx);
+                    context.Generator.EmitLdloc(context.GetLocalIndex(reg));
+
+                    context.Generator.Emit(OpCodes.Stfld, reg.GetField());
+                }
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/ChocolArm64/Translation/IoType.cs b/ChocolArm64/Translation/IoType.cs
index 290b159f..c7710e0c 100644
--- a/ChocolArm64/Translation/IoType.cs
+++ b/ChocolArm64/Translation/IoType.cs
@@ -1,15 +1,10 @@
-using System;
-
 namespace ChocolArm64.Translation
 {
-    [Flags]
     enum IoType
     {
         Arg,
-        Fields,
         Flag,
         Int,
-        Float,
         Vector
     }
 }
\ No newline at end of file
diff --git a/ChocolArm64/Translation/LocalAlloc.cs b/ChocolArm64/Translation/LocalAlloc.cs
index 8237bd54..763be619 100644
--- a/ChocolArm64/Translation/LocalAlloc.cs
+++ b/ChocolArm64/Translation/LocalAlloc.cs
@@ -1,3 +1,4 @@
+using System;
 using System.Collections.Generic;
 
 namespace ChocolArm64.Translation
@@ -65,11 +66,41 @@ namespace ChocolArm64.Translation
             public long VecInputs;
             public long IntOutputs;
             public long VecOutputs;
+
+            public override bool Equals(object obj)
+            {
+                if (!(obj is BlockIo other))
+                {
+                    return false;
+                }
+
+                return other.Block      == Block      &&
+                       other.Entry      == Entry      &&
+                       other.IntInputs  == IntInputs  &&
+                       other.VecInputs  == VecInputs  &&
+                       other.IntOutputs == IntOutputs &&
+                       other.VecOutputs == VecOutputs;
+            }
+
+            public override int GetHashCode()
+            {
+                return HashCode.Combine(Block, Entry, IntInputs, VecInputs, IntOutputs, VecOutputs);
+            }
+
+            public static bool operator ==(BlockIo lhs, BlockIo rhs)
+            {
+                return lhs.Equals(rhs);
+            }
+
+            public static bool operator !=(BlockIo lhs, BlockIo rhs)
+            {
+                return !(lhs == rhs);
+            }
         }
 
         private const int MaxOptGraphLength = 40;
 
-        public LocalAlloc(ILBlock[] graph, ILBlock root)
+        public LocalAlloc(ILBlock[] graph, ILBlock entry)
         {
             _intPaths = new Dictionary<ILBlock, PathIo>();
             _vecPaths = new Dictionary<ILBlock, PathIo>();
@@ -77,7 +108,7 @@ namespace ChocolArm64.Translation
             if (graph.Length > 1 &&
                 graph.Length < MaxOptGraphLength)
             {
-                InitializeOptimal(graph, root);
+                InitializeOptimal(graph, entry);
             }
             else
             {
@@ -85,7 +116,7 @@ namespace ChocolArm64.Translation
             }
         }
 
-        private void InitializeOptimal(ILBlock[] graph, ILBlock root)
+        private void InitializeOptimal(ILBlock[] graph, ILBlock entry)
         {
             //This will go through all possible paths on the graph,
             //and store all inputs/outputs for each block. A register
@@ -93,7 +124,7 @@ namespace ChocolArm64.Translation
             //When a block can be reached by more than one path, then the
             //output from all paths needs to be set for this block, and
             //only outputs present in all of the parent blocks can be considered
-            //when doing input elimination. Each block chain have a root, that's where
+            //when doing input elimination. Each block chain have a entry, that's where
             //the code starts executing. They are present on the subroutine start point,
             //and on call return points too (address written to X30 by BL).
             HashSet<BlockIo> visited = new HashSet<BlockIo>();
@@ -112,8 +143,8 @@ namespace ChocolArm64.Translation
 
             Enqueue(new BlockIo()
             {
-                Block = root,
-                Entry = root
+                Block = entry,
+                Entry = entry
             });
 
             while (unvisited.Count > 0)
@@ -146,22 +177,22 @@ namespace ChocolArm64.Translation
 
                 void EnqueueFromCurrent(ILBlock block, bool retTarget)
                 {
-                    BlockIo blkIO = new BlockIo() { Block = block };
+                    BlockIo blockIo = new BlockIo() { Block = block };
 
                     if (retTarget)
                     {
-                        blkIO.Entry = block;
+                        blockIo.Entry = block;
                     }
                     else
                     {
-                        blkIO.Entry      = current.Entry;
-                        blkIO.IntInputs  = current.IntInputs;
-                        blkIO.VecInputs  = current.VecInputs;
-                        blkIO.IntOutputs = current.IntOutputs;
-                        blkIO.VecOutputs = current.VecOutputs;
+                        blockIo.Entry      = current.Entry;
+                        blockIo.IntInputs  = current.IntInputs;
+                        blockIo.VecInputs  = current.VecInputs;
+                        blockIo.IntOutputs = current.IntOutputs;
+                        blockIo.VecOutputs = current.VecOutputs;
                     }
 
-                    Enqueue(blkIO);
+                    Enqueue(blockIo);
                 }
 
                 if (current.Block.Next != null)
@@ -179,7 +210,7 @@ namespace ChocolArm64.Translation
         private void InitializeFast(ILBlock[] graph)
         {
             //This is WAY faster than InitializeOptimal, but results in
-            //uneeded loads and stores, so the resulting code will be slower.
+            //unneeded loads and stores, so the resulting code will be slower.
             long intInputs = 0, intOutputs = 0;
             long vecInputs = 0, vecOutputs = 0;