1 files changed, 198 insertions, 0 deletions
diff --git a/src/ARMeilleure/Translation/Cache/JitCache.cs b/src/ARMeilleure/Translation/Cache/JitCache.cs
new file mode 100644
index 00000000..f496a8e9
--- /dev/null
+++ b/src/ARMeilleure/Translation/Cache/JitCache.cs
@@ -0,0 +1,198 @@
+using ARMeilleure.CodeGen;
+using ARMeilleure.CodeGen.Unwinding;
+using ARMeilleure.Memory;
+using ARMeilleure.Native;
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.Runtime.InteropServices;
+
+namespace ARMeilleure.Translation.Cache
+{
+    static class JitCache
+    {
+        private const int PageSize = 4 * 1024;
+        private const int PageMask = PageSize - 1;
+
+        private const int CodeAlignment = 4; // Bytes.
+        private const int CacheSize = 2047 * 1024 * 1024;
+
+        private static ReservedRegion _jitRegion;
+        private static JitCacheInvalidation _jitCacheInvalidator;
+
+        private static CacheMemoryAllocator _cacheAllocator;
+
+        private static readonly List<CacheEntry> _cacheEntries = new List<CacheEntry>();
+
+        private static readonly object _lock = new object();
+        private static bool _initialized;
+
+        public static void Initialize(IJitMemoryAllocator allocator)
+        {
+            if (_initialized) return;
+
+            lock (_lock)
+            {
+                if (_initialized) return;
+
+                _jitRegion = new ReservedRegion(allocator, CacheSize);
+                _jitCacheInvalidator = new JitCacheInvalidation(allocator);
+
+                _cacheAllocator = new CacheMemoryAllocator(CacheSize);
+
+                if (OperatingSystem.IsWindows())
+                {
+                    JitUnwindWindows.InstallFunctionTableHandler(_jitRegion.Pointer, CacheSize, _jitRegion.Pointer + Allocate(PageSize));
+                }
+
+                _initialized = true;
+            }
+        }
+
+        public static IntPtr Map(CompiledFunction func)
+        {
+            byte[] code = func.Code;
+
+            lock (_lock)
+            {
+                Debug.Assert(_initialized);
+
+                int funcOffset = Allocate(code.Length);
+
+                IntPtr funcPtr = _jitRegion.Pointer + funcOffset;
+
+                if (OperatingSystem.IsMacOS() && RuntimeInformation.ProcessArchitecture == Architecture.Arm64)
+                {
+                    unsafe
+                    {
+                        fixed (byte *codePtr = code)
+                        {
+                            JitSupportDarwin.Copy(funcPtr, (IntPtr)codePtr, (ulong)code.Length);
+                        }
+                    }
+                }
+                else
+                {
+                    ReprotectAsWritable(funcOffset, code.Length);
+                    Marshal.Copy(code, 0, funcPtr, code.Length);
+                    ReprotectAsExecutable(funcOffset, code.Length);
+
+                    _jitCacheInvalidator.Invalidate(funcPtr, (ulong)code.Length);
+                }
+
+                Add(funcOffset, code.Length, func.UnwindInfo);
+
+                return funcPtr;
+            }
+        }
+
+        public static void Unmap(IntPtr pointer)
+        {
+            lock (_lock)
+            {
+                Debug.Assert(_initialized);
+
+                int funcOffset = (int)(pointer.ToInt64() - _jitRegion.Pointer.ToInt64());
+
+                bool result = TryFind(funcOffset, out CacheEntry entry);
+                Debug.Assert(result);
+
+                _cacheAllocator.Free(funcOffset, AlignCodeSize(entry.Size));
+
+                Remove(funcOffset);
+            }
+        }
+
+        private static void ReprotectAsWritable(int offset, int size)
+        {
+            int endOffs = offset + size;
+
+            int regionStart = offset & ~PageMask;
+            int regionEnd = (endOffs + PageMask) & ~PageMask;
+
+            _jitRegion.Block.MapAsRwx((ulong)regionStart, (ulong)(regionEnd - regionStart));
+        }
+
+        private static void ReprotectAsExecutable(int offset, int size)
+        {
+            int endOffs = offset + size;
+
+            int regionStart = offset & ~PageMask;
+            int regionEnd = (endOffs + PageMask) & ~PageMask;
+
+            _jitRegion.Block.MapAsRx((ulong)regionStart, (ulong)(regionEnd - regionStart));
+        }
+
+        private static int Allocate(int codeSize)
+        {
+            codeSize = AlignCodeSize(codeSize);
+
+            int allocOffset = _cacheAllocator.Allocate(codeSize);
+
+            if (allocOffset < 0)
+            {
+                throw new OutOfMemoryException("JIT Cache exhausted.");
+            }
+
+            _jitRegion.ExpandIfNeeded((ulong)allocOffset + (ulong)codeSize);
+
+            return allocOffset;
+        }
+
+        private static int AlignCodeSize(int codeSize)
+        {
+            return checked(codeSize + (CodeAlignment - 1)) & ~(CodeAlignment - 1);
+        }
+
+        private static void Add(int offset, int size, UnwindInfo unwindInfo)
+        {
+            CacheEntry entry = new CacheEntry(offset, size, unwindInfo);
+
+            int index = _cacheEntries.BinarySearch(entry);
+
+            if (index < 0)
+            {
+                index = ~index;
+            }
+
+            _cacheEntries.Insert(index, entry);
+        }
+
+        private static void Remove(int offset)
+        {
+            int index = _cacheEntries.BinarySearch(new CacheEntry(offset, 0, default));
+
+            if (index < 0)
+            {
+                index = ~index - 1;
+            }
+
+            if (index >= 0)
+            {
+                _cacheEntries.RemoveAt(index);
+            }
+        }
+
+        public static bool TryFind(int offset, out CacheEntry entry)
+        {
+            lock (_lock)
+            {
+                int index = _cacheEntries.BinarySearch(new CacheEntry(offset, 0, default));
+
+                if (index < 0)
+                {
+                    index = ~index - 1;
+                }
+
+                if (index >= 0)
+                {
+                    entry = _cacheEntries[index];
+                    return true;
+                }
+            }
+
+            entry = default;
+            return false;
+        }
+    }
+}
+\ No newline at end of file