Clear JIT cache on exit (#1518)

* Initial cache memory allocator implementation

* Get rid of CallFlag

* Perform cache cleanup on exit

* Basic cache invalidation

* Thats not how conditionals works in C# it seems

* Set PTC version to PR number

* Address PR feedback

* Update InstEmitFlowHelper.cs

* Flag clear on address is no longer needed

* Do not include exit block in function size calculation

* Dispose jump table

* For future use

* InternalVersion = 1519 (force retest).

Co-authored-by: LDj3SNuD <35856442+LDj3SNuD@users.noreply.github.com>

2 files changed, 94 insertions, 56 deletions

diff --git a/ARMeilleure/Translation/PTC/Ptc.cs b/ARMeilleure/Translation/PTC/Ptc.cs
index 05dd352e..1cdaa8fe 100644
--- a/ARMeilleure/Translation/PTC/Ptc.cs
+++ b/ARMeilleure/Translation/PTC/Ptc.cs
@@ -2,6 +2,7 @@ using ARMeilleure.CodeGen;
 using ARMeilleure.CodeGen.Unwinding;
 using ARMeilleure.CodeGen.X86;
 using ARMeilleure.Memory;
+using ARMeilleure.Translation.Cache;
 using Ryujinx.Common.Configuration;
 using Ryujinx.Common.Logging;
 using System;
@@ -21,7 +22,7 @@ namespace ARMeilleure.Translation.PTC
     {
         private const string HeaderMagic = "PTChd";
 
-        private const int InternalVersion = 1801; //! To be incremented manually for each change to the ARMeilleure project.
+        private const int InternalVersion = 1519; //! To be incremented manually for each change to the ARMeilleure project.
 
         private const string ActualDir = "0";
         private const string BackupDir = "1";
@@ -433,9 +434,9 @@ namespace ARMeilleure.Translation.PTC
 
                     UnwindInfo unwindInfo = ReadUnwindInfo(unwindInfosReader);
 
-                    TranslatedFunction func = FastTranslate(code, unwindInfo, infoEntry.HighCq);
+                    TranslatedFunction func = FastTranslate(code, infoEntry.GuestSize, unwindInfo, infoEntry.HighCq);
 
-                    funcs.AddOrUpdate((ulong)infoEntry.Address, func, (key, oldFunc) => func.HighCq && !oldFunc.HighCq ? func : oldFunc);
+                    funcs.AddOrUpdate(infoEntry.Address, func, (key, oldFunc) => func.HighCq && !oldFunc.HighCq ? func : oldFunc);
                 }
             }
 
@@ -457,7 +458,8 @@ namespace ARMeilleure.Translation.PTC
         {
             InfoEntry infoEntry = new InfoEntry();
 
-            infoEntry.Address = infosReader.ReadInt64();
+            infoEntry.Address = infosReader.ReadUInt64();
+            infoEntry.GuestSize = infosReader.ReadUInt64();
             infoEntry.HighCq = infosReader.ReadBoolean();
             infoEntry.CodeLen = infosReader.ReadInt32();
             infoEntry.RelocEntriesCount = infosReader.ReadInt32();
@@ -541,7 +543,7 @@ namespace ARMeilleure.Translation.PTC
             return new UnwindInfo(pushEntries, prologueSize);
         }
 
-        private static TranslatedFunction FastTranslate(byte[] code, UnwindInfo unwindInfo, bool highCq)
+        private static TranslatedFunction FastTranslate(byte[] code, ulong guestSize, UnwindInfo unwindInfo, bool highCq)
         {
             CompiledFunction cFunc = new CompiledFunction(code, unwindInfo);
 
@@ -549,7 +551,7 @@ namespace ARMeilleure.Translation.PTC
 
             GuestFunction gFunc = Marshal.GetDelegateForFunctionPointer<GuestFunction>(codePtr);
 
-            TranslatedFunction tFunc = new TranslatedFunction(gFunc, highCq);
+            TranslatedFunction tFunc = new TranslatedFunction(gFunc, guestSize, highCq);
 
             return tFunc;
         }
@@ -632,12 +634,13 @@ namespace ARMeilleure.Translation.PTC
             Logger.Info?.Print(LogClass.Ptc, $"{funcsCount + _translateCount} of {ProfiledFuncsCount} functions to translate - {_rejitCount} functions rejited");
         }
 
-        internal static void WriteInfoCodeReloc(long address, bool highCq, PtcInfo ptcInfo)
+        internal static void WriteInfoCodeReloc(ulong address, ulong guestSize, bool highCq, PtcInfo ptcInfo)
         {
             lock (_lock)
             {
                 // WriteInfo.
-                _infosWriter.Write((long)address); // InfoEntry.Address
+                _infosWriter.Write((ulong)address); // InfoEntry.Address
+                _infosWriter.Write((ulong)guestSize); // InfoEntry.GuestSize
                 _infosWriter.Write((bool)highCq); // InfoEntry.HighCq
                 _infosWriter.Write((int)ptcInfo.CodeStream.Length); // InfoEntry.CodeLen
                 _infosWriter.Write((int)ptcInfo.RelocEntriesCount); // InfoEntry.RelocEntriesCount
@@ -673,9 +676,10 @@ namespace ARMeilleure.Translation.PTC
 
         private struct InfoEntry
         {
-            public const int Stride = 17; // Bytes.
+            public const int Stride = 25; // Bytes.
 
-            public long Address;
+            public ulong Address;
+            public ulong GuestSize;
             public bool HighCq;
             public int CodeLen;
             public int RelocEntriesCount;
diff --git a/ARMeilleure/Translation/PTC/PtcJumpTable.cs b/ARMeilleure/Translation/PTC/PtcJumpTable.cs
index 0a3ae240..e4af68d6 100644
--- a/ARMeilleure/Translation/PTC/PtcJumpTable.cs
+++ b/ARMeilleure/Translation/PTC/PtcJumpTable.cs
@@ -1,3 +1,4 @@
+using ARMeilleure.Translation.Cache;
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
@@ -8,25 +9,53 @@ namespace ARMeilleure.Translation.PTC
     [Serializable]
     class PtcJumpTable
     {
-        private readonly List<KeyValuePair<long, DirectHostAddress>>   _jumpTable;
-        private readonly List<KeyValuePair<long, IndirectHostAddress>> _dynamicTable;
+        [Serializable]
+        private struct TableEntry<TAddress>
+        {
+            public int EntryIndex;
+            public long GuestAddress;
+            public TAddress HostAddress;
 
-        private readonly List<ulong> _targets;
-        private readonly Dictionary<ulong, LinkedList<int>> _dependants;
+            public TableEntry(int entryIndex, long guestAddress, TAddress hostAddress)
+            {
+                EntryIndex = entryIndex;
+                GuestAddress = guestAddress;
+                HostAddress = hostAddress;
+            }
+        }
+
+        private enum DirectHostAddress
+        {
+            CallStub,
+            TailCallStub,
+            Host
+        }
+
+        private enum IndirectHostAddress
+        {
+            CallStub,
+            TailCallStub
+        }
+
+        private readonly List<TableEntry<DirectHostAddress>> _jumpTable;
+        private readonly List<TableEntry<IndirectHostAddress>> _dynamicTable;
 
-        public int TableEnd    => _jumpTable.Count;
-        public int DynTableEnd => _dynamicTable.Count;
+        private readonly List<ulong> _targets;
+        private readonly Dictionary<ulong, List<int>> _dependants;
+        private readonly Dictionary<ulong, List<int>> _owners;
 
         public List<ulong> Targets => _targets;
-        public Dictionary<ulong, LinkedList<int>> Dependants => _dependants;
+        public Dictionary<ulong, List<int>> Dependants => _dependants;
+        public Dictionary<ulong, List<int>> Owners => _owners;
 
         public PtcJumpTable()
         {
-            _jumpTable    = new List<KeyValuePair<long, DirectHostAddress>>();
-            _dynamicTable = new List<KeyValuePair<long, IndirectHostAddress>>();
+            _jumpTable = new List<TableEntry<DirectHostAddress>>();
+            _dynamicTable = new List<TableEntry<IndirectHostAddress>>();
 
-            _targets    = new List<ulong>();
-            _dependants = new Dictionary<ulong, LinkedList<int>>();
+            _targets = new List<ulong>();
+            _dependants = new Dictionary<ulong, List<int>>();
+            _owners = new Dictionary<ulong, List<int>>();
         }
 
         public void Initialize(JumpTable jumpTable)
@@ -42,7 +71,14 @@ namespace ARMeilleure.Translation.PTC
 
             foreach (var item in jumpTable.Dependants)
             {
-                _dependants.Add(item.Key, new LinkedList<int>(item.Value));
+                _dependants.Add(item.Key, new List<int>(item.Value));
+            }
+
+            _owners.Clear();
+
+            foreach (var item in jumpTable.Owners)
+            {
+                _owners.Add(item.Key, new List<int>(item.Value));
             }
         }
 
@@ -53,20 +89,17 @@ namespace ARMeilleure.Translation.PTC
 
             _targets.Clear();
             _dependants.Clear();
+            _owners.Clear();
         }
 
         public void WriteJumpTable(JumpTable jumpTable, ConcurrentDictionary<ulong, TranslatedFunction> funcs)
         {
-            jumpTable.ExpandIfNeededJumpTable(TableEnd);
-
-            int entry = 0;
+            // Writes internal state to jump table in-memory, after PTC was loaded.
 
             foreach (var item in _jumpTable)
             {
-                entry += 1;
-
-                long guestAddress = item.Key;
-                DirectHostAddress directHostAddress = item.Value;
+                long guestAddress = item.GuestAddress;
+                DirectHostAddress directHostAddress = item.HostAddress;
 
                 long hostAddress;
 
@@ -94,6 +127,11 @@ namespace ARMeilleure.Translation.PTC
                     throw new InvalidOperationException(nameof(directHostAddress));
                 }
 
+                int entry = item.EntryIndex;
+
+                jumpTable.Table.SetEntry(entry);
+                jumpTable.ExpandIfNeededJumpTable(entry);
+
                 IntPtr addr = jumpTable.GetEntryAddressJumpTable(entry);
 
                 Marshal.WriteInt64(addr, 0, guestAddress);
@@ -103,21 +141,17 @@ namespace ARMeilleure.Translation.PTC
 
         public void WriteDynamicTable(JumpTable jumpTable)
         {
+            // Writes internal state to jump table in-memory, after PTC was loaded.
+
             if (JumpTable.DynamicTableElems > 1)
             {
                 throw new NotSupportedException();
             }
 
-            jumpTable.ExpandIfNeededDynamicTable(DynTableEnd);
-
-            int entry = 0;
-
             foreach (var item in _dynamicTable)
             {
-                entry += 1;
-
-                long guestAddress = item.Key;
-                IndirectHostAddress indirectHostAddress = item.Value;
+                long guestAddress = item.GuestAddress;
+                IndirectHostAddress indirectHostAddress = item.HostAddress;
 
                 long hostAddress;
 
@@ -134,6 +168,11 @@ namespace ARMeilleure.Translation.PTC
                     throw new InvalidOperationException(nameof(indirectHostAddress));
                 }
 
+                int entry = item.EntryIndex;
+
+                jumpTable.DynTable.SetEntry(entry);
+                jumpTable.ExpandIfNeededDynamicTable(entry);
+
                 IntPtr addr = jumpTable.GetEntryAddressDynamicTable(entry);
 
                 Marshal.WriteInt64(addr, 0, guestAddress);
@@ -143,14 +182,18 @@ namespace ARMeilleure.Translation.PTC
 
         public void ReadJumpTable(JumpTable jumpTable)
         {
+            // Reads in-memory jump table state and store internally for PTC serialization.
+
             _jumpTable.Clear();
 
-            for (int entry = 1; entry <= jumpTable.TableEnd; entry++)
+            IEnumerable<int> entries = jumpTable.Table.GetEntries();
+
+            foreach (int entry in entries)
             {
                 IntPtr addr = jumpTable.GetEntryAddressJumpTable(entry);
 
                 long guestAddress = Marshal.ReadInt64(addr, 0);
-                long hostAddress  = Marshal.ReadInt64(addr, 8);
+                long hostAddress = Marshal.ReadInt64(addr, 8);
 
                 DirectHostAddress directHostAddress;
 
@@ -167,12 +210,14 @@ namespace ARMeilleure.Translation.PTC
                     directHostAddress = DirectHostAddress.Host;
                 }
 
-                _jumpTable.Add(new KeyValuePair<long, DirectHostAddress>(guestAddress, directHostAddress));
+                _jumpTable.Add(new TableEntry<DirectHostAddress>(entry, guestAddress, directHostAddress));
             }
         }
 
         public void ReadDynamicTable(JumpTable jumpTable)
         {
+            // Reads in-memory jump table state and store internally for PTC serialization.
+
             if (JumpTable.DynamicTableElems > 1)
             {
                 throw new NotSupportedException();
@@ -180,12 +225,14 @@ namespace ARMeilleure.Translation.PTC
 
             _dynamicTable.Clear();
 
-            for (int entry = 1; entry <= jumpTable.DynTableEnd; entry++)
+            IEnumerable<int> entries = jumpTable.DynTable.GetEntries();
+
+            foreach (int entry in entries)
             {
                 IntPtr addr = jumpTable.GetEntryAddressDynamicTable(entry);
 
                 long guestAddress = Marshal.ReadInt64(addr, 0);
-                long hostAddress  = Marshal.ReadInt64(addr, 8);
+                long hostAddress = Marshal.ReadInt64(addr, 8);
 
                 IndirectHostAddress indirectHostAddress;
 
@@ -202,21 +249,8 @@ namespace ARMeilleure.Translation.PTC
                     throw new InvalidOperationException($"({nameof(hostAddress)} = 0x{hostAddress:X16})");
                 }
 
-                _dynamicTable.Add(new KeyValuePair<long, IndirectHostAddress>(guestAddress, indirectHostAddress));
+                _dynamicTable.Add(new TableEntry<IndirectHostAddress>(entry, guestAddress, indirectHostAddress));
             }
         }
-
-        private enum DirectHostAddress
-        {
-            CallStub,
-            TailCallStub,
-            Host
-        }
-
-        private enum IndirectHostAddress
-        {
-            CallStub,
-            TailCallStub
-        }
     }
 }
\ No newline at end of file