using System; using System.Collections; using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; namespace Ryujinx.Common.Collections { /// /// Dictionary that provides the ability for O(logN) Lookups for keys that exist in the Dictionary, and O(logN) lookups for keys immediately greater than or less than a specified key. /// /// Key /// Value public class TreeDictionary : IDictionary where K : IComparable { private const bool Black = true; private const bool Red = false; private Node _root = null; private int _count = 0; public TreeDictionary() { } #region Public Methods /// /// Returns the value of the node whose key is , or the default value if no such node exists. /// /// Key of the node value to get /// Value associated w/ /// is null public V Get(K key) { if (key == null) { throw new ArgumentNullException(nameof(key)); } Node node = GetNode(key); if (node == null) { return default; } return node.Value; } /// /// Adds a new node into the tree whose key is key and value is . ///

/// Note: Adding the same key multiple times will cause the value for that key to be overwritten. ///
/// Key of the node to add /// Value of the node to add /// or are null public void Add(K key, V value) { if (key == null) { throw new ArgumentNullException(nameof(key)); } if (null == value) { throw new ArgumentNullException(nameof(value)); } Insert(key, value); } /// /// Removes the node whose key is from the tree. /// /// Key of the node to remove /// is null public void Remove(K key) { if (key == null) { throw new ArgumentNullException(nameof(key)); } if (Delete(key) != null) { _count--; } } /// /// Returns the value whose key is equal to or immediately less than . /// /// Key for which to find the floor value of /// Key of node immediately less than /// is null public K Floor(K key) { Node node = FloorNode(key); if (node != null) { return node.Key; } return default; } /// /// Returns the node whose key is equal to or immediately greater than . /// /// Key for which to find the ceiling node of /// Key of node immediately greater than /// is null public K Ceiling(K key) { Node node = CeilingNode(key); if (node != null) { return node.Key; } return default; } /// /// Finds the value whose key is immediately greater than . /// /// Key to find the successor of /// Value public K SuccessorOf(K key) { Node node = GetNode(key); if (node != null) { Node successor = SuccessorOf(node); return successor != null ? successor.Key : default; } return default; } /// /// Finds the value whose key is immediately less than . /// /// Key to find the predecessor of /// Value public K PredecessorOf(K key) { Node node = GetNode(key); if (node != null) { Node predecessor = PredecessorOf(node); return predecessor != null ? predecessor.Key : default; } return default; } /// /// Adds all the nodes in the dictionary as key/value pairs into . ///

/// The key/value pairs will be added in Level Order. ///
/// List to add the tree pairs into public List> AsLevelOrderList() { List> list = new List>(); Queue> nodes = new Queue>(); if (this._root != null) { nodes.Enqueue(this._root); } while (nodes.Count > 0) { Node node = nodes.Dequeue(); list.Add(new KeyValuePair(node.Key, node.Value)); if (node.Left != null) { nodes.Enqueue(node.Left); } if (node.Right != null) { nodes.Enqueue(node.Right); } } return list; } /// /// Adds all the nodes in the dictionary into . /// /// A list of all KeyValuePairs sorted by Key Order public List> AsList() { List> list = new List>(); AddToList(_root, list); return list; } #endregion #region Private Methods (BST) /// /// Adds all nodes that are children of or contained within into , in Key Order. /// /// The node to search for nodes within /// The list to add node to private void AddToList(Node node, List> list) { if (node == null) { return; } AddToList(node.Left, list); list.Add(new KeyValuePair(node.Key, node.Value)); AddToList(node.Right, list); } /// /// Retrieve the node reference whose key is , or null if no such node exists. /// /// Key of the node to get /// Node reference in the tree /// is null private Node GetNode(K key) { if (key == null) { throw new ArgumentNullException(nameof(key)); } Node node = _root; while (node != null) { int cmp = key.CompareTo(node.Key); if (cmp < 0) { node = node.Left; } else if (cmp > 0) { node = node.Right; } else { return node; } } return null; } /// /// Inserts a new node into the tree whose key is and value is . ///

/// Adding the same key multiple times will overwrite the previous value. ///
/// Key of the node to insert /// Value of the node to insert private void Insert(K key, V value) { Node newNode = BSTInsert(key, value); RestoreBalanceAfterInsertion(newNode); } /// /// Insertion Mechanism for a Binary Search Tree (BST). ///

/// Iterates the tree starting from the root and inserts a new node where all children in the left subtree are less than , and all children in the right subtree are greater than . ///

/// Note: If a node whose key is already exists, it's value will be overwritten. ///
/// Key of the node to insert /// Value of the node to insert /// The inserted Node private Node BSTInsert(K key, V value) { Node parent = null; Node node = _root; while (node != null) { parent = node; int cmp = key.CompareTo(node.Key); if (cmp < 0) { node = node.Left; } else if (cmp > 0) { node = node.Right; } else { node.Value = value; return node; } } Node newNode = new Node(key, value, parent); if (newNode.Parent == null) { _root = newNode; } else if (key.CompareTo(parent.Key) < 0) { parent.Left = newNode; } else { parent.Right = newNode; } _count++; return newNode; } /// /// Removes from the dictionary, if it exists. /// /// Key of the node to delete /// The deleted Node private Node Delete(K key) { // O(1) Retrieval Node nodeToDelete = GetNode(key); if (nodeToDelete == null) return null; Node replacementNode; if (LeftOf(nodeToDelete) == null || RightOf(nodeToDelete) == null) { replacementNode = nodeToDelete; } else { replacementNode = PredecessorOf(nodeToDelete); } Node tmp = LeftOf(replacementNode) ?? RightOf(replacementNode); if (tmp != null) { tmp.Parent = ParentOf(replacementNode); } if (ParentOf(replacementNode) == null) { _root = tmp; } else if (replacementNode == LeftOf(ParentOf(replacementNode))) { ParentOf(replacementNode).Left = tmp; } else { ParentOf(replacementNode).Right = tmp; } if (replacementNode != nodeToDelete) { nodeToDelete.Key = replacementNode.Key; nodeToDelete.Value = replacementNode.Value; } if (tmp != null && ColorOf(replacementNode) == Black) { RestoreBalanceAfterRemoval(tmp); } return replacementNode; } /// /// Returns the node with the largest key where is considered the root node. /// /// Root Node /// Node with the maximum key in the tree of private static Node Maximum(Node node) { Node tmp = node; while (tmp.Right != null) { tmp = tmp.Right; } return tmp; } /// /// Returns the node with the smallest key where is considered the root node. /// /// Root Node /// Node with the minimum key in the tree of /// is null private static Node Minimum(Node node) { if (node == null) { throw new ArgumentNullException(nameof(node)); } Node tmp = node; while (tmp.Left != null) { tmp = tmp.Left; } return tmp; } /// /// Returns the node whose key immediately less than or equal to . /// /// Key for which to find the floor node of /// Node whose key is immediately less than or equal to , or null if no such node is found. /// is null private Node FloorNode(K key) { if (key == null) { throw new ArgumentNullException(nameof(key)); } Node tmp = _root; while (tmp != null) { int cmp = key.CompareTo(tmp.Key); if (cmp > 0) { if (tmp.Right != null) { tmp = tmp.Right; } else { return tmp; } } else if (cmp < 0) { if (tmp.Left != null) { tmp = tmp.Left; } else { Node parent = tmp.Parent; Node ptr = tmp; while (parent != null && ptr == parent.Left) { ptr = parent; parent = parent.Parent; } return parent; } } else { return tmp; } } return null; } /// /// Returns the node whose key is immediately greater than or equal to than . /// /// Key for which to find the ceiling node of /// Node whose key is immediately greater than or equal to , or null if no such node is found. /// is null private Node CeilingNode(K key) { if (key == null) { throw new ArgumentNullException(nameof(key)); } Node tmp = _root; while (tmp != null) { int cmp = key.CompareTo(tmp.Key); if (cmp < 0) { if (tmp.Left != null) { tmp = tmp.Left; } else { return tmp; } } else if (cmp > 0) { if (tmp.Right != null) { tmp = tmp.Right; } else { Node parent = tmp.Parent; Node ptr = tmp; while (parent != null && ptr == parent.Right) { ptr = parent; parent = parent.Parent; } return parent; } } else { return tmp; } } return null; } /// /// Finds the node with the key is immediately greater than . /// /// Node to find the successor of /// Successor of private static Node SuccessorOf(Node node) { if (node.Right != null) { return Minimum(node.Right); } Node parent = node.Parent; while (parent != null && node == parent.Right) { node = parent; parent = parent.Parent; } return parent; } /// /// Finds the node whose key is immediately less than . /// /// Node to find the predecessor of /// Predecessor of private static Node PredecessorOf(Node node) { if (node.Left != null) { return Maximum(node.Left); } Node parent = node.Parent; while (parent != null && node == parent.Left) { node = parent; parent = parent.Parent; } return parent; } #endregion #region Private Methods (RBL) private void RestoreBalanceAfterRemoval(Node balanceNode) { Node ptr = balanceNode; while (ptr != _root && ColorOf(ptr) == Black) { if (ptr == LeftOf(ParentOf(ptr))) { Node sibling = RightOf(ParentOf(ptr)); if (ColorOf(sibling) == Red) { SetColor(sibling, Black); SetColor(ParentOf(ptr), Red); RotateLeft(ParentOf(ptr)); sibling = RightOf(ParentOf(ptr)); } if (ColorOf(LeftOf(sibling)) == Black && ColorOf(RightOf(sibling)) == Black) { SetColor(sibling, Red); ptr = ParentOf(ptr); } else { if (ColorOf(RightOf(sibling)) == Black) { SetColor(LeftOf(sibling), Black); SetColor(sibling, Red); RotateRight(sibling); sibling = RightOf(ParentOf(ptr)); } SetColor(sibling, ColorOf(ParentOf(ptr))); SetColor(ParentOf(ptr), Black); SetColor(RightOf(sibling), Black); RotateLeft(ParentOf(ptr)); ptr = _root; } } else { Node sibling = LeftOf(ParentOf(ptr)); if (ColorOf(sibling) == Red) { SetColor(sibling, Black); SetColor(ParentOf(ptr), Red); RotateRight(ParentOf(ptr)); sibling = LeftOf(ParentOf(ptr)); } if (ColorOf(RightOf(sibling)) == Black && ColorOf(LeftOf(sibling)) == Black) { SetColor(sibling, Red); ptr = ParentOf(ptr); } else { if (ColorOf(LeftOf(sibling)) == Black) { SetColor(RightOf(sibling), Black); SetColor(sibling, Red); RotateLeft(sibling); sibling = LeftOf(ParentOf(ptr)); } SetColor(sibling, ColorOf(ParentOf(ptr))); SetColor(ParentOf(ptr), Black); SetColor(LeftOf(sibling), Black); RotateRight(ParentOf(ptr)); ptr = _root; } } } SetColor(ptr, Black); } private void RestoreBalanceAfterInsertion(Node balanceNode) { SetColor(balanceNode, Red); while (balanceNode != null && balanceNode != _root && ColorOf(ParentOf(balanceNode)) == Red) { if (ParentOf(balanceNode) == LeftOf(ParentOf(ParentOf(balanceNode)))) { Node sibling = RightOf(ParentOf(ParentOf(balanceNode))); if (ColorOf(sibling) == Red) { SetColor(ParentOf(balanceNode), Black); SetColor(sibling, Black); SetColor(ParentOf(ParentOf(balanceNode)), Red); balanceNode = ParentOf(ParentOf(balanceNode)); } else { if (balanceNode == RightOf(ParentOf(balanceNode))) { balanceNode = ParentOf(balanceNode); RotateLeft(balanceNode); } SetColor(ParentOf(balanceNode), Black); SetColor(ParentOf(ParentOf(balanceNode)), Red); RotateRight(ParentOf(ParentOf(balanceNode))); } } else { Node sibling = LeftOf(ParentOf(ParentOf(balanceNode))); if (ColorOf(sibling) == Red) { SetColor(ParentOf(balanceNode), Black); SetColor(sibling, Black); SetColor(ParentOf(ParentOf(balanceNode)), Red); balanceNode = ParentOf(ParentOf(balanceNode)); } else { if (balanceNode == LeftOf(ParentOf(balanceNode))) { balanceNode = ParentOf(balanceNode); RotateRight(balanceNode); } SetColor(ParentOf(balanceNode), Black); SetColor(ParentOf(ParentOf(balanceNode)), Red); RotateLeft(ParentOf(ParentOf(balanceNode))); } } } SetColor(_root, Black); } private void RotateLeft(Node node) { if (node != null) { Node right = RightOf(node); node.Right = LeftOf(right); if (LeftOf(right) != null) { LeftOf(right).Parent = node; } right.Parent = ParentOf(node); if (ParentOf(node) == null) { _root = right; } else if (node == LeftOf(ParentOf(node))) { ParentOf(node).Left = right; } else { ParentOf(node).Right = right; } right.Left = node; node.Parent = right; } } private void RotateRight(Node node) { if (node != null) { Node left = LeftOf(node); node.Left = RightOf(left); if (RightOf(left) != null) { RightOf(left).Parent = node; } left.Parent = node.Parent; if (ParentOf(node) == null) { _root = left; } else if (node == RightOf(ParentOf(node))) { ParentOf(node).Right = left; } else { ParentOf(node).Left = left; } left.Right = node; node.Parent = left; } } #endregion #region Safety-Methods // These methods save memory by allowing us to forego sentinel nil nodes, as well as serve as protection against NullReferenceExceptions. /// /// Returns the color of , or Black if it is null. /// /// Node /// The boolean color of , or black if null private static bool ColorOf(Node node) { return node == null || node.Color; } /// /// Sets the color of node to . ///

/// This method does nothing if is null. ///
/// Node to set the color of /// Color (Boolean) private static void SetColor(Node node, bool color) { if (node != null) { node.Color = color; } } /// /// This method returns the left node of , or null if is null. /// /// Node to retrieve the left child from /// Left child of private static Node LeftOf(Node node) { return node?.Left; } /// /// This method returns the right node of , or null if is null. /// /// Node to retrieve the right child from /// Right child of private static Node RightOf(Node node) { return node?.Right; } /// /// Returns the parent node of , or null if is null. /// /// Node to retrieve the parent from /// Parent of private static Node ParentOf(Node node) { return node?.Parent; } #endregion #region Interface Implementations // Method descriptions are not provided as they are already included as part of the interface. public bool ContainsKey(K key) { if (key == null) { throw new ArgumentNullException(nameof(key)); } return GetNode(key) != null; } bool IDictionary.Remove(K key) { int count = _count; Remove(key); return count > _count; } public bool TryGetValue(K key, [MaybeNullWhen(false)] out V value) { if (null == key) { throw new ArgumentNullException(nameof(key)); } Node node = GetNode(key); value = node != null ? node.Value : default; return node != null; } public void Add(KeyValuePair item) { if (item.Key == null) { throw new ArgumentNullException(nameof(item.Key)); } Add(item.Key, item.Value); } public void Clear() { _root = null; _count = 0; } public bool Contains(KeyValuePair item) { if (item.Key == null) { return false; } Node node = GetNode(item.Key); if (node != null) { return node.Key.Equals(item.Key) && node.Value.Equals(item.Value); } return false; } public void CopyTo(KeyValuePair[] array, int arrayIndex) { if (arrayIndex < 0 || array.Length - arrayIndex < this.Count) { throw new ArgumentOutOfRangeException(nameof(arrayIndex)); } SortedList list = GetKeyValues(); int offset = 0; for (int i = arrayIndex; i < array.Length && offset < list.Count; i++) { array[i] = new KeyValuePair(list.Keys[i], list.Values[i]); offset++; } } public bool Remove(KeyValuePair item) { Node node = GetNode(item.Key); if (node == null) { return false; } if (node.Value.Equals(item.Value)) { int count = _count; Remove(item.Key); return count > _count; } return false; } public IEnumerator> GetEnumerator() { return GetKeyValues().GetEnumerator(); } IEnumerator IEnumerable.GetEnumerator() { return GetKeyValues().GetEnumerator(); } public int Count => _count; public ICollection Keys => GetKeyValues().Keys; public ICollection Values => GetKeyValues().Values; public bool IsReadOnly => false; public V this[K key] { get => Get(key); set => Add(key, value); } #endregion #region Private Interface Helper Methods /// /// Returns a sorted list of all the node keys / values in the tree. /// /// List of node keys private SortedList GetKeyValues() { SortedList set = new SortedList(); Queue> queue = new Queue>(); if (_root != null) { queue.Enqueue(_root); } while (queue.Count > 0) { Node node = queue.Dequeue(); set.Add(node.Key, node.Value); if (null != node.Left) { queue.Enqueue(node.Left); } if (null != node.Right) { queue.Enqueue(node.Right); } } return set; } #endregion } /// /// Represents a node in the TreeDictionary which contains a key and value of generic type K and V, respectively. /// /// Key of the node /// Value of the node class Node { public bool Color = true; public Node Left = null; public Node Right = null; public Node Parent = null; public K Key; public V Value; public Node(K key, V value, Node parent) { Key = key; Value = value; Parent = parent; } } }