2 files changed, 120 insertions, 4 deletions
diff --git a/libs/kstd/kstd/flat_map.hpp b/libs/kstd/kstd/flat_map.hpp
index 3eda42d8..f00956a5 100644
--- a/libs/kstd/kstd/flat_map.hpp
+++ b/libs/kstd/kstd/flat_map.hpp
@@ -820,6 +820,58 @@ namespace kstd
       return const_iterator{m_containers, offset};
     }
 
+    //! Get an iterator to the first element greater than the given key.
+    //!
+    //! @param key The key to look for.
+    //! @return An iterator to the first element greater than the given iff. such an element exists, the end iterator
+    //!         otherwise.
+    [[nodiscard]] constexpr auto upper_bound(key_type const & key) noexcept -> iterator
+    {
+      auto found = std::ranges::upper_bound(m_containers.keys, key, m_comparator);
+      auto offset = static_cast<std::size_t>(std::ranges::distance(m_containers.keys.begin(), found));
+      return iterator{m_containers, offset};
+    }
+
+    //! Get an iterator to the first element greater than the given key.
+    //!
+    //! @param key The key to look for.
+    //! @return An iterator to the first element greater than the given iff. such an element exists, the end iterator
+    //!         otherwise.
+    [[nodiscard]] constexpr auto upper_bound(key_type const & key) const noexcept -> const_iterator
+    {
+      auto found = std::ranges::upper_bound(m_containers.keys, key, m_comparator);
+      auto offset = static_cast<std::size_t>(std::ranges::distance(m_containers.keys.begin(), found));
+      return const_iterator{m_containers, offset};
+    }
+
+    //! Get an iterator to the first element greater than the given key.
+    //!
+    //! @param key The key to look for.
+    //! @return An iterator to the first element greater than the given iff. such an element exists, the end iterator
+    //!         otherwise.
+    template<typename K>
+      requires requires { typename key_compare::is_transparent; }
+    [[nodiscard]] constexpr auto upper_bound(K const & key) noexcept -> iterator
+    {
+      auto found = std::ranges::upper_bound(m_containers.keys, key, m_comparator);
+      auto offset = static_cast<std::size_t>(std::ranges::distance(m_containers.keys.begin(), found));
+      return iterator{m_containers, offset};
+    }
+
+    //! Get an iterator to the first element greater than the given key.
+    //!
+    //! @param key The key to look for.
+    //! @return An iterator to the first element greater than the given iff. such an element exists, the end iterator
+    //!         otherwise.
+    template<typename K>
+      requires requires { typename key_compare::is_transparent; }
+    [[nodiscard]] constexpr auto upper_bound(K const & key) const noexcept -> const_iterator
+    {
+      auto found = std::ranges::upper_bound(m_containers.keys, key, m_comparator);
+      auto offset = static_cast<std::size_t>(std::ranges::distance(m_containers.keys.begin(), found));
+      return iterator{m_containers, offset};
+    }
+
     //! Get the key comparator of this flat map.
     //!
     //! @return The key comparator of this flat map.
diff --git a/libs/kstd/kstd/flat_map.tests.cpp b/libs/kstd/kstd/flat_map.tests.cpp
index 5f29fa14..41281e93 100644
--- a/libs/kstd/kstd/flat_map.tests.cpp
+++ b/libs/kstd/kstd/flat_map.tests.cpp
@@ -736,8 +736,8 @@ SCENARIO("Flat Map queries", "[flat_map]")
       THEN("it returns an iterator to the element")
       {
         REQUIRE(it != map.end());
-        REQUIRE((*it).first == 'c');
-        REQUIRE((*it).second == 20);
+        REQUIRE(it->first == 'c');
+        REQUIRE(it->second == 20);
       }
     }
 
@@ -760,6 +760,38 @@ SCENARIO("Flat Map queries", "[flat_map]")
         REQUIRE(it == map.begin());
       }
     }
+
+    WHEN("finding the upper bound of a key present in the map")
+    {
+      auto it = map.upper_bound('b');
+
+      THEN("it returns an iterator to the next element")
+      {
+        REQUIRE(it != map.end());
+        REQUIRE(it->first == 'c');
+        REQUIRE(it->second == 20);
+      }
+    }
+
+    WHEN("finding the upper bound of a key larger than any in the map")
+    {
+      auto it = map.upper_bound('e');
+
+      THEN("it returns an iterator to the end of the map")
+      {
+        REQUIRE(it == map.end());
+      }
+    }
+
+    WHEN("finding the upper bound of a key smaller than any in the map")
+    {
+      auto it = map.upper_bound('a');
+
+      THEN("it returns an iterator to the first element")
+      {
+        REQUIRE(it == map.begin());
+      }
+    }
   }
 
   GIVEN("A populated Flat Map with a transparent comparator")
@@ -776,8 +808,8 @@ SCENARIO("Flat Map queries", "[flat_map]")
       THEN("it returns an iterator to the element")
       {
         REQUIRE(it != map.end());
-        REQUIRE((*it).first == "c");
-        REQUIRE((*it).second == 20);
+        REQUIRE(it->first == "c");
+        REQUIRE(it->second == 20);
       }
     }
 
@@ -800,5 +832,37 @@ SCENARIO("Flat Map queries", "[flat_map]")
         REQUIRE(it == map.begin());
       }
     }
+
+    WHEN("finding the upper bound of a key present in the map")
+    {
+      auto it = map.upper_bound("b");
+
+      THEN("it returns an iterator to the next element")
+      {
+        REQUIRE(it != map.end());
+        REQUIRE(it->first == "c");
+        REQUIRE(it->second == 20);
+      }
+    }
+
+    WHEN("finding the upper bound of a key larger than any in the map")
+    {
+      auto it = map.upper_bound("e");
+
+      THEN("it returns an iterator to the end of the map")
+      {
+        REQUIRE(it == map.end());
+      }
+    }
+
+    WHEN("finding the upper bound of a key smaller than any in the map")
+    {
+      auto it = map.upper_bound("a");
+
+      THEN("it returns an iterator to the first element")
+      {
+        REQUIRE(it == map.begin());
+      }
+    }
   }
 }