AK: Add RBTree::find_smallest_above_iterator(Key)

1 files changed, 22 insertions, 0 deletions

diff --git a/AK/RedBlackTree.h b/AK/RedBlackTree.h
index 248b257df4..acdc2b9ed4 100644
--- a/AK/RedBlackTree.h
+++ b/AK/RedBlackTree.h
@@ -131,6 +131,20 @@ protected:
         return candidate;
     }
 
+    static Node* find_smallest_not_below(Node* node, K key)
+    {
+        while (node) {
+            if (node->key >= key && (!node->left_child || node->left_child->key < key))
+                return node;
+
+            if (node->key <= key)
+                node = node->right_child;
+            else
+                node = node->left_child;
+        }
+        return node;
+    }
+
     void insert(Node* node)
     {
         VERIFY(node);
@@ -493,6 +507,14 @@ public:
         return ConstIterator(node, static_cast<Node*>(BaseTree::predecessor(node)));
     }
 
+    ConstIterator find_smallest_not_below_iterator(K key) const
+    {
+        auto node = static_cast<Node*>(BaseTree::find_smallest_not_below(this->m_root, key));
+        if (!node)
+            return end();
+        return ConstIterator(node, static_cast<Node*>(BaseTree::predecessor(node)));
+    }
+
     V unsafe_remove(K key)
     {
         auto* node = BaseTree::find(this->m_root, key);