Optimize kernel memory block lookup and consolidate RBTree implementations (#3410)

* Implement intrusive red-black tree, use it for HLE kernel block manager

* Implement TreeDictionary using IntrusiveRedBlackTree

* Implement IntervalTree using IntrusiveRedBlackTree

* Implement IntervalTree (on Ryujinx.Memory) using IntrusiveRedBlackTree

* Make PredecessorOf and SuccessorOf internal, expose Predecessor and Successor properties on the node itself

* Allocation free tree node lookup

3 files changed, 129 insertions, 137 deletions

diff --git a/Ryujinx.HLE/HOS/Kernel/Memory/KMemoryBlock.cs b/Ryujinx.HLE/HOS/Kernel/Memory/KMemoryBlock.cs
index b612022c..e082105b 100644
--- a/Ryujinx.HLE/HOS/Kernel/Memory/KMemoryBlock.cs
+++ b/Ryujinx.HLE/HOS/Kernel/Memory/KMemoryBlock.cs
@@ -1,42 +1,43 @@
+using Ryujinx.Common.Collections;
 using System;
 
 namespace Ryujinx.HLE.HOS.Kernel.Memory
 {
-    class KMemoryBlock
+    class KMemoryBlock : IntrusiveRedBlackTreeNode<KMemoryBlock>, IComparable<KMemoryBlock>, IComparable<ulong>
     {
         public ulong BaseAddress { get; private set; }
-        public ulong PagesCount  { get; private set; }
+        public ulong PagesCount { get; private set; }
 
-        public MemoryState      State            { get; private set; }
-        public KMemoryPermission Permission       { get; private set; }
-        public MemoryAttribute  Attribute        { get; private set; }
+        public MemoryState State { get; private set; }
+        public KMemoryPermission Permission { get; private set; }
+        public MemoryAttribute Attribute { get; private set; }
         public KMemoryPermission SourcePermission { get; private set; }
 
-        public int IpcRefCount    { get; private set; }
+        public int IpcRefCount { get; private set; }
         public int DeviceRefCount { get; private set; }
 
         public KMemoryBlock(
-            ulong            baseAddress,
-            ulong            pagesCount,
-            MemoryState      state,
+            ulong baseAddress,
+            ulong pagesCount,
+            MemoryState state,
             KMemoryPermission permission,
-            MemoryAttribute  attribute,
-            int              ipcRefCount    = 0,
-            int              deviceRefCount = 0)
+            MemoryAttribute attribute,
+            int ipcRefCount = 0,
+            int deviceRefCount = 0)
         {
-            BaseAddress    = baseAddress;
-            PagesCount     = pagesCount;
-            State          = state;
-            Attribute      = attribute;
-            Permission     = permission;
-            IpcRefCount    = ipcRefCount;
+            BaseAddress = baseAddress;
+            PagesCount = pagesCount;
+            State = state;
+            Attribute = attribute;
+            Permission = permission;
+            IpcRefCount = ipcRefCount;
             DeviceRefCount = deviceRefCount;
         }
 
         public void SetState(KMemoryPermission permission, MemoryState state, MemoryAttribute attribute)
         {
             Permission = permission;
-            State      = state;
+            State = state;
             Attribute &= MemoryAttribute.IpcAndDeviceMapped;
             Attribute |= attribute;
         }
@@ -55,7 +56,7 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
                 SourcePermission = Permission;
 
                 Permission &= ~KMemoryPermission.ReadAndWrite;
-                Permission |=  KMemoryPermission.ReadAndWrite & newPermission;
+                Permission |= KMemoryPermission.ReadAndWrite & newPermission;
             }
 
             Attribute |= MemoryAttribute.IpcMapped;
@@ -119,5 +120,37 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
                 IpcRefCount,
                 DeviceRefCount);
         }
+
+        public int CompareTo(KMemoryBlock other)
+        {
+            if (BaseAddress < other.BaseAddress)
+            {
+                return -1;
+            }
+            else if (BaseAddress <= other.BaseAddress + other.PagesCount * KPageTableBase.PageSize - 1UL)
+            {
+                return 0;
+            }
+            else
+            {
+                return 1;
+            }
+        }
+
+        public int CompareTo(ulong address)
+        {
+            if (address < BaseAddress)
+            {
+                return 1;
+            }
+            else if (address <= BaseAddress + PagesCount * KPageTableBase.PageSize - 1UL)
+            {
+                return 0;
+            }
+            else
+            {
+                return -1;
+            }
+        }
     }
 }
\ No newline at end of file
diff --git a/Ryujinx.HLE/HOS/Kernel/Memory/KMemoryBlockManager.cs b/Ryujinx.HLE/HOS/Kernel/Memory/KMemoryBlockManager.cs
index c0d11a95..9cfdfda3 100644
--- a/Ryujinx.HLE/HOS/Kernel/Memory/KMemoryBlockManager.cs
+++ b/Ryujinx.HLE/HOS/Kernel/Memory/KMemoryBlockManager.cs
@@ -1,5 +1,5 @@
-using Ryujinx.HLE.HOS.Kernel.Common;
-using System.Collections.Generic;
+using Ryujinx.Common.Collections;
+using Ryujinx.HLE.HOS.Kernel.Common;
 using System.Diagnostics;
 
 namespace Ryujinx.HLE.HOS.Kernel.Memory
@@ -8,26 +8,27 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
     {
         private const int PageSize = KPageTableBase.PageSize;
 
-        private readonly LinkedList<KMemoryBlock> _blocks;
+        private readonly IntrusiveRedBlackTree<KMemoryBlock> _blockTree;
 
-        public int BlocksCount => _blocks.Count;
+        public int BlocksCount => _blockTree.Count;
 
         private KMemoryBlockSlabManager _slabManager;
 
+        private ulong _addrSpaceStart;
         private ulong _addrSpaceEnd;
 
         public KMemoryBlockManager()
         {
-            _blocks = new LinkedList<KMemoryBlock>();
+            _blockTree = new IntrusiveRedBlackTree<KMemoryBlock>();
         }
 
         public KernelResult Initialize(ulong addrSpaceStart, ulong addrSpaceEnd, KMemoryBlockSlabManager slabManager)
         {
             _slabManager = slabManager;
+            _addrSpaceStart = addrSpaceStart;
             _addrSpaceEnd = addrSpaceEnd;
 
-            // First insertion will always need only a single block,
-            // because there's nothing else to split.
+            // First insertion will always need only a single block, because there's nothing to split.
             if (!slabManager.CanAllocate(1))
             {
                 return KernelResult.OutOfResource;
@@ -35,7 +36,7 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
 
             ulong addrSpacePagesCount = (addrSpaceEnd - addrSpaceStart) / PageSize;
 
-            _blocks.AddFirst(new KMemoryBlock(
+            _blockTree.Add(new KMemoryBlock(
                 addrSpaceStart,
                 addrSpacePagesCount,
                 MemoryState.Unmapped,
@@ -58,20 +59,17 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
             // Insert new block on the list only on areas where the state
             // of the block matches the state specified on the old* state
             // arguments, otherwise leave it as is.
-            int oldCount = _blocks.Count;
+
+            int oldCount = _blockTree.Count;
 
             oldAttribute |= MemoryAttribute.IpcAndDeviceMapped;
 
             ulong endAddr = baseAddress + pagesCount * PageSize;
 
-            LinkedListNode<KMemoryBlock> node = _blocks.First;
+            KMemoryBlock currBlock = FindBlock(baseAddress);
 
-            while (node != null)
+            while (currBlock != null)
             {
-                LinkedListNode<KMemoryBlock> newNode = node;
-
-                KMemoryBlock currBlock = node.Value;
-
                 ulong currBaseAddr = currBlock.BaseAddress;
                 ulong currEndAddr = currBlock.PagesCount * PageSize + currBaseAddr;
 
@@ -83,24 +81,27 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
                         currBlock.Permission != oldPermission ||
                         currBlockAttr != oldAttribute)
                     {
-                        node = node.Next;
+                        currBlock = currBlock.Successor;
 
                         continue;
                     }
 
                     if (baseAddress > currBaseAddr)
                     {
-                        _blocks.AddBefore(node, currBlock.SplitRightAtAddress(baseAddress));
+                        KMemoryBlock newBlock = currBlock.SplitRightAtAddress(baseAddress);
+                        _blockTree.Add(newBlock);
                     }
 
                     if (endAddr < currEndAddr)
                     {
-                        newNode = _blocks.AddBefore(node, currBlock.SplitRightAtAddress(endAddr));
+                        KMemoryBlock newBlock = currBlock.SplitRightAtAddress(endAddr);
+                        _blockTree.Add(newBlock);
+                        currBlock = newBlock;
                     }
 
-                    newNode.Value.SetState(newPermission, newState, newAttribute);
+                    currBlock.SetState(newPermission, newState, newAttribute);
 
-                    newNode = MergeEqualStateNeighbors(newNode);
+                    currBlock = MergeEqualStateNeighbors(currBlock);
                 }
 
                 if (currEndAddr - 1 >= endAddr - 1)
@@ -108,10 +109,10 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
                     break;
                 }
 
-                node = newNode.Next;
+                currBlock = currBlock.Successor;
             }
 
-            _slabManager.Count += _blocks.Count - oldCount;
+            _slabManager.Count += _blockTree.Count - oldCount;
 
             ValidateInternalState();
         }
@@ -125,18 +126,15 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
         {
             // Inserts new block at the list, replacing and splitting
             // existing blocks as needed.
-            int oldCount = _blocks.Count;
+
+            int oldCount = _blockTree.Count;
 
             ulong endAddr = baseAddress + pagesCount * PageSize;
 
-            LinkedListNode<KMemoryBlock> node = _blocks.First;
+            KMemoryBlock currBlock = FindBlock(baseAddress);
 
-            while (node != null)
+            while (currBlock != null)
             {
-                LinkedListNode<KMemoryBlock> newNode = node;
-
-                KMemoryBlock currBlock = node.Value;
-
                 ulong currBaseAddr = currBlock.BaseAddress;
                 ulong currEndAddr = currBlock.PagesCount * PageSize + currBaseAddr;
 
@@ -144,17 +142,20 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
                 {
                     if (baseAddress > currBaseAddr)
                     {
-                        _blocks.AddBefore(node, currBlock.SplitRightAtAddress(baseAddress));
+                        KMemoryBlock newBlock = currBlock.SplitRightAtAddress(baseAddress);
+                        _blockTree.Add(newBlock);
                     }
 
                     if (endAddr < currEndAddr)
                     {
-                        newNode = _blocks.AddBefore(node, currBlock.SplitRightAtAddress(endAddr));
+                        KMemoryBlock newBlock = currBlock.SplitRightAtAddress(endAddr);
+                        _blockTree.Add(newBlock);
+                        currBlock = newBlock;
                     }
 
-                    newNode.Value.SetState(permission, state, attribute);
+                    currBlock.SetState(permission, state, attribute);
 
-                    newNode = MergeEqualStateNeighbors(newNode);
+                    currBlock = MergeEqualStateNeighbors(currBlock);
                 }
 
                 if (currEndAddr - 1 >= endAddr - 1)
@@ -162,10 +163,10 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
                     break;
                 }
 
-                node = newNode.Next;
+                currBlock = currBlock.Successor;
             }
 
-            _slabManager.Count += _blocks.Count - oldCount;
+            _slabManager.Count += _blockTree.Count - oldCount;
 
             ValidateInternalState();
         }
@@ -181,18 +182,15 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
             // Inserts new block at the list, replacing and splitting
             // existing blocks as needed, then calling the callback
             // function on the new block.
-            int oldCount = _blocks.Count;
+
+            int oldCount = _blockTree.Count;
 
             ulong endAddr = baseAddress + pagesCount * PageSize;
 
-            LinkedListNode<KMemoryBlock> node = _blocks.First;
+            KMemoryBlock currBlock = FindBlock(baseAddress);
 
-            while (node != null)
+            while (currBlock != null)
             {
-                LinkedListNode<KMemoryBlock> newNode = node;
-
-                KMemoryBlock currBlock = node.Value;
-
                 ulong currBaseAddr = currBlock.BaseAddress;
                 ulong currEndAddr = currBlock.PagesCount * PageSize + currBaseAddr;
 
@@ -200,19 +198,20 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
                 {
                     if (baseAddress > currBaseAddr)
                     {
-                        _blocks.AddBefore(node, currBlock.SplitRightAtAddress(baseAddress));
+                        KMemoryBlock newBlock = currBlock.SplitRightAtAddress(baseAddress);
+                        _blockTree.Add(newBlock);
                     }
 
                     if (endAddr < currEndAddr)
                     {
-                        newNode = _blocks.AddBefore(node, currBlock.SplitRightAtAddress(endAddr));
+                        KMemoryBlock newBlock = currBlock.SplitRightAtAddress(endAddr);
+                        _blockTree.Add(newBlock);
+                        currBlock = newBlock;
                     }
 
-                    KMemoryBlock newBlock = newNode.Value;
+                    blockMutate(currBlock, permission);
 
-                    blockMutate(newBlock, permission);
-
-                    newNode = MergeEqualStateNeighbors(newNode);
+                    currBlock = MergeEqualStateNeighbors(currBlock);
                 }
 
                 if (currEndAddr - 1 >= endAddr - 1)
@@ -220,10 +219,10 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
                     break;
                 }
 
-                node = newNode.Next;
+                currBlock = currBlock.Successor;
             }
 
-            _slabManager.Count += _blocks.Count - oldCount;
+            _slabManager.Count += _blockTree.Count - oldCount;
 
             ValidateInternalState();
         }
@@ -233,58 +232,42 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
         {
             ulong expectedAddress = 0;
 
-            LinkedListNode<KMemoryBlock> node = _blocks.First;
+            KMemoryBlock currBlock = FindBlock(_addrSpaceStart);
 
-            while (node != null)
+            while (currBlock != null)
             {
-                LinkedListNode<KMemoryBlock> newNode = node;
-
-                KMemoryBlock currBlock = node.Value;
-
                 Debug.Assert(currBlock.BaseAddress == expectedAddress);
 
                 expectedAddress = currBlock.BaseAddress + currBlock.PagesCount * PageSize;
 
-                node = newNode.Next;
+                currBlock = currBlock.Successor;
             }
 
             Debug.Assert(expectedAddress == _addrSpaceEnd);
         }
 
-        private LinkedListNode<KMemoryBlock> MergeEqualStateNeighbors(LinkedListNode<KMemoryBlock> node)
+        private KMemoryBlock MergeEqualStateNeighbors(KMemoryBlock block)
         {
-            KMemoryBlock block = node.Value;
+            KMemoryBlock previousBlock = block.Predecessor;
+            KMemoryBlock nextBlock = block.Successor;
 
-            if (node.Previous != null)
+            if (previousBlock != null && BlockStateEquals(block, previousBlock))
             {
-                KMemoryBlock previousBlock = node.Previous.Value;
-
-                if (BlockStateEquals(block, previousBlock))
-                {
-                    LinkedListNode<KMemoryBlock> previousNode = node.Previous;
-
-                    _blocks.Remove(node);
+                _blockTree.Remove(block);
 
-                    previousBlock.AddPages(block.PagesCount);
+                previousBlock.AddPages(block.PagesCount);
 
-                    node = previousNode;
-                    block = previousBlock;
-                }
+                block = previousBlock;
             }
 
-            if (node.Next != null)
+            if (nextBlock != null && BlockStateEquals(block, nextBlock))
             {
-                KMemoryBlock nextBlock = node.Next.Value;
+                _blockTree.Remove(nextBlock);
 
-                if (BlockStateEquals(block, nextBlock))
-                {
-                    _blocks.Remove(node.Next);
-
-                    block.AddPages(nextBlock.PagesCount);
-                }
+                block.AddPages(nextBlock.PagesCount);
             }
 
-            return node;
+            return block;
         }
 
         private static bool BlockStateEquals(KMemoryBlock lhs, KMemoryBlock rhs)
@@ -299,31 +282,7 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
 
         public KMemoryBlock FindBlock(ulong address)
         {
-            return FindBlockNode(address)?.Value;
-        }
-
-        public LinkedListNode<KMemoryBlock> FindBlockNode(ulong address)
-        {
-            lock (_blocks)
-            {
-                LinkedListNode<KMemoryBlock> node = _blocks.First;
-
-                while (node != null)
-                {
-                    KMemoryBlock block = node.Value;
-
-                    ulong currEndAddr = block.PagesCount * PageSize + block.BaseAddress;
-
-                    if (block.BaseAddress <= address && currEndAddr - 1 >= address)
-                    {
-                        return node;
-                    }
-
-                    node = node.Next;
-                }
-            }
-
-            return null;
+            return _blockTree.GetNodeByKey(address);
         }
     }
 }
diff --git a/Ryujinx.HLE/HOS/Kernel/Memory/KPageTableBase.cs b/Ryujinx.HLE/HOS/Kernel/Memory/KPageTableBase.cs
index ab43b477..857be7a6 100644
--- a/Ryujinx.HLE/HOS/Kernel/Memory/KPageTableBase.cs
+++ b/Ryujinx.HLE/HOS/Kernel/Memory/KPageTableBase.cs
@@ -2447,9 +2447,9 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
         {
             ulong endAddr = address + size;
 
-            LinkedListNode<KMemoryBlock> node = _blockManager.FindBlockNode(address);
+            KMemoryBlock currBlock = _blockManager.FindBlock(address);
 
-            KMemoryInfo info = node.Value.GetInfo();
+            KMemoryInfo info = currBlock.GetInfo();
 
             MemoryState firstState = info.State;
             KMemoryPermission firstPermission = info.Permission;
@@ -2457,7 +2457,7 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
 
             do
             {
-                info = node.Value.GetInfo();
+                info = currBlock.GetInfo();
 
                 // Check if the block state matches what we expect.
                 if (firstState != info.State ||
@@ -2474,7 +2474,7 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
                     return false;
                 }
             }
-            while (info.Address + info.Size - 1 < endAddr - 1 && (node = node.Next) != null);
+            while (info.Address + info.Size - 1 < endAddr - 1 && (currBlock = currBlock.Successor) != null);
 
             outState = firstState;
             outPermission = firstPermission;
@@ -2509,17 +2509,17 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
 
         private IEnumerable<KMemoryInfo> IterateOverRange(ulong start, ulong end)
         {
-            LinkedListNode<KMemoryBlock> node = _blockManager.FindBlockNode(start);
+            KMemoryBlock currBlock = _blockManager.FindBlock(start);
 
             KMemoryInfo info;
 
             do
             {
-                info = node.Value.GetInfo();
+                info = currBlock.GetInfo();
 
                 yield return info;
             }
-            while (info.Address + info.Size - 1 < end - 1 && (node = node.Next) != null);
+            while (info.Address + info.Size - 1 < end - 1 && (currBlock = currBlock.Successor) != null);
         }
 
         private ulong AllocateVa(ulong regionStart, ulong regionPagesCount, ulong neededPagesCount, int alignment)
@@ -2605,9 +2605,9 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
 
             ulong regionEndAddr = regionStart + regionPagesCount * PageSize;
 
-            LinkedListNode<KMemoryBlock> node = _blockManager.FindBlockNode(regionStart);
+            KMemoryBlock currBlock = _blockManager.FindBlock(regionStart);
 
-            KMemoryInfo info = node.Value.GetInfo();
+            KMemoryInfo info = currBlock.GetInfo();
 
             while (regionEndAddr >= info.Address)
             {
@@ -2636,14 +2636,14 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
                     }
                 }
 
-                node = node.Next;
+                currBlock = currBlock.Successor;
 
-                if (node == null)
+                if (currBlock == null)
                 {
                     break;
                 }
 
-                info = node.Value.GetInfo();
+                info = currBlock.GetInfo();
             }
 
             return 0;