1 files changed, 27 insertions, 8 deletions
diff --git a/ARMeilleure/Translation/Translator.cs b/ARMeilleure/Translation/Translator.cs
index 4962a846..611716e2 100644
--- a/ARMeilleure/Translation/Translator.cs
+++ b/ARMeilleure/Translation/Translator.cs
@@ -49,7 +49,7 @@ namespace ARMeilleure.Translation
         private readonly AutoResetEvent _backgroundTranslatorEvent;
         private readonly ReaderWriterLock _backgroundTranslatorLock;
 
-        internal ConcurrentDictionary<ulong, TranslatedFunction> Functions { get; }
+        internal TranslatorCache<TranslatedFunction> Functions { get; }
         internal AddressTable<ulong> FunctionTable { get; }
         internal EntryTable<uint> CountTable { get; }
         internal TranslatorStubs Stubs { get; }
@@ -75,7 +75,7 @@ namespace ARMeilleure.Translation
             JitCache.Initialize(allocator);
 
             CountTable = new EntryTable<uint>();
-            Functions = new ConcurrentDictionary<ulong, TranslatedFunction>();
+            Functions = new TranslatorCache<TranslatedFunction>();
             FunctionTable = new AddressTable<ulong>(for64Bits ? Levels64Bit : Levels32Bit);
             Stubs = new TranslatorStubs(this);
 
@@ -93,12 +93,12 @@ namespace ARMeilleure.Translation
             {
                 _backgroundTranslatorLock.AcquireReaderLock(Timeout.Infinite);
 
-                if (_backgroundStack.TryPop(out RejitRequest request) && 
+                if (_backgroundStack.TryPop(out RejitRequest request) &&
                     _backgroundSet.TryRemove(request.Address, out _))
                 {
                     TranslatedFunction func = Translate(request.Address, request.Mode, highCq: true);
 
-                    Functions.AddOrUpdate(request.Address, func, (key, oldFunc) =>
+                    Functions.AddOrUpdate(request.Address, func.GuestSize, func, (key, oldFunc) =>
                     {
                         EnqueueForDeletion(key, oldFunc);
                         return func;
@@ -196,7 +196,7 @@ namespace ARMeilleure.Translation
             }
         }
 
-        public ulong ExecuteSingle(State.ExecutionContext context, ulong address)
+        private ulong ExecuteSingle(State.ExecutionContext context, ulong address)
         {
             TranslatedFunction func = GetOrTranslate(address, context.ExecutionMode);
 
@@ -215,7 +215,7 @@ namespace ARMeilleure.Translation
             {
                 func = Translate(address, mode, highCq: false);
 
-                TranslatedFunction oldFunc = Functions.GetOrAdd(address, func);
+                TranslatedFunction oldFunc = Functions.GetOrAdd(address, func.GuestSize, func);
 
                 if (oldFunc != func)
                 {
@@ -471,7 +471,24 @@ namespace ARMeilleure.Translation
             // If rejit is running, stop it as it may be trying to rejit a function on the invalidated region.
             ClearRejitQueue(allowRequeue: true);
 
-            // TODO: Completely remove functions overlapping the specified range from the cache.
+            ulong[] overlapAddresses = Array.Empty<ulong>();
+
+            int overlapsCount = Functions.GetOverlaps(address, size, ref overlapAddresses);
+
+            for (int index = 0; index < overlapsCount; index++)
+            {
+                ulong overlapAddress = overlapAddresses[index];
+
+                if (Functions.TryGetValue(overlapAddress, out TranslatedFunction overlap))
+                {
+                    Functions.Remove(overlapAddress);
+                    Volatile.Write(ref FunctionTable.GetValue(overlapAddress), FunctionTable.Fill);
+                    EnqueueForDeletion(overlapAddress, overlap);
+                }
+            }
+
+            // TODO: Remove overlapping functions from the JitCache aswell.
+            // This should be done safely, with a mechanism to ensure the function is not being executed.
         }
 
         internal void EnqueueForRejit(ulong guestAddress, ExecutionMode mode)
@@ -493,7 +510,9 @@ namespace ARMeilleure.Translation
             // Ensure no attempt will be made to compile new functions due to rejit.
             ClearRejitQueue(allowRequeue: false);
 
-            foreach (var func in Functions.Values)
+            List<TranslatedFunction> functions = Functions.AsList();
+
+            foreach (var func in functions)
             {
                 JitCache.Unmap(func.FuncPtr);