[CPU] Speed up translation a little bit

4 files changed, 48 insertions, 33 deletions

diff --git a/ChocolArm64/ATranslatedSub.cs b/ChocolArm64/ATranslatedSub.cs
index 414038ab..9dbc378e 100644
--- a/ChocolArm64/ATranslatedSub.cs
+++ b/ChocolArm64/ATranslatedSub.cs
@@ -3,6 +3,7 @@ using ChocolArm64.State;
 using System;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
+using System.Linq;
 using System.Reflection;
 using System.Reflection.Emit;
 
@@ -23,7 +24,7 @@ namespace ChocolArm64
 
         public ReadOnlyCollection<ARegister> Params { get; private set; }
 
-        private HashSet<long> Callees;
+        private HashSet<long> Callers;
 
         private ATranslatedSubType Type;
 
@@ -33,7 +34,7 @@ namespace ChocolArm64
 
         private int MinCallCountForReJit = 250;
 
-        public ATranslatedSub(DynamicMethod Method, List<ARegister> Params, HashSet<long> Callees)
+        public ATranslatedSub(DynamicMethod Method, List<ARegister> Params)
         {
             if (Method == null)
             {
@@ -45,14 +46,10 @@ namespace ChocolArm64
                 throw new ArgumentNullException(nameof(Params));
             }
 
-            if (Callees == null)
-            {
-                throw new ArgumentNullException(nameof(Callees));
-            }
-
             this.Method  = Method;
             this.Params  = Params.AsReadOnly();
-            this.Callees = Callees;
+
+            Callers = new HashSet<long>();
 
             PrepareDelegate();
         }
@@ -107,17 +104,14 @@ namespace ChocolArm64
 
         public bool ShouldReJit()
         {
-            if (Type == ATranslatedSubType.SubTier0)
+            if (NeedsReJit && CallCount < MinCallCountForReJit)
             {
-                if (CallCount < MinCallCountForReJit)
-                {
-                    CallCount++;
-                }
+                CallCount++;
 
-                return CallCount == MinCallCountForReJit;
+                return false;
             }
 
-            return Type == ATranslatedSubType.SubTier1 && NeedsReJit;
+            return NeedsReJit;
         }
 
         public long Execute(AThreadState ThreadState, AMemory Memory)
@@ -125,10 +119,32 @@ namespace ChocolArm64
             return ExecDelegate(ThreadState, Memory);
         }
 
-        public void SetType(ATranslatedSubType Type) => this.Type = Type;
+        public void AddCaller(long Position)
+        {
+            lock (Callers)
+            {
+                Callers.Add(Position);
+            }
+        }
+
+        public long[] GetCallerPositions()
+        {
+            lock (Callers)
+            {
+                return Callers.ToArray();
+            }
+        }
 
-        public bool HasCallee(long Position) => Callees.Contains(Position);
+        public void SetType(ATranslatedSubType Type)
+        {
+            this.Type = Type;
+
+            if (Type == ATranslatedSubType.SubTier0)
+            {
+                NeedsReJit = true;
+            }
+        }
 
-        public void MarkForReJit() => NeedsReJit = true;        
+        public void MarkForReJit() => NeedsReJit = true;
     }
 }
\ No newline at end of file
diff --git a/ChocolArm64/ATranslator.cs b/ChocolArm64/ATranslator.cs
index f1bc2cff..e46750fc 100644
--- a/ChocolArm64/ATranslator.cs
+++ b/ChocolArm64/ATranslator.cs
@@ -160,11 +160,14 @@ namespace ChocolArm64
 
             //Mark all methods that calls this method for ReJiting,
             //since we can now call it directly which is faster.
-            foreach (ATranslatedSub TS in CachedSubs.Values)
+            if (CachedSubs.TryGetValue(Position, out ATranslatedSub OldSub))
             {
-                if (TS.HasCallee(Position))
+                foreach (long CallerPos in OldSub.GetCallerPositions())
                 {
-                    TS.MarkForReJit();
+                    if (CachedSubs.TryGetValue(Position, out ATranslatedSub CallerSub))
+                    {
+                        CallerSub.MarkForReJit();
+                    }
                 }
             }
 
diff --git a/ChocolArm64/Translation/AILEmitter.cs b/ChocolArm64/Translation/AILEmitter.cs
index af37a6c7..55b1751f 100644
--- a/ChocolArm64/Translation/AILEmitter.cs
+++ b/ChocolArm64/Translation/AILEmitter.cs
@@ -60,11 +60,11 @@ namespace ChocolArm64.Translation
 
         public AILBlock GetILBlock(int Index) => ILBlocks[Index];
 
-        public ATranslatedSub GetSubroutine(HashSet<long> Callees)
+        public ATranslatedSub GetSubroutine()
         {
             LocalAlloc = new ALocalAlloc(ILBlocks, Root);
 
-            InitSubroutine(Callees);
+            InitSubroutine();
             InitLocals();
 
             foreach (AILBlock ILBlock in ILBlocks)
@@ -75,7 +75,7 @@ namespace ChocolArm64.Translation
             return Subroutine;
         }
 
-        private void InitSubroutine(HashSet<long> Callees)
+        private void InitSubroutine()
         {
             List<ARegister> Params = new List<ARegister>();
 
@@ -99,7 +99,7 @@ namespace ChocolArm64.Translation
 
             Generator = Mthd.GetILGenerator();
 
-            Subroutine = new ATranslatedSub(Mthd, Params, Callees);
+            Subroutine = new ATranslatedSub(Mthd, Params);
         }
 
         private void InitLocals()
@@ -115,7 +115,7 @@ namespace ChocolArm64.Translation
                 Generator.EmitLdarg(Index + ParamsStart);
                 Generator.EmitStloc(GetLocalIndex(Reg));
             }
-        }        
+        }
 
         private Type[] GetParamTypes(IList<ARegister> Params)
         {
diff --git a/ChocolArm64/Translation/AILEmitterCtx.cs b/ChocolArm64/Translation/AILEmitterCtx.cs
index 03b06610..2f4a67e1 100644
--- a/ChocolArm64/Translation/AILEmitterCtx.cs
+++ b/ChocolArm64/Translation/AILEmitterCtx.cs
@@ -12,8 +12,6 @@ namespace ChocolArm64.Translation
     {
         private ATranslator Translator;
 
-        private HashSet<long> Callees;
-
         private Dictionary<long, AILLabel> Labels;
 
         private int BlkIndex;
@@ -66,8 +64,6 @@ namespace ChocolArm64.Translation
             this.Graph      = Graph;
             this.Root       = Root;
 
-            Callees = new HashSet<long>();
-
             Labels = new Dictionary<long, AILLabel>();
 
             Emitter = new AILEmitter(Graph, Root, SubName);
@@ -84,7 +80,7 @@ namespace ChocolArm64.Translation
 
         public ATranslatedSub GetSubroutine()
         {
-            return Emitter.GetSubroutine(Callees);
+            return Emitter.GetSubroutine();
         }
 
         public bool AdvanceOpCode()
@@ -123,8 +119,6 @@ namespace ChocolArm64.Translation
 
         public bool TryOptEmitSubroutineCall()
         {
-            Callees.Add(((AOpCodeBImm)CurrOp).Imm);
-
             if (CurrBlock.Next == null)
             {
                 return false;
@@ -152,6 +146,8 @@ namespace ChocolArm64.Translation
 
             EmitCall(Sub.Method);
 
+            Sub.AddCaller(Root.Position);
+
             return true;
         }