avoid allocations when doing address space lookup

src/Ryujinx.Common/Collections/IntrusiveRedBlackTreeNode.cs

@@ -5,10 +5,10 @@ namespace Ryujinx.Common.Collections
     /// </summary>
     public class IntrusiveRedBlackTreeNode<T> where T : IntrusiveRedBlackTreeNode<T>
     {
-        internal bool Color = true;
+        public bool Color = true;
-        internal T Left;
+        public T Left;
-        internal T Right;
+        public T Right;
-        internal T Parent;
+        public T Parent;
 
         public T Predecessor => IntrusiveRedBlackTreeImpl<T>.PredecessorOf((T)this);
         public T Successor => IntrusiveRedBlackTreeImpl<T>.SuccessorOf((T)this);

src/Ryujinx.Cpu/Jit/AddressIntrusiveRedBlackTree.cs

@@ -0,0 +1,35 @@
+using Ryujinx.Common.Collections;
+using System;
+
+namespace Ryujinx.Cpu.Jit
+{
+    internal class AddressIntrusiveRedBlackTree<T> : IntrusiveRedBlackTree<T> where T : IntrusiveRedBlackTreeNode<T>, IComparable<T>, IComparable<ulong>
+    {
+        /// <summary>
+        /// Retrieve the node that is considered equal to the specified address by the comparator.
+        /// </summary>
+        /// <param name="address">Address to compare with</param>
+        /// <returns>Node that is equal to <paramref name="address"/></returns>
+        public T GetNode(ulong address)
+        {
+            T node = Root;
+            while (node != null)
+            {
+                int cmp = node.CompareTo(address);
+                if (cmp < 0)
+                {
+                    node = node.Left;
+                }
+                else if (cmp > 0)
+                {
+                    node = node.Right;
+                }
+                else
+                {
+                    return node;
+                }
+            }
+            return null;
+        }
+    }
+}

src/Ryujinx.Cpu/Jit/AddressSpacePageProtections.cs

@@ -10,12 +10,15 @@ namespace Ryujinx.Cpu.Jit
     {
         private const ulong GuestPageSize = 0x1000;
 
-        class PageProtection : IntrusiveRedBlackTreeNode<PageProtection>, IComparable<PageProtection>
+        [ThreadStatic]
+        private static PageProtection _dummyProtection;
+
+        class PageProtection : IntrusiveRedBlackTreeNode<PageProtection>, IComparable<PageProtection>, IComparable<ulong>
         {
             public readonly AddressSpacePartitionAllocation Memory;
             public readonly ulong Offset;
-            public readonly ulong Address;
+            public ulong Address;
-            public readonly ulong Size;
+            public ulong Size;
 
             private MemoryBlock _viewBlock;
 
@@ -63,9 +66,25 @@ namespace Ryujinx.Cpu.Jit
                     return 1;
                 }
             }
+
+            public int CompareTo(ulong address)
+            {
+                if (address < Address)
+                {
+                    return -1;
+                }
+                else if (address <= Address + Size - 1UL)
+                {
+                    return 0;
+                }
+                else
+                {
+                    return 1;
+                }
+            }
         }
 
-        private readonly IntrusiveRedBlackTree<PageProtection> _protectionTree;
+        private readonly AddressIntrusiveRedBlackTree<PageProtection> _protectionTree;
 
         public AddressSpacePageProtections()
         {
@@ -99,7 +118,7 @@ namespace Ryujinx.Cpu.Jit
         {
             ulong pageSize = MemoryBlock.GetPageSize();
 
-            PageProtection pageProtection = _protectionTree.GetNode(new PageProtection(default, 0, va, 1));
+            PageProtection pageProtection = _protectionTree.GetNode(va);
 
             if (pageProtection == null)
             {
@@ -310,7 +329,21 @@ namespace Ryujinx.Cpu.Jit
 
         private PageProtection GetLowestOverlap(ulong va, ulong size)
         {
-            PageProtection pageProtection = _protectionTree.GetNode(new PageProtection(default, 0, va, size));
+            PageProtection lookup = _dummyProtection;
+
+            if (lookup == null)
+            {
+                lookup = new(default, 0, va, size);
+                
+                _dummyProtection = lookup;
+            }
+            else
+            {
+                lookup.Address = va;
+                lookup.Size = size;
+            }
+
+            PageProtection pageProtection = _protectionTree.GetNode(lookup);
 
             if (pageProtection == null)
             {

src/Ryujinx.Cpu/Jit/AddressSpacePartition.cs

@@ -35,7 +35,7 @@ namespace Ryujinx.Cpu.Jit
             Private,
         }
 
-        private class Mapping : IntrusiveRedBlackTreeNode<Mapping>, IComparable<Mapping>
+        private class Mapping : IntrusiveRedBlackTreeNode<Mapping>, IComparable<Mapping>, IComparable<ulong>
         {
             public ulong Address { get; private set; }
             public ulong Size { get; private set; }
@@ -87,9 +87,25 @@ namespace Ryujinx.Cpu.Jit
                     return 1;
                 }
             }
+
+            public int CompareTo(ulong address)
+            {
+                if (address < Address)
+                {
+                    return -1;
+                }
+                else if (address <= EndAddress - 1UL)
+                {
+                    return 0;
+                }
+                else
+                {
+                    return 1;
+                }
+            }
         }
 
-        private class PrivateMapping : IntrusiveRedBlackTreeNode<PrivateMapping>, IComparable<PrivateMapping>
+        private class PrivateMapping : IntrusiveRedBlackTreeNode<PrivateMapping>, IComparable<PrivateMapping>, IComparable<ulong>
         {
             public ulong Address { get; private set; }
             public ulong Size { get; private set; }
@@ -158,13 +174,29 @@ namespace Ryujinx.Cpu.Jit
                     return 1;
                 }
             }
+
+            public int CompareTo(ulong address)
+            {
+                if (address < Address)
+                {
+                    return -1;
+                }
+                else if (address <= EndAddress - 1UL)
+                {
+                    return 0;
+                }
+                else
+                {
+                    return 1;
+                }
+            }
         }
 
         private readonly MemoryBlock _backingMemory;
         private readonly AddressSpacePartitionAllocation _baseMemory;
         private readonly PrivateMemoryAllocator _privateMemoryAllocator;
-        private readonly IntrusiveRedBlackTree<Mapping> _mappingTree;
+        private readonly AddressIntrusiveRedBlackTree<Mapping> _mappingTree;
-        private readonly IntrusiveRedBlackTree<PrivateMapping> _privateTree;
+        private readonly AddressIntrusiveRedBlackTree<PrivateMapping> _privateTree;
         private readonly AddressSpacePageProtections _pageProtections;
 
         private readonly ReaderWriterLockSlim _treeLock;
@@ -185,8 +217,8 @@ namespace Ryujinx.Cpu.Jit
         public AddressSpacePartition(AddressSpacePartitionAllocation baseMemory, MemoryBlock backingMemory, ulong address, ulong size)
         {
             _privateMemoryAllocator = new PrivateMemoryAllocator(DefaultBlockAlignment, MemoryAllocationFlags.Mirrorable);
-            _mappingTree = new IntrusiveRedBlackTree<Mapping>();
+            _mappingTree = new AddressIntrusiveRedBlackTree<Mapping>();
-            _privateTree = new IntrusiveRedBlackTree<PrivateMapping>();
+            _privateTree = new AddressIntrusiveRedBlackTree<PrivateMapping>();
             _pageProtections = new AddressSpacePageProtections();
             _treeLock = new ReaderWriterLockSlim();
 
@@ -212,7 +244,7 @@ namespace Ryujinx.Cpu.Jit
 
             try
             {
-                Mapping map = _mappingTree.GetNode(new Mapping(Address, Size, MappingType.None));
+                Mapping map = _mappingTree.GetNode(Address);
 
                 return map != null && map.Address == Address && map.Size == Size && map.Type == MappingType.None;
             }
@@ -375,7 +407,7 @@ namespace Ryujinx.Cpu.Jit
 
             try
             {
-                PrivateMapping map = _privateTree.GetNode(new PrivateMapping(Address, 1UL, default));
+                PrivateMapping map = _privateTree.GetNode(Address);
 
                 if (map != null && map.PrivateAllocation.IsValid)
                 {
@@ -398,7 +430,7 @@ namespace Ryujinx.Cpu.Jit
             {
                 ulong pageAddress = EndAddress - _hostPageSize;
 
-                PrivateMapping map = _privateTree.GetNode(new PrivateMapping(pageAddress, 1UL, default));
+                PrivateMapping map = _privateTree.GetNode(pageAddress);
 
                 if (map != null && map.PrivateAllocation.IsValid)
                 {
@@ -419,7 +451,7 @@ namespace Ryujinx.Cpu.Jit
 
             try
             {
-                PrivateMapping map = _privateTree.GetNode(new PrivateMapping(va, 1UL, default));
+                PrivateMapping map = _privateTree.GetNode(va);
 
                 if (map != null && map.PrivateAllocation.IsValid)
                 {
@@ -440,7 +472,7 @@ namespace Ryujinx.Cpu.Jit
 
             try
             {
-                Mapping map = _mappingTree.GetNode(new Mapping(va, 1UL, MappingType.None));
+                Mapping map = _mappingTree.GetNode(va);
 
                 Update(map, va, pa, size, type);
             }
@@ -534,7 +566,7 @@ namespace Ryujinx.Cpu.Jit
             ulong vaAligned = BitUtils.AlignDown(va, alignment);
             ulong endAddressAligned = BitUtils.AlignUp(endAddress, alignment);
 
-            PrivateMapping map = _privateTree.GetNode(new PrivateMapping(va, 1UL, default));
+            PrivateMapping map = _privateTree.GetNode(va);
 
             for (; map != null; map = map.Successor)
             {
@@ -578,7 +610,7 @@ namespace Ryujinx.Cpu.Jit
                 return;
             }
 
-            PrivateMapping map = _privateTree.GetNode(new PrivateMapping(vaAligned, 1UL, default));
+            PrivateMapping map = _privateTree.GetNode(vaAligned);
 
             for (; map != null; map = map.Successor)
             {
@@ -639,7 +671,7 @@ namespace Ryujinx.Cpu.Jit
 
             try
             {
-                PrivateMapping map = _privateTree.GetNode(new PrivateMapping(va, 1UL, default));
+                PrivateMapping map = _privateTree.GetNode(va);
 
                 nextVa = map.EndAddress;
 
@@ -667,7 +699,7 @@ namespace Ryujinx.Cpu.Jit
 
             try
             {
-                PrivateMapping map = _privateTree.GetNode(new PrivateMapping(va, size, default));
+                PrivateMapping map = _privateTree.GetNode(va);
 
                 if (map != null && map.PrivateAllocation.IsValid)
                 {

src/Ryujinx.Cpu/Jit/AddressSpacePartitionAllocator.cs

@@ -61,7 +61,7 @@ namespace Ryujinx.Cpu.Jit
             private readonly MemoryTracking _tracking;
             private readonly MemoryEhMeilleure _memoryEh;
 
-            private class Mapping : IntrusiveRedBlackTreeNode<Mapping>, IComparable<Mapping>
+            private class Mapping : IntrusiveRedBlackTreeNode<Mapping>, IComparable<Mapping>, IComparable<ulong>
             {
                 public ulong Address { get; }
                 public ulong Size { get; }
@@ -94,9 +94,25 @@ namespace Ryujinx.Cpu.Jit
                         return 1;
                     }
                 }
+
+                public int CompareTo(ulong address)
+                {
+                    if (address < Address)
+                    {
+                        return -1;
+                    }
+                    else if (address <= EndAddress - 1UL)
+                    {
+                        return 0;
+                    }
+                    else
+                    {
+                        return 1;
+                    }
+                }
             }
 
-            private readonly IntrusiveRedBlackTree<Mapping> _mappingTree;
+            private readonly AddressIntrusiveRedBlackTree<Mapping> _mappingTree;
             private readonly object _lock;
 
             public Block(MemoryTracking tracking, MemoryBlock memory, ulong size, object locker) : base(memory, size)
@@ -114,7 +130,7 @@ namespace Ryujinx.Cpu.Jit
 
             public void RemoveMapping(ulong offset, ulong size)
             {
-                _mappingTree.Remove(_mappingTree.GetNode(new Mapping(offset, size, 0, 0, 0)));
+                _mappingTree.Remove(_mappingTree.GetNode(offset));
             }
 
             private bool VirtualMemoryEvent(ulong address, ulong size, bool write)
@@ -123,7 +139,7 @@ namespace Ryujinx.Cpu.Jit
 
                 lock (_lock)
                 {
-                    map = _mappingTree.GetNode(new Mapping(address, size, 0, 0, 0));
+                    map = _mappingTree.GetNode(address);
                 }
 
                 if (map == null)