Implement support for page sizes > 4KB (#4252)

* Implement support for page sizes > 4KB

* Check and work around more alignment issues

* Was not meant to change this

* Use MemoryBlock.GetPageSize() value for signal handler code

* Do not take the path for private allocations if host supports 4KB pages

* Add Flags attribute on MemoryMapFlags

* Fix dirty region size with 16kb pages

Would accidentally report a size that was too high (generally 16k instead of 4k, uploading 4x as much data)

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

6 files changed, 859 insertions, 28 deletions

diff --git a/Ryujinx.Cpu/AddressSpace.cs b/Ryujinx.Cpu/AddressSpace.cs
new file mode 100644
index 00000000..cea3b56d
--- /dev/null
+++ b/Ryujinx.Cpu/AddressSpace.cs
@@ -0,0 +1,470 @@
+using Ryujinx.Common;
+using Ryujinx.Common.Collections;
+using Ryujinx.Memory;
+using System;
+
+namespace Ryujinx.Cpu
+{
+    class AddressSpace : IDisposable
+    {
+        private const ulong PageSize = 0x1000;
+
+        private const int DefaultBlockAlignment = 1 << 20;
+
+        private enum MappingType : byte
+        {
+            None,
+            Private,
+            Shared
+        }
+
+        private class Mapping : IntrusiveRedBlackTreeNode<Mapping>, IComparable<Mapping>
+        {
+            public ulong Address { get; private set; }
+            public ulong Size { get; private set; }
+            public ulong EndAddress => Address + Size;
+            public MappingType Type { get; private set; }
+
+            public Mapping(ulong address, ulong size, MappingType type)
+            {
+                Address = address;
+                Size = size;
+                Type = type;
+            }
+
+            public Mapping Split(ulong splitAddress)
+            {
+                ulong leftSize = splitAddress - Address;
+                ulong rightSize = EndAddress - splitAddress;
+
+                Mapping left = new Mapping(Address, leftSize, Type);
+
+                Address = splitAddress;
+                Size = rightSize;
+
+                return left;
+            }
+
+            public void UpdateState(MappingType newType)
+            {
+                Type = newType;
+            }
+
+            public void Extend(ulong sizeDelta)
+            {
+                Size += sizeDelta;
+            }
+
+            public int CompareTo(Mapping other)
+            {
+                if (Address < other.Address)
+                {
+                    return -1;
+                }
+                else if (Address <= other.EndAddress - 1UL)
+                {
+                    return 0;
+                }
+                else
+                {
+                    return 1;
+                }
+            }
+        }
+
+        private class PrivateMapping : IntrusiveRedBlackTreeNode<PrivateMapping>, IComparable<PrivateMapping>
+        {
+            public ulong Address { get; private set; }
+            public ulong Size { get; private set; }
+            public ulong EndAddress => Address + Size;
+            public PrivateMemoryAllocation PrivateAllocation { get; private set; }
+
+            public PrivateMapping(ulong address, ulong size, PrivateMemoryAllocation privateAllocation)
+            {
+                Address = address;
+                Size = size;
+                PrivateAllocation = privateAllocation;
+            }
+
+            public PrivateMapping Split(ulong splitAddress)
+            {
+                ulong leftSize = splitAddress - Address;
+                ulong rightSize = EndAddress - splitAddress;
+
+                (var leftAllocation, PrivateAllocation) = PrivateAllocation.Split(leftSize);
+
+                PrivateMapping left = new PrivateMapping(Address, leftSize, leftAllocation);
+
+                Address = splitAddress;
+                Size = rightSize;
+
+                return left;
+            }
+
+            public void Map(MemoryBlock baseBlock, MemoryBlock mirrorBlock, PrivateMemoryAllocation newAllocation)
+            {
+                baseBlock.MapView(newAllocation.Memory, newAllocation.Offset, Address, Size);
+                mirrorBlock.MapView(newAllocation.Memory, newAllocation.Offset, Address, Size);
+                PrivateAllocation = newAllocation;
+            }
+
+            public void Unmap(MemoryBlock baseBlock, MemoryBlock mirrorBlock)
+            {
+                if (PrivateAllocation.IsValid)
+                {
+                    baseBlock.UnmapView(PrivateAllocation.Memory, Address, Size);
+                    mirrorBlock.UnmapView(PrivateAllocation.Memory, Address, Size);
+                    PrivateAllocation.Dispose();
+                }
+
+                PrivateAllocation = default;
+            }
+
+            public void Extend(ulong sizeDelta)
+            {
+                Size += sizeDelta;
+            }
+
+            public int CompareTo(PrivateMapping other)
+            {
+                if (Address < other.Address)
+                {
+                    return -1;
+                }
+                else if (Address <= other.EndAddress - 1UL)
+                {
+                    return 0;
+                }
+                else
+                {
+                    return 1;
+                }
+            }
+        }
+
+        private readonly MemoryBlock _backingMemory;
+        private readonly PrivateMemoryAllocator _privateMemoryAllocator;
+        private readonly IntrusiveRedBlackTree<Mapping> _mappingTree;
+        private readonly IntrusiveRedBlackTree<PrivateMapping> _privateTree;
+
+        private readonly object _treeLock;
+
+        private readonly bool _supports4KBPages;
+
+        public MemoryBlock Base { get; }
+        public MemoryBlock Mirror { get; }
+
+        public AddressSpace(MemoryBlock backingMemory, ulong asSize, bool supports4KBPages)
+        {
+            if (!supports4KBPages)
+            {
+                _privateMemoryAllocator = new PrivateMemoryAllocator(DefaultBlockAlignment, MemoryAllocationFlags.Mirrorable | MemoryAllocationFlags.NoMap);
+                _mappingTree = new IntrusiveRedBlackTree<Mapping>();
+                _privateTree = new IntrusiveRedBlackTree<PrivateMapping>();
+                _treeLock = new object();
+
+                _mappingTree.Add(new Mapping(0UL, asSize, MappingType.None));
+                _privateTree.Add(new PrivateMapping(0UL, asSize, default));
+            }
+
+            _backingMemory = backingMemory;
+            _supports4KBPages = supports4KBPages;
+
+            MemoryAllocationFlags asFlags = MemoryAllocationFlags.Reserve | MemoryAllocationFlags.ViewCompatible;
+
+            Base = new MemoryBlock(asSize, asFlags);
+            Mirror = new MemoryBlock(asSize, asFlags);
+        }
+
+        public void Map(ulong va, ulong pa, ulong size, MemoryMapFlags flags)
+        {
+            if (_supports4KBPages)
+            {
+                Base.MapView(_backingMemory, pa, va, size);
+                Mirror.MapView(_backingMemory, pa, va, size);
+
+                return;
+            }
+
+            lock (_treeLock)
+            {
+                ulong alignment = MemoryBlock.GetPageSize();
+                bool isAligned = ((va | pa | size) & (alignment - 1)) == 0;
+
+                if (flags.HasFlag(MemoryMapFlags.Private) && !isAligned)
+                {
+                    Update(va, pa, size, MappingType.Private);
+                }
+                else
+                {
+                    // The update method assumes that shared mappings are already aligned.
+
+                    if (!flags.HasFlag(MemoryMapFlags.Private))
+                    {
+                        if ((va & (alignment - 1)) != (pa & (alignment - 1)))
+                        {
+                            throw new InvalidMemoryRegionException($"Virtual address 0x{va:X} and physical address 0x{pa:X} are misaligned and can't be aligned.");
+                        }
+
+                        ulong endAddress = va + size;
+                        va = BitUtils.AlignDown(va, alignment);
+                        pa = BitUtils.AlignDown(pa, alignment);
+                        size = BitUtils.AlignUp(endAddress, alignment) - va;
+                    }
+
+                    Update(va, pa, size, MappingType.Shared);
+                }
+            }
+        }
+
+        public void Unmap(ulong va, ulong size)
+        {
+            if (_supports4KBPages)
+            {
+                Base.UnmapView(_backingMemory, va, size);
+                Mirror.UnmapView(_backingMemory, va, size);
+
+                return;
+            }
+
+            lock (_treeLock)
+            {
+                Update(va, 0UL, size, MappingType.None);
+            }
+        }
+
+        private void Update(ulong va, ulong pa, ulong size, MappingType type)
+        {
+            Mapping map = _mappingTree.GetNode(new Mapping(va, 1UL, MappingType.None));
+
+            Update(map, va, pa, size, type);
+        }
+
+        private Mapping Update(Mapping map, ulong va, ulong pa, ulong size, MappingType type)
+        {
+            ulong endAddress = va + size;
+
+            for (; map != null; map = map.Successor)
+            {
+                if (map.Address < va)
+                {
+                    _mappingTree.Add(map.Split(va));
+                }
+
+                if (map.EndAddress > endAddress)
+                {
+                    Mapping newMap = map.Split(endAddress);
+                    _mappingTree.Add(newMap);
+                    map = newMap;
+                }
+
+                switch (type)
+                {
+                    case MappingType.None:
+                        if (map.Type == MappingType.Shared)
+                        {
+                            ulong startOffset = map.Address - va;
+                            ulong mapVa = va + startOffset;
+                            ulong mapSize = Math.Min(size - startOffset, map.Size);
+                            ulong mapEndAddress = mapVa + mapSize;
+                            ulong alignment = MemoryBlock.GetPageSize();
+
+                            mapVa = BitUtils.AlignDown(mapVa, alignment);
+                            mapEndAddress = BitUtils.AlignUp(mapEndAddress, alignment);
+
+                            mapSize = mapEndAddress - mapVa;
+
+                            Base.UnmapView(_backingMemory, mapVa, mapSize);
+                            Mirror.UnmapView(_backingMemory, mapVa, mapSize);
+                        }
+                        else
+                        {
+                            UnmapPrivate(va, size);
+                        }
+                        break;
+                    case MappingType.Private:
+                        if (map.Type == MappingType.Shared)
+                        {
+                            throw new InvalidMemoryRegionException($"Private mapping request at 0x{va:X} with size 0x{size:X} overlaps shared mapping at 0x{map.Address:X} with size 0x{map.Size:X}.");
+                        }
+                        else
+                        {
+                            MapPrivate(va, size);
+                        }
+                        break;
+                    case MappingType.Shared:
+                        if (map.Type != MappingType.None)
+                        {
+                            throw new InvalidMemoryRegionException($"Shared mapping request at 0x{va:X} with size 0x{size:X} overlaps mapping at 0x{map.Address:X} with size 0x{map.Size:X}.");
+                        }
+                        else
+                        {
+                            ulong startOffset = map.Address - va;
+                            ulong mapPa = pa + startOffset;
+                            ulong mapVa = va + startOffset;
+                            ulong mapSize = Math.Min(size - startOffset, map.Size);
+
+                            Base.MapView(_backingMemory, mapPa, mapVa, mapSize);
+                            Mirror.MapView(_backingMemory, mapPa, mapVa, mapSize);
+                        }
+                        break;
+                }
+
+                map.UpdateState(type);
+                map = TryCoalesce(map);
+
+                if (map.EndAddress >= endAddress)
+                {
+                    break;
+                }
+            }
+
+            return map;
+        }
+
+        private Mapping TryCoalesce(Mapping map)
+        {
+            Mapping previousMap = map.Predecessor;
+            Mapping nextMap = map.Successor;
+
+            if (previousMap != null && CanCoalesce(previousMap, map))
+            {
+                previousMap.Extend(map.Size);
+                _mappingTree.Remove(map);
+                map = previousMap;
+            }
+
+            if (nextMap != null && CanCoalesce(map, nextMap))
+            {
+                map.Extend(nextMap.Size);
+                _mappingTree.Remove(nextMap);
+            }
+
+            return map;
+        }
+
+        private static bool CanCoalesce(Mapping left, Mapping right)
+        {
+            return left.Type == right.Type;
+        }
+
+        private void MapPrivate(ulong va, ulong size)
+        {
+            ulong endAddress = va + size;
+
+            ulong alignment = MemoryBlock.GetPageSize();
+
+            // Expand the range outwards based on page size to ensure that at least the requested region is mapped.
+            ulong vaAligned = BitUtils.AlignDown(va, alignment);
+            ulong endAddressAligned = BitUtils.AlignUp(endAddress, alignment);
+
+            ulong sizeAligned = endAddressAligned - vaAligned;
+
+            PrivateMapping map = _privateTree.GetNode(new PrivateMapping(va, 1UL, default));
+
+            for (; map != null; map = map.Successor)
+            {
+                if (!map.PrivateAllocation.IsValid)
+                {
+                    if (map.Address < vaAligned)
+                    {
+                        _privateTree.Add(map.Split(vaAligned));
+                    }
+
+                    if (map.EndAddress > endAddressAligned)
+                    {
+                        PrivateMapping newMap = map.Split(endAddressAligned);
+                        _privateTree.Add(newMap);
+                        map = newMap;
+                    }
+
+                    map.Map(Base, Mirror, _privateMemoryAllocator.Allocate(map.Size, MemoryBlock.GetPageSize()));
+                }
+
+                if (map.EndAddress >= endAddressAligned)
+                {
+                    break;
+                }
+            }
+        }
+
+        private void UnmapPrivate(ulong va, ulong size)
+        {
+            ulong endAddress = va + size;
+
+            ulong alignment = MemoryBlock.GetPageSize();
+
+            // Shrink the range inwards based on page size to ensure we won't unmap memory that might be still in use.
+            ulong vaAligned = BitUtils.AlignUp(va, alignment);
+            ulong endAddressAligned = BitUtils.AlignDown(endAddress, alignment);
+
+            if (endAddressAligned <= vaAligned)
+            {
+                return;
+            }
+
+            ulong alignedSize = endAddressAligned - vaAligned;
+
+            PrivateMapping map = _privateTree.GetNode(new PrivateMapping(va, 1UL, default));
+
+            for (; map != null; map = map.Successor)
+            {
+                if (map.PrivateAllocation.IsValid)
+                {
+                    if (map.Address < vaAligned)
+                    {
+                        _privateTree.Add(map.Split(vaAligned));
+                    }
+
+                    if (map.EndAddress > endAddressAligned)
+                    {
+                        PrivateMapping newMap = map.Split(endAddressAligned);
+                        _privateTree.Add(newMap);
+                        map = newMap;
+                    }
+
+                    map.Unmap(Base, Mirror);
+                    map = TryCoalesce(map);
+                }
+
+                if (map.EndAddress >= endAddressAligned)
+                {
+                    break;
+                }
+            }
+        }
+
+        private PrivateMapping TryCoalesce(PrivateMapping map)
+        {
+            PrivateMapping previousMap = map.Predecessor;
+            PrivateMapping nextMap = map.Successor;
+
+            if (previousMap != null && CanCoalesce(previousMap, map))
+            {
+                previousMap.Extend(map.Size);
+                _privateTree.Remove(map);
+                map = previousMap;
+            }
+
+            if (nextMap != null && CanCoalesce(map, nextMap))
+            {
+                map.Extend(nextMap.Size);
+                _privateTree.Remove(nextMap);
+            }
+
+            return map;
+        }
+
+        private static bool CanCoalesce(PrivateMapping left, PrivateMapping right)
+        {
+            return !left.PrivateAllocation.IsValid && !right.PrivateAllocation.IsValid;
+        }
+
+        public void Dispose()
+        {
+            _privateMemoryAllocator.Dispose();
+            Base.Dispose();
+            Mirror.Dispose();
+        }
+    }
+}
\ No newline at end of file
diff --git a/Ryujinx.Cpu/Jit/JitMemoryAllocator.cs b/Ryujinx.Cpu/Jit/JitMemoryAllocator.cs
index 0cf35c17..4aa78d06 100644
--- a/Ryujinx.Cpu/Jit/JitMemoryAllocator.cs
+++ b/Ryujinx.Cpu/Jit/JitMemoryAllocator.cs
@@ -7,5 +7,7 @@ namespace Ryujinx.Cpu.Jit
     {
         public IJitMemoryBlock Allocate(ulong size) => new JitMemoryBlock(size, MemoryAllocationFlags.None);
         public IJitMemoryBlock Reserve(ulong size) => new JitMemoryBlock(size, MemoryAllocationFlags.Reserve | MemoryAllocationFlags.Jit);
+
+        public ulong GetPageSize() => MemoryBlock.GetPageSize();
     }
 }
diff --git a/Ryujinx.Cpu/Jit/MemoryManager.cs b/Ryujinx.Cpu/Jit/MemoryManager.cs
index 21c50d51..014d843b 100644
--- a/Ryujinx.Cpu/Jit/MemoryManager.cs
+++ b/Ryujinx.Cpu/Jit/MemoryManager.cs
@@ -28,6 +28,9 @@ namespace Ryujinx.Cpu.Jit
         private readonly MemoryBlock _backingMemory;
         private readonly InvalidAccessHandler _invalidAccessHandler;
 
+        /// <inheritdoc/>
+        public bool Supports4KBPages => true;
+
         /// <summary>
         /// Address space width in bits.
         /// </summary>
@@ -76,7 +79,7 @@ namespace Ryujinx.Cpu.Jit
         }
 
         /// <inheritdoc/>
-        public void Map(ulong va, ulong pa, ulong size)
+        public void Map(ulong va, ulong pa, ulong size, MemoryMapFlags flags)
         {
             AssertValidAddressAndSize(va, size);
 
@@ -91,10 +94,17 @@ namespace Ryujinx.Cpu.Jit
                 pa += PageSize;
                 remainingSize -= PageSize;
             }
+
             Tracking.Map(oVa, size);
         }
 
         /// <inheritdoc/>
+        public void MapForeign(ulong va, nuint hostPointer, ulong size)
+        {
+            throw new NotSupportedException();
+        }
+
+        /// <inheritdoc/>
         public void Unmap(ulong va, ulong size)
         {
             // If size is 0, there's nothing to unmap, just exit early.
@@ -379,6 +389,32 @@ namespace Ryujinx.Cpu.Jit
         }
 
         /// <inheritdoc/>
+        public IEnumerable<HostMemoryRange> GetHostRegions(ulong va, ulong size)
+        {
+            if (size == 0)
+            {
+                return Enumerable.Empty<HostMemoryRange>();
+            }
+
+            var guestRegions = GetPhysicalRegionsImpl(va, size);
+            if (guestRegions == null)
+            {
+                return null;
+            }
+
+            var regions = new HostMemoryRange[guestRegions.Count];
+
+            for (int i = 0; i < regions.Length; i++)
+            {
+                var guestRegion = guestRegions[i];
+                IntPtr pointer = _backingMemory.GetPointer(guestRegion.Address, guestRegion.Size);
+                regions[i] = new HostMemoryRange((nuint)(ulong)pointer, guestRegion.Size);
+            }
+
+            return regions;
+        }
+
+        /// <inheritdoc/>
         public IEnumerable<MemoryRange> GetPhysicalRegions(ulong va, ulong size)
         {
             if (size == 0)
@@ -386,6 +422,11 @@ namespace Ryujinx.Cpu.Jit
                 return Enumerable.Empty<MemoryRange>();
             }
 
+            return GetPhysicalRegionsImpl(va, size);
+        }
+
+        private List<MemoryRange> GetPhysicalRegionsImpl(ulong va, ulong size)
+        {
             if (!ValidateAddress(va) || !ValidateAddressAndSize(va, size))
             {
                 return null;
diff --git a/Ryujinx.Cpu/Jit/MemoryManagerHostMapped.cs b/Ryujinx.Cpu/Jit/MemoryManagerHostMapped.cs
index c4e59db9..856b6b9b 100644
--- a/Ryujinx.Cpu/Jit/MemoryManagerHostMapped.cs
+++ b/Ryujinx.Cpu/Jit/MemoryManagerHostMapped.cs
@@ -5,6 +5,7 @@ using Ryujinx.Memory.Range;
 using Ryujinx.Memory.Tracking;
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Threading;
 
@@ -37,20 +38,21 @@ namespace Ryujinx.Cpu.Jit
         private readonly InvalidAccessHandler _invalidAccessHandler;
         private readonly bool _unsafeMode;
 
-        private readonly MemoryBlock _addressSpace;
-        private readonly MemoryBlock _addressSpaceMirror;
+        private readonly AddressSpace _addressSpace;
         private readonly ulong _addressSpaceSize;
 
-        private readonly MemoryBlock _backingMemory;
         private readonly PageTable<ulong> _pageTable;
 
         private readonly MemoryEhMeilleure _memoryEh;
 
         private readonly ulong[] _pageBitmap;
 
+        /// <inheritdoc/>
+        public bool Supports4KBPages => MemoryBlock.GetPageSize() == PageSize;
+
         public int AddressSpaceBits { get; }
 
-        public IntPtr PageTablePointer => _addressSpace.Pointer;
+        public IntPtr PageTablePointer => _addressSpace.Base.Pointer;
 
         public MemoryManagerType Type => _unsafeMode ? MemoryManagerType.HostMappedUnsafe : MemoryManagerType.HostMapped;
 
@@ -67,7 +69,6 @@ namespace Ryujinx.Cpu.Jit
         /// <param name="invalidAccessHandler">Optional function to handle invalid memory accesses</param>
         public MemoryManagerHostMapped(MemoryBlock backingMemory, ulong addressSpaceSize, bool unsafeMode, InvalidAccessHandler invalidAccessHandler = null)
         {
-            _backingMemory = backingMemory;
             _pageTable = new PageTable<ulong>();
             _invalidAccessHandler = invalidAccessHandler;
             _unsafeMode = unsafeMode;
@@ -86,13 +87,10 @@ namespace Ryujinx.Cpu.Jit
 
             _pageBitmap = new ulong[1 << (AddressSpaceBits - (PageBits + PageToPteShift))];
 
-            MemoryAllocationFlags asFlags = MemoryAllocationFlags.Reserve | MemoryAllocationFlags.ViewCompatible;
+            _addressSpace = new AddressSpace(backingMemory, asSize, Supports4KBPages);
 
-            _addressSpace = new MemoryBlock(asSize, asFlags);
-            _addressSpaceMirror = new MemoryBlock(asSize, asFlags);
-
-            Tracking = new MemoryTracking(this, PageSize, invalidAccessHandler);
-            _memoryEh = new MemoryEhMeilleure(_addressSpace, _addressSpaceMirror, Tracking);
+            Tracking = new MemoryTracking(this, (int)MemoryBlock.GetPageSize(), invalidAccessHandler);
+            _memoryEh = new MemoryEhMeilleure(_addressSpace.Base, _addressSpace.Mirror, Tracking);
         }
 
         /// <summary>
@@ -145,12 +143,11 @@ namespace Ryujinx.Cpu.Jit
         }
 
         /// <inheritdoc/>
-        public void Map(ulong va, ulong pa, ulong size)
+        public void Map(ulong va, ulong pa, ulong size, MemoryMapFlags flags)
         {
             AssertValidAddressAndSize(va, size);
 
-            _addressSpace.MapView(_backingMemory, pa, va, size);
-            _addressSpaceMirror.MapView(_backingMemory, pa, va, size);
+            _addressSpace.Map(va, pa, size, flags);
             AddMapping(va, size);
             PtMap(va, pa, size);
 
@@ -158,6 +155,12 @@ namespace Ryujinx.Cpu.Jit
         }
 
         /// <inheritdoc/>
+        public void MapForeign(ulong va, nuint hostPointer, ulong size)
+        {
+            throw new NotSupportedException();
+        }
+
+        /// <inheritdoc/>
         public void Unmap(ulong va, ulong size)
         {
             AssertValidAddressAndSize(va, size);
@@ -167,8 +170,7 @@ namespace Ryujinx.Cpu.Jit
 
             RemoveMapping(va, size);
             PtUnmap(va, size);
-            _addressSpace.UnmapView(_backingMemory, va, size);
-            _addressSpaceMirror.UnmapView(_backingMemory, va, size);
+            _addressSpace.Unmap(va, size);
         }
 
         private void PtMap(ulong va, ulong pa, ulong size)
@@ -201,7 +203,7 @@ namespace Ryujinx.Cpu.Jit
             {
                 AssertMapped(va, (ulong)Unsafe.SizeOf<T>());
 
-                return _addressSpaceMirror.Read<T>(va);
+                return _addressSpace.Mirror.Read<T>(va);
             }
             catch (InvalidMemoryRegionException)
             {
@@ -241,7 +243,7 @@ namespace Ryujinx.Cpu.Jit
             {
                 AssertMapped(va, (ulong)data.Length);
 
-                _addressSpaceMirror.Read(va, data);
+                _addressSpace.Mirror.Read(va, data);
             }
             catch (InvalidMemoryRegionException)
             {
@@ -260,7 +262,7 @@ namespace Ryujinx.Cpu.Jit
             {
                 SignalMemoryTracking(va, (ulong)Unsafe.SizeOf<T>(), write: true);
 
-                _addressSpaceMirror.Write(va, value);
+                _addressSpace.Mirror.Write(va, value);
             }
             catch (InvalidMemoryRegionException)
             {
@@ -278,7 +280,7 @@ namespace Ryujinx.Cpu.Jit
             {
                 SignalMemoryTracking(va, (ulong)data.Length, write: true);
 
-                _addressSpaceMirror.Write(va, data);
+                _addressSpace.Mirror.Write(va, data);
             }
             catch (InvalidMemoryRegionException)
             {
@@ -296,7 +298,7 @@ namespace Ryujinx.Cpu.Jit
             {
                 AssertMapped(va, (ulong)data.Length);
 
-                _addressSpaceMirror.Write(va, data);
+                _addressSpace.Mirror.Write(va, data);
             }
             catch (InvalidMemoryRegionException)
             {
@@ -314,7 +316,7 @@ namespace Ryujinx.Cpu.Jit
             {
                 SignalMemoryTracking(va, (ulong)data.Length, false);
 
-                Span<byte> target = _addressSpaceMirror.GetSpan(va, data.Length);
+                Span<byte> target = _addressSpace.Mirror.GetSpan(va, data.Length);
                 bool changed = !data.SequenceEqual(target);
 
                 if (changed)
@@ -347,7 +349,7 @@ namespace Ryujinx.Cpu.Jit
                 AssertMapped(va, (ulong)size);
             }
 
-            return _addressSpaceMirror.GetSpan(va, size);
+            return _addressSpace.Mirror.GetSpan(va, size);
         }
 
         /// <inheritdoc/>
@@ -362,7 +364,7 @@ namespace Ryujinx.Cpu.Jit
                 AssertMapped(va, (ulong)size);
             }
 
-            return _addressSpaceMirror.GetWritableRegion(va, size);
+            return _addressSpace.Mirror.GetWritableRegion(va, size);
         }
 
         /// <inheritdoc/>
@@ -370,7 +372,7 @@ namespace Ryujinx.Cpu.Jit
         {
             SignalMemoryTracking(va, (ulong)Unsafe.SizeOf<T>(), true);
 
-            return ref _addressSpaceMirror.GetRef<T>(va);
+            return ref _addressSpace.Mirror.GetRef<T>(va);
         }
 
         /// <inheritdoc/>
@@ -455,6 +457,14 @@ namespace Ryujinx.Cpu.Jit
         }
 
         /// <inheritdoc/>
+        public IEnumerable<HostMemoryRange> GetHostRegions(ulong va, ulong size)
+        {
+            AssertValidAddressAndSize(va, size);
+
+            return Enumerable.Repeat(new HostMemoryRange((nuint)(ulong)_addressSpace.Mirror.GetPointer(va, size), size), 1);
+        }
+
+        /// <inheritdoc/>
         public IEnumerable<MemoryRange> GetPhysicalRegions(ulong va, ulong size)
         {
             int pages = GetPagesCount(va, (uint)size, out va);
@@ -692,7 +702,7 @@ namespace Ryujinx.Cpu.Jit
                 _ => MemoryPermission.None
             };
 
-            _addressSpace.Reprotect(va, size, protection, false);
+            _addressSpace.Base.Reprotect(va, size, protection, false);
         }
 
         /// <inheritdoc/>
@@ -799,7 +809,6 @@ namespace Ryujinx.Cpu.Jit
         protected override void Destroy()
         {
             _addressSpace.Dispose();
-            _addressSpaceMirror.Dispose();
             _memoryEh.Dispose();
         }
 
diff --git a/Ryujinx.Cpu/PrivateMemoryAllocation.cs b/Ryujinx.Cpu/PrivateMemoryAllocation.cs
new file mode 100644
index 00000000..1327880e
--- /dev/null
+++ b/Ryujinx.Cpu/PrivateMemoryAllocation.cs
@@ -0,0 +1,41 @@
+using Ryujinx.Memory;
+using System;
+
+namespace Ryujinx.Cpu
+{
+    struct PrivateMemoryAllocation : IDisposable
+    {
+        private readonly PrivateMemoryAllocator _owner;
+        private readonly PrivateMemoryAllocator.Block _block;
+
+        public bool IsValid => _owner != null;
+        public MemoryBlock Memory => _block?.Memory;
+        public ulong Offset { get; }
+        public ulong Size { get; }
+
+        public PrivateMemoryAllocation(
+            PrivateMemoryAllocator owner,
+            PrivateMemoryAllocator.Block block,
+            ulong offset,
+            ulong size)
+        {
+            _owner = owner;
+            _block = block;
+            Offset = offset;
+            Size = size;
+        }
+
+        public (PrivateMemoryAllocation, PrivateMemoryAllocation) Split(ulong splitOffset)
+        {
+            PrivateMemoryAllocation left = new PrivateMemoryAllocation(_owner, _block, Offset, splitOffset);
+            PrivateMemoryAllocation right = new PrivateMemoryAllocation(_owner, _block, Offset + splitOffset, Size - splitOffset);
+
+            return (left, right);
+        }
+
+        public void Dispose()
+        {
+            _owner.Free(_block, Offset, Size);
+        }
+    }
+}
diff --git a/Ryujinx.Cpu/PrivateMemoryAllocator.cs b/Ryujinx.Cpu/PrivateMemoryAllocator.cs
new file mode 100644
index 00000000..cbf1f1d9
--- /dev/null
+++ b/Ryujinx.Cpu/PrivateMemoryAllocator.cs
@@ -0,0 +1,268 @@
+using Ryujinx.Common;
+using Ryujinx.Memory;
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+
+namespace Ryujinx.Cpu
+{
+    class PrivateMemoryAllocator : PrivateMemoryAllocatorImpl<PrivateMemoryAllocator.Block>
+    {
+        public const ulong InvalidOffset = ulong.MaxValue;
+
+        public class Block : IComparable<Block>
+        {
+            public MemoryBlock Memory { get; private set; }
+            public ulong Size { get; }
+
+            private struct Range : IComparable<Range>
+            {
+                public ulong Offset { get; }
+                public ulong Size { get; }
+
+                public Range(ulong offset, ulong size)
+                {
+                    Offset = offset;
+                    Size = size;
+                }
+
+                public int CompareTo(Range other)
+                {
+                    return Offset.CompareTo(other.Offset);
+                }
+            }
+
+            private readonly List<Range> _freeRanges;
+
+            public Block(MemoryBlock memory, ulong size)
+            {
+                Memory = memory;
+                Size = size;
+                _freeRanges = new List<Range>
+                {
+                    new Range(0, size)
+                };
+            }
+
+            public ulong Allocate(ulong size, ulong alignment)
+            {
+                for (int i = 0; i < _freeRanges.Count; i++)
+                {
+                    var range = _freeRanges[i];
+
+                    ulong alignedOffset = BitUtils.AlignUp(range.Offset, alignment);
+                    ulong sizeDelta = alignedOffset - range.Offset;
+                    ulong usableSize = range.Size - sizeDelta;
+
+                    if (sizeDelta < range.Size && usableSize >= size)
+                    {
+                        _freeRanges.RemoveAt(i);
+
+                        if (sizeDelta != 0)
+                        {
+                            InsertFreeRange(range.Offset, sizeDelta);
+                        }
+
+                        ulong endOffset = range.Offset + range.Size;
+                        ulong remainingSize = endOffset - (alignedOffset + size);
+                        if (remainingSize != 0)
+                        {
+                            InsertFreeRange(endOffset - remainingSize, remainingSize);
+                        }
+
+                        return alignedOffset;
+                    }
+                }
+
+                return InvalidOffset;
+            }
+
+            public void Free(ulong offset, ulong size)
+            {
+                InsertFreeRangeComingled(offset, size);
+            }
+
+            private void InsertFreeRange(ulong offset, ulong size)
+            {
+                var range = new Range(offset, size);
+                int index = _freeRanges.BinarySearch(range);
+                if (index < 0)
+                {
+                    index = ~index;
+                }
+
+                _freeRanges.Insert(index, range);
+            }
+
+            private void InsertFreeRangeComingled(ulong offset, ulong size)
+            {
+                ulong endOffset = offset + size;
+                var range = new Range(offset, size);
+                int index = _freeRanges.BinarySearch(range);
+                if (index < 0)
+                {
+                    index = ~index;
+                }
+
+                if (index < _freeRanges.Count && _freeRanges[index].Offset == endOffset)
+                {
+                    endOffset = _freeRanges[index].Offset + _freeRanges[index].Size;
+                    _freeRanges.RemoveAt(index);
+                }
+
+                if (index > 0 && _freeRanges[index - 1].Offset + _freeRanges[index - 1].Size == offset)
+                {
+                    offset = _freeRanges[index - 1].Offset;
+                    _freeRanges.RemoveAt(--index);
+                }
+
+                range = new Range(offset, endOffset - offset);
+
+                _freeRanges.Insert(index, range);
+            }
+
+            public bool IsTotallyFree()
+            {
+                if (_freeRanges.Count == 1 && _freeRanges[0].Size == Size)
+                {
+                    Debug.Assert(_freeRanges[0].Offset == 0);
+                    return true;
+                }
+
+                return false;
+            }
+
+            public int CompareTo(Block other)
+            {
+                return Size.CompareTo(other.Size);
+            }
+
+            public virtual void Destroy()
+            {
+                Memory.Dispose();
+            }
+        }
+
+        public PrivateMemoryAllocator(int blockAlignment, MemoryAllocationFlags allocationFlags) : base(blockAlignment, allocationFlags)
+        {
+        }
+
+        public PrivateMemoryAllocation Allocate(ulong size, ulong alignment)
+        {
+            var allocation = Allocate(size, alignment, CreateBlock);
+
+            return new PrivateMemoryAllocation(this, allocation.Block, allocation.Offset, allocation.Size);
+        }
+
+        private Block CreateBlock(MemoryBlock memory, ulong size)
+        {
+            return new Block(memory, size);
+        }
+    }
+
+    class PrivateMemoryAllocatorImpl<T> : IDisposable where T : PrivateMemoryAllocator.Block
+    {
+        private const ulong InvalidOffset = ulong.MaxValue;
+
+        public struct Allocation
+        {
+            public T Block { get; }
+            public ulong Offset { get; }
+            public ulong Size { get; }
+
+            public Allocation(T block, ulong offset, ulong size)
+            {
+                Block = block;
+                Offset = offset;
+                Size = size;
+            }
+        }
+
+        private readonly List<T> _blocks;
+
+        private readonly int _blockAlignment;
+        private readonly MemoryAllocationFlags _allocationFlags;
+
+        public PrivateMemoryAllocatorImpl(int blockAlignment, MemoryAllocationFlags allocationFlags)
+        {
+            _blocks = new List<T>();
+            _blockAlignment = blockAlignment;
+            _allocationFlags = allocationFlags;
+        }
+
+        protected Allocation Allocate(ulong size, ulong alignment, Func<MemoryBlock, ulong, T> createBlock)
+        {
+            // Ensure we have a sane alignment value.
+            if ((ulong)(int)alignment != alignment || (int)alignment <= 0)
+            {
+                throw new ArgumentOutOfRangeException(nameof(alignment), $"Invalid alignment 0x{alignment:X}.");
+            }
+
+            for (int i = 0; i < _blocks.Count; i++)
+            {
+                var block = _blocks[i];
+
+                if (block.Size >= size)
+                {
+                    ulong offset = block.Allocate(size, alignment);
+                    if (offset != InvalidOffset)
+                    {
+                        return new Allocation(block, offset, size);
+                    }
+                }
+            }
+
+            ulong blockAlignedSize = BitUtils.AlignUp(size, (ulong)_blockAlignment);
+
+            var memory = new MemoryBlock(blockAlignedSize, _allocationFlags);
+            var newBlock = createBlock(memory, blockAlignedSize);
+
+            InsertBlock(newBlock);
+
+            ulong newBlockOffset = newBlock.Allocate(size, alignment);
+            Debug.Assert(newBlockOffset != InvalidOffset);
+
+            return new Allocation(newBlock, newBlockOffset, size);
+        }
+
+        public void Free(PrivateMemoryAllocator.Block block, ulong offset, ulong size)
+        {
+            block.Free(offset, size);
+
+            if (block.IsTotallyFree())
+            {
+                for (int i = 0; i < _blocks.Count; i++)
+                {
+                    if (_blocks[i] == block)
+                    {
+                        _blocks.RemoveAt(i);
+                        break;
+                    }
+                }
+
+                block.Destroy();
+            }
+        }
+
+        private void InsertBlock(T block)
+        {
+            int index = _blocks.BinarySearch(block);
+            if (index < 0)
+            {
+                index = ~index;
+            }
+
+            _blocks.Insert(index, block);
+        }
+
+        public void Dispose()
+        {
+            for (int i = 0; i < _blocks.Count; i++)
+            {
+                _blocks[i].Destroy();
+            }
+
+            _blocks.Clear();
+        }
+    }
+}
\ No newline at end of file