From e7dea10381578628e492fa2d6dc6b77e9821f306 Mon Sep 17 00:00:00 2001
From: Ali Mohammad Pur <ali.mpfard@gmail.com>
Date: Wed, 9 Feb 2022 16:33:44 +0330
Subject: [PATCH] AK: Add RBTree::find_smallest_above_iterator(Key)

---
 AK/RedBlackTree.h | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

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);