Implement host tracked memory manager mode (#6356)

* Add host tracked memory manager mode

* Skipping flush is no longer needed

* Formatting + revert unrelated change

* LightningJit: Ensure that dest register is saved for load ops that do partial updates

* avoid allocations when doing address space lookup

Add missing improvement

* IsRmwMemory -> IsPartialRegisterUpdateMemory

* Ensure we iterate all private allocations in range

* PR feedback and potential fixes

* Simplified bridges a lot

* Skip calling SignalMappingChanged if Guest is true

* Late map bridge too

* Force address masking for prefetch instructions

* Reprotection for bridges

* Move partition list validation to separate debug method

* Move host tracked related classes to HostTracked folder

* New HostTracked namespace

* Move host tracked modes to the end of enum to avoid PPTC invalidation

---------

Co-authored-by: riperiperi <rhy3756547@hotmail.com>

3 files changed, 139 insertions, 15 deletions

diff --git a/src/ARMeilleure/Instructions/InstEmitMemoryHelper.cs b/src/ARMeilleure/Instructions/InstEmitMemoryHelper.cs
index a807eed5..ace6fe1c 100644
--- a/src/ARMeilleure/Instructions/InstEmitMemoryHelper.cs
+++ b/src/ARMeilleure/Instructions/InstEmitMemoryHelper.cs
@@ -157,7 +157,7 @@ namespace ARMeilleure.Instructions
 
             context.Copy(temp, value);
 
-            if (!context.Memory.Type.IsHostMapped())
+            if (!context.Memory.Type.IsHostMappedOrTracked())
             {
                 context.Branch(lblEnd);
 
@@ -198,7 +198,7 @@ namespace ARMeilleure.Instructions
 
             SetInt(context, rt, value);
 
-            if (!context.Memory.Type.IsHostMapped())
+            if (!context.Memory.Type.IsHostMappedOrTracked())
             {
                 context.Branch(lblEnd);
 
@@ -265,7 +265,7 @@ namespace ARMeilleure.Instructions
 
             context.Copy(GetVec(rt), value);
 
-            if (!context.Memory.Type.IsHostMapped())
+            if (!context.Memory.Type.IsHostMappedOrTracked())
             {
                 context.Branch(lblEnd);
 
@@ -312,7 +312,7 @@ namespace ARMeilleure.Instructions
                     break;
             }
 
-            if (!context.Memory.Type.IsHostMapped())
+            if (!context.Memory.Type.IsHostMappedOrTracked())
             {
                 context.Branch(lblEnd);
 
@@ -385,7 +385,7 @@ namespace ARMeilleure.Instructions
                     break;
             }
 
-            if (!context.Memory.Type.IsHostMapped())
+            if (!context.Memory.Type.IsHostMappedOrTracked())
             {
                 context.Branch(lblEnd);
 
@@ -403,6 +403,27 @@ namespace ARMeilleure.Instructions
             {
                 return EmitHostMappedPointer(context, address);
             }
+            else if (context.Memory.Type.IsHostTracked())
+            {
+                if (address.Type == OperandType.I32)
+                {
+                    address = context.ZeroExtend32(OperandType.I64, address);
+                }
+
+                if (context.Memory.Type == MemoryManagerType.HostTracked)
+                {
+                    Operand mask = Const(ulong.MaxValue >> (64 - context.Memory.AddressSpaceBits));
+                    address = context.BitwiseAnd(address, mask);
+                }
+
+                Operand ptBase = !context.HasPtc
+                    ? Const(context.Memory.PageTablePointer.ToInt64())
+                    : Const(context.Memory.PageTablePointer.ToInt64(), Ptc.PageTableSymbol);
+
+                Operand ptOffset = context.ShiftRightUI(address, Const(PageBits));
+
+                return context.Add(address, context.Load(OperandType.I64, context.Add(ptBase, context.ShiftLeft(ptOffset, Const(3)))));
+            }
 
             int ptLevelBits = context.Memory.AddressSpaceBits - PageBits;
             int ptLevelSize = 1 << ptLevelBits;
diff --git a/src/ARMeilleure/Memory/MemoryManagerType.cs b/src/ARMeilleure/Memory/MemoryManagerType.cs
index b1cdbb06..bc8ae263 100644
--- a/src/ARMeilleure/Memory/MemoryManagerType.cs
+++ b/src/ARMeilleure/Memory/MemoryManagerType.cs
@@ -29,6 +29,18 @@ namespace ARMeilleure.Memory
         /// Allows invalid access from JIT code to the rest of the program, but is faster.
         /// </summary>
         HostMappedUnsafe,
+
+        /// <summary>
+        /// High level implementation using a software flat page table for address translation
+        /// with no support for handling invalid or non-contiguous memory access.
+        /// </summary>
+        HostTracked,
+
+        /// <summary>
+        /// High level implementation using a software flat page table for address translation
+        /// without masking the address and no support for handling invalid or non-contiguous memory access.
+        /// </summary>
+        HostTrackedUnsafe,
     }
 
     public static class MemoryManagerTypeExtensions
@@ -37,5 +49,15 @@ namespace ARMeilleure.Memory
         {
             return type == MemoryManagerType.HostMapped || type == MemoryManagerType.HostMappedUnsafe;
         }
+
+        public static bool IsHostTracked(this MemoryManagerType type)
+        {
+            return type == MemoryManagerType.HostTracked || type == MemoryManagerType.HostTrackedUnsafe;
+        }
+
+        public static bool IsHostMappedOrTracked(this MemoryManagerType type)
+        {
+            return type.IsHostMapped() || type.IsHostTracked();
+        }
     }
 }
diff --git a/src/ARMeilleure/Signal/NativeSignalHandlerGenerator.cs b/src/ARMeilleure/Signal/NativeSignalHandlerGenerator.cs
index c5e708e1..2ec5bc1b 100644
--- a/src/ARMeilleure/Signal/NativeSignalHandlerGenerator.cs
+++ b/src/ARMeilleure/Signal/NativeSignalHandlerGenerator.cs
@@ -21,10 +21,8 @@ namespace ARMeilleure.Signal
 
         private const uint EXCEPTION_ACCESS_VIOLATION = 0xc0000005;
 
-        private static Operand EmitGenericRegionCheck(EmitterContext context, IntPtr signalStructPtr, Operand faultAddress, Operand isWrite, int rangeStructSize, ulong pageSize)
+        private static Operand EmitGenericRegionCheck(EmitterContext context, IntPtr signalStructPtr, Operand faultAddress, Operand isWrite, int rangeStructSize)
         {
-            ulong pageMask = pageSize - 1;
-
             Operand inRegionLocal = context.AllocateLocal(OperandType.I32);
             context.Copy(inRegionLocal, Const(0));
 
@@ -51,7 +49,7 @@ namespace ARMeilleure.Signal
                 // Only call tracking if in range.
                 context.BranchIfFalse(nextLabel, inRange, BasicBlockFrequency.Cold);
 
-                Operand offset = context.BitwiseAnd(context.Subtract(faultAddress, rangeAddress), Const(~pageMask));
+                Operand offset = context.Subtract(faultAddress, rangeAddress);
 
                 // Call the tracking action, with the pointer's relative offset to the base address.
                 Operand trackingActionPtr = context.Load(OperandType.I64, Const((ulong)signalStructPtr + rangeBaseOffset + 20));
@@ -62,8 +60,10 @@ namespace ARMeilleure.Signal
 
                 // Tracking action should be non-null to call it, otherwise assume false return.
                 context.BranchIfFalse(skipActionLabel, trackingActionPtr);
-                Operand result = context.Call(trackingActionPtr, OperandType.I32, offset, Const(pageSize), isWrite);
-                context.Copy(inRegionLocal, result);
+                Operand result = context.Call(trackingActionPtr, OperandType.I64, offset, Const(1UL), isWrite);
+                context.Copy(inRegionLocal, context.ICompareNotEqual(result, Const(0UL)));
+
+                GenerateFaultAddressPatchCode(context, faultAddress, result);
 
                 context.MarkLabel(skipActionLabel);
 
@@ -155,7 +155,7 @@ namespace ARMeilleure.Signal
             throw new PlatformNotSupportedException();
         }
 
-        public static byte[] GenerateUnixSignalHandler(IntPtr signalStructPtr, int rangeStructSize, ulong pageSize)
+        public static byte[] GenerateUnixSignalHandler(IntPtr signalStructPtr, int rangeStructSize)
         {
             EmitterContext context = new();
 
@@ -168,7 +168,7 @@ namespace ARMeilleure.Signal
 
             Operand isWrite = context.ICompareNotEqual(writeFlag, Const(0L)); // Normalize to 0/1.
 
-            Operand isInRegion = EmitGenericRegionCheck(context, signalStructPtr, faultAddress, isWrite, rangeStructSize, pageSize);
+            Operand isInRegion = EmitGenericRegionCheck(context, signalStructPtr, faultAddress, isWrite, rangeStructSize);
 
             Operand endLabel = Label();
 
@@ -203,7 +203,7 @@ namespace ARMeilleure.Signal
             return Compiler.Compile(cfg, argTypes, OperandType.None, CompilerOptions.HighCq, RuntimeInformation.ProcessArchitecture).Code;
         }
 
-        public static byte[] GenerateWindowsSignalHandler(IntPtr signalStructPtr, int rangeStructSize, ulong pageSize)
+        public static byte[] GenerateWindowsSignalHandler(IntPtr signalStructPtr, int rangeStructSize)
         {
             EmitterContext context = new();
 
@@ -232,7 +232,7 @@ namespace ARMeilleure.Signal
 
             Operand isWrite = context.ICompareNotEqual(writeFlag, Const(0L)); // Normalize to 0/1.
 
-            Operand isInRegion = EmitGenericRegionCheck(context, signalStructPtr, faultAddress, isWrite, rangeStructSize, pageSize);
+            Operand isInRegion = EmitGenericRegionCheck(context, signalStructPtr, faultAddress, isWrite, rangeStructSize);
 
             Operand endLabel = Label();
 
@@ -256,5 +256,86 @@ namespace ARMeilleure.Signal
 
             return Compiler.Compile(cfg, argTypes, OperandType.I32, CompilerOptions.HighCq, RuntimeInformation.ProcessArchitecture).Code;
         }
+
+        private static void GenerateFaultAddressPatchCode(EmitterContext context, Operand faultAddress, Operand newAddress)
+        {
+            if (RuntimeInformation.ProcessArchitecture == Architecture.Arm64)
+            {
+                if (SupportsFaultAddressPatchingForHostOs())
+                {
+                    Operand lblSkip = Label();
+
+                    context.BranchIf(lblSkip, faultAddress, newAddress, Comparison.Equal);
+
+                    Operand ucontextPtr = context.LoadArgument(OperandType.I64, 2);
+                    Operand pcCtxAddress = default;
+                    ulong baseRegsOffset = 0;
+
+                    if (OperatingSystem.IsLinux())
+                    {
+                        pcCtxAddress = context.Add(ucontextPtr, Const(440UL));
+                        baseRegsOffset = 184UL;
+                    }
+                    else if (OperatingSystem.IsMacOS() || OperatingSystem.IsIOS())
+                    {
+                        ucontextPtr = context.Load(OperandType.I64, context.Add(ucontextPtr, Const(48UL)));
+
+                        pcCtxAddress = context.Add(ucontextPtr, Const(272UL));
+                        baseRegsOffset = 16UL;
+                    }
+
+                    Operand pc = context.Load(OperandType.I64, pcCtxAddress);
+
+                    Operand reg = GetAddressRegisterFromArm64Instruction(context, pc);
+                    Operand reg64 = context.ZeroExtend32(OperandType.I64, reg);
+                    Operand regCtxAddress = context.Add(ucontextPtr, context.Add(context.ShiftLeft(reg64, Const(3)), Const(baseRegsOffset)));
+                    Operand regAddress = context.Load(OperandType.I64, regCtxAddress);
+
+                    Operand addressDelta = context.Subtract(regAddress, faultAddress);
+
+                    context.Store(regCtxAddress, context.Add(newAddress, addressDelta));
+
+                    context.MarkLabel(lblSkip);
+                }
+            }
+        }
+
+        private static Operand GetAddressRegisterFromArm64Instruction(EmitterContext context, Operand pc)
+        {
+            Operand inst = context.Load(OperandType.I32, pc);
+            Operand reg = context.AllocateLocal(OperandType.I32);
+
+            Operand isSysInst = context.ICompareEqual(context.BitwiseAnd(inst, Const(0xFFF80000)), Const(0xD5080000));
+
+            Operand lblSys = Label();
+            Operand lblEnd = Label();
+
+            context.BranchIfTrue(lblSys, isSysInst, BasicBlockFrequency.Cold);
+
+            context.Copy(reg, context.BitwiseAnd(context.ShiftRightUI(inst, Const(5)), Const(0x1F)));
+            context.Branch(lblEnd);
+
+            context.MarkLabel(lblSys);
+            context.Copy(reg, context.BitwiseAnd(inst, Const(0x1F)));
+
+            context.MarkLabel(lblEnd);
+
+            return reg;
+        }
+
+        public static bool SupportsFaultAddressPatchingForHost()
+        {
+            return SupportsFaultAddressPatchingForHostArch() && SupportsFaultAddressPatchingForHostOs();
+        }
+
+        private static bool SupportsFaultAddressPatchingForHostArch()
+        {
+            return RuntimeInformation.ProcessArchitecture == Architecture.Arm64;
+        }
+
+        private static bool SupportsFaultAddressPatchingForHostOs()
+        {
+            return OperatingSystem.IsLinux() || OperatingSystem.IsMacOS() || OperatingSystem.IsIOS();
+        }
     }
 }