kstd: add more subscript operator overloads to flat map

3 files changed, 171 insertions, 29 deletions

diff --git a/libs/kstd/kstd/bits/concepts.hpp b/libs/kstd/kstd/bits/concepts.hpp
index 065b6f07..3e89caa9 100644
--- a/libs/kstd/kstd/bits/concepts.hpp
+++ b/libs/kstd/kstd/bits/concepts.hpp
@@ -23,6 +23,9 @@ namespace kstd::bits
     suspect.deallocate(suspect.allocate(1), 1);
   };
 
+  template<typename Comparator>
+  concept transparent = requires { typename Comparator::is_transparent; };
+
 }  // namespace kstd::bits
 
 #endif
diff --git a/libs/kstd/kstd/flat_map.hpp b/libs/kstd/kstd/flat_map.hpp
index 1ed999ac..5da34083 100644
--- a/libs/kstd/kstd/flat_map.hpp
+++ b/libs/kstd/kstd/flat_map.hpp
@@ -455,7 +455,7 @@ namespace kstd
     //! @param x The key to look up.
     //! @return A reference to the mapped value.
     template<typename K>
-      requires requires { typename key_compare::is_transparent; }
+      requires bits::transparent<KeyCompare>
     [[nodiscard]] constexpr auto at(K const & x) -> mapped_type &
     {
       auto found = find(x);
@@ -473,7 +473,7 @@ namespace kstd
     //! @param x The key to look up.
     //! @return A const reference to the mapped value.
     template<typename K>
-      requires requires { typename key_compare::is_transparent; }
+      requires bits::transparent<KeyCompare>
     [[nodiscard]] auto at(K const & x) const -> mapped_type const &
     {
       auto found = find(x);
@@ -494,6 +494,26 @@ namespace kstd
       return result.first->second;
     }
 
+    //! Get a reference to the mapped value associated with the given key, or insert a default one if none exists.
+    //!
+    //! @param key The key to look up.
+    [[nodiscard]] constexpr auto operator[](key_type && key) -> mapped_type &
+    {
+      auto result = try_emplace(std::move(key));
+      return result.first->second;
+    }
+
+    //! Get a reference to the mapped value associated with the given key, or insert a default one if none exists.
+    //!
+    //! @param key The key to look up.
+    template<typename K>
+      requires bits::transparent<KeyCompare>
+    [[nodiscard]] constexpr auto operator[](K && key) -> mapped_type &
+    {
+      auto result = try_emplace(std::forward<K>(key));
+      return result.first->second;
+    }
+
     //! Get an iterator to the first element.
     [[nodiscard]] auto begin() noexcept -> iterator
     {
@@ -594,10 +614,11 @@ namespace kstd
       requires std::constructible_from<value_type, Args...>
     {
       auto value = value_type{std::forward<Args>(args)...};
-      auto found = std::ranges::lower_bound(m_containers.keys, value.first, m_comparator);
-      auto offset = std::distance(m_containers.keys.begin(), found);
+      auto key_iterator = std::ranges::lower_bound(m_containers.keys, value.first, m_comparator);
+      auto offset = std::distance(m_containers.keys.begin(), key_iterator);
 
-      if (found != m_containers.keys.cend() && !m_comparator(value.first, *found) && !m_comparator(*found, value.first))
+      if (key_iterator != m_containers.keys.cend() && !m_comparator(value.first, *key_iterator) &&
+          !m_comparator(*key_iterator, value.first))
       {
         return {
           iterator{m_containers, offset},
@@ -605,14 +626,11 @@ namespace kstd
         };
       }
 
-      auto key_iterator = m_containers.keys.begin() + offset;
-      auto mapped_iterator = m_containers.values.begin() + offset;
-
-      auto inserted_key = m_containers.keys.insert(key_iterator, std::move(value.first));
-      auto inserted_mapped = m_containers.values.insert(mapped_iterator, std::move(value.second));
+      m_containers.keys.insert(key_iterator, std::move(value.first));
+      m_containers.values.insert(m_containers.values.begin() + offset, std::move(value.second));
 
       return {
-        iterator{m_containers, inserted_key - m_containers.keys.begin()},
+        iterator{m_containers, offset},
         true
       };
     }
@@ -626,10 +644,11 @@ namespace kstd
     template<typename... Args>
     auto try_emplace(key_type const & key, Args &&... args) -> std::pair<iterator, bool>
     {
-      auto found = std::ranges::lower_bound(m_containers.keys, key, m_comparator);
-      auto offset = std::distance(m_containers.keys.begin(), found);
+      auto key_iterator = std::ranges::lower_bound(m_containers.keys, key, m_comparator);
+      auto offset = std::distance(m_containers.keys.begin(), key_iterator);
 
-      if (found != m_containers.keys.cend() && !m_comparator(*found, key) && !m_comparator(key, *found))
+      if (key_iterator != m_containers.keys.cend() && !m_comparator(*key_iterator, key) &&
+          !m_comparator(key, *key_iterator))
       {
         return {
           iterator{m_containers, offset},
@@ -637,14 +656,76 @@ namespace kstd
         };
       }
 
-      auto key_iterator = m_containers.keys.begin() + offset;
-      auto mapped_iterator = m_containers.values.begin() + offset;
+      m_containers.keys.emplace(key_iterator, key);
+      m_containers.values.emplace(m_containers.values.begin() + offset, std::forward<Args>(args)...);
+
+      return {
+        iterator{m_containers, offset},
+        true
+      };
+    }
+
+    //! Try to insert a element for the given key into this map.
+    //!
+    //! This function does nothing if the key is already present.
+    //!
+    //! @param key The key to insert a value for.
+    //! @param args The arguments to use to construct the mapped value.
+    template<typename... Args>
+    auto try_emplace(key_type && key, Args &&... args) -> std::pair<iterator, bool>
+    {
+      auto key_iterator = std::ranges::lower_bound(m_containers.keys, key, m_comparator);
+      auto offset = std::distance(m_containers.keys.begin(), key_iterator);
+
+      if (key_iterator != m_containers.keys.cend() && !m_comparator(*key_iterator, key) &&
+          !m_comparator(key, *key_iterator))
+      {
+        return {
+          iterator{m_containers, offset},
+          false
+        };
+      }
+
+      m_containers.keys.emplace(key_iterator, std::move(key));
+      m_containers.values.emplace(m_containers.values.begin() + offset, std::forward<Args>(args)...);
+
+      return {
+        iterator{m_containers, offset},
+        true
+      };
+    }
+
+    //! Try to insert a element for the given key into this map.
+    //!
+    //! This function does nothing if the key is already present.
+    //!
+    //! @param key The key to insert a value for.
+    //! @param args The arguments to use to construct the mapped value.
+    template<typename K, typename... Args>
+      requires bits::transparent<KeyCompare> &&                  //
+               std::constructible_from<key_type, K> &&           //
+               std::constructible_from<mapped_type, Args...> &&  //
+               (!std::convertible_to<K &&, const_iterator>) &&   //
+               (!std::convertible_to<K &&, iterator>)            //
+    auto try_emplace(K && key, Args &&... args) -> std::pair<iterator, bool>
+    {
+      auto key_iterator = std::ranges::lower_bound(m_containers.keys, key, m_comparator);
+      auto offset = std::distance(m_containers.keys.begin(), key_iterator);
+
+      if (key_iterator != m_containers.keys.cend() && !m_comparator(*key_iterator, key) &&
+          !m_comparator(key, *key_iterator))
+      {
+        return {
+          iterator{m_containers, offset},
+          false
+        };
+      }
 
-      auto inserted_key = m_containers.keys.emplace(key_iterator, key);
-      auto inserted_mapped = m_containers.values.emplace(key_iterator, std::forward<Args>(args)...);
+      m_containers.keys.emplace(key_iterator, std::forward<K>(key));
+      m_containers.values.emplace(m_containers.values.begin() + offset, std::forward<Args>(args)...);
 
       return {
-        iterator{m_containers, inserted_key - m_containers.keys.begin()},
+        iterator{m_containers, offset},
         true
       };
     }
@@ -701,7 +782,7 @@ namespace kstd
     //! @param x The key to look up.
     //! @return An iterator to the element with the equivalent key, or end() if no such element is found.
     template<typename K>
-      requires requires { typename key_compare::is_transparent; }
+      requires bits::transparent<KeyCompare>
     [[nodiscard]] auto find(K const & x) noexcept -> iterator
     {
       auto found = std::ranges::lower_bound(m_containers.keys, x, m_comparator);
@@ -718,7 +799,7 @@ namespace kstd
     //! @param x The key to look up.
     //! @return An iterator to the element with the equivalent key, or end() if no such element is found.
     template<typename K>
-      requires requires { typename key_compare::is_transparent; }
+      requires bits::transparent<KeyCompare>
     [[nodiscard]] auto find(K const & x) const noexcept -> const_iterator
     {
       auto found = std::ranges::lower_bound(m_containers.keys, x, m_comparator);
@@ -749,7 +830,7 @@ namespace kstd
     //! @param key The key to look for.
     //! @return 0 if the key does not exist, 1 otherwise.
     template<typename K>
-      requires requires { typename key_compare::is_transparent; }
+      requires bits::transparent<KeyCompare>
     [[nodiscard]] constexpr auto count(key_type const & key) const noexcept -> std::size_t
     {
       if (contains(key))
@@ -774,7 +855,7 @@ namespace kstd
     //! @param key The key to check.
     //! @return true iff. the key is found, false otherwise.
     template<typename K>
-      requires requires { typename key_compare::is_transparent; }
+      requires bits::transparent<KeyCompare>
     [[nodiscard]] constexpr auto contains(K const & key) const noexcept -> bool
     {
       return find(key) != cend();
@@ -810,7 +891,7 @@ namespace kstd
     //! @return An iterator to the first element not less than the given key iff. such an element exists, the end
     //!         iterator otherwise.
     template<typename K>
-      requires requires { typename key_compare::is_transparent; }
+      requires bits::transparent<KeyCompare>
     [[nodiscard]] constexpr auto lower_bound(K const & key) noexcept -> iterator
     {
       auto found = std::ranges::lower_bound(m_containers.keys, key, m_comparator);
@@ -824,7 +905,7 @@ namespace kstd
     //! @return An iterator to the first element not less than the given key iff. such an element exists, the end
     //!         iterator otherwise.
     template<typename K>
-      requires requires { typename key_compare::is_transparent; }
+      requires bits::transparent<KeyCompare>
     [[nodiscard]] constexpr auto lower_bound(K const & key) const noexcept -> const_iterator
     {
       auto found = std::ranges::lower_bound(m_containers.keys, key, m_comparator);
@@ -862,7 +943,7 @@ namespace kstd
     //! @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; }
+      requires bits::transparent<KeyCompare>
     [[nodiscard]] constexpr auto upper_bound(K const & key) noexcept -> iterator
     {
       auto found = std::ranges::upper_bound(m_containers.keys, key, m_comparator);
@@ -876,7 +957,7 @@ namespace kstd
     //! @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; }
+      requires bits::transparent<KeyCompare>
     [[nodiscard]] constexpr auto upper_bound(K const & key) const noexcept -> const_iterator
     {
       auto found = std::ranges::upper_bound(m_containers.keys, key, m_comparator);
@@ -914,7 +995,7 @@ namespace kstd
     //! @param key The key to look for.
     //! @return A pair of iterators describing the, possibly empty, range of elements.
     template<typename K>
-      requires requires { typename key_compare::is_transparent; }
+      requires bits::transparent<KeyCompare>
     [[nodiscard]] constexpr auto equal_range(K const & key) noexcept -> std::pair<iterator, iterator>
     {
       auto found = std::ranges::equal_range(m_containers.keys, key, m_comparator);
@@ -928,7 +1009,7 @@ namespace kstd
     //! @param key The key to look for.
     //! @return A pair of iterators describing the, possibly empty, range of elements.
     template<typename K>
-      requires requires { typename key_compare::is_transparent; }
+      requires bits::transparent<KeyCompare>
     [[nodiscard]] constexpr auto equal_range(K const & key) const noexcept -> std::pair<const_iterator, const_iterator>
     {
       auto found = std::ranges::equal_range(m_containers.keys, key, m_comparator);
diff --git a/libs/kstd/kstd/flat_map.tests.cpp b/libs/kstd/kstd/flat_map.tests.cpp
index f0e0f813..a53b73db 100644
--- a/libs/kstd/kstd/flat_map.tests.cpp
+++ b/libs/kstd/kstd/flat_map.tests.cpp
@@ -205,6 +205,28 @@ SCENARIO("Flat Map modifiers", "[flat_map]")
       }
     }
 
+    WHEN("emplacing multiple elements out of order")
+    {
+      auto [it, inserted] = map.emplace(2, 200);
+      auto [it2, inserted2] = map.emplace(1, 100);
+      auto [it3, inserted3] = map.emplace(3, 300);
+
+      THEN("the map contains the new elements")
+      {
+        REQUIRE(inserted);
+        REQUIRE(inserted2);
+        REQUIRE(inserted3);
+        REQUIRE(map.contains(1));
+        REQUIRE(map.contains(2));
+        REQUIRE(map.contains(3));
+      }
+
+      THEN("the elements are ordered correctly in the map")
+      {
+        REQUIRE(map.keys() == kstd::vector{1, 2, 3});
+      }
+    }
+
     WHEN("emplacing an existing element")
     {
       map.emplace(1, 100);
@@ -303,6 +325,22 @@ SCENARIO("Flat Map element access", "[flat_map]")
         REQUIRE_THROWS_AS(map.at(4), kstd::tests::os_panic);
       }
     }
+
+    WHEN("accessing a existing element using the subscript operator")
+    {
+      auto & val = map[2];
+
+      THEN("it returns a reference to the mapped value")
+      {
+        REQUIRE(val == 20);
+      }
+
+      THEN("the mapped value can be modified")
+      {
+        val = 200;
+        REQUIRE(map[2] == 200);
+      }
+    }
   }
 
   GIVEN("A const populated Flat Map")
@@ -955,5 +993,25 @@ SCENARIO("Flat Map queries", "[flat_map]")
         REQUIRE(end == map.end());
       }
     }
+
+    WHEN("accessing an existing element using the subscript operator")
+    {
+      auto & val = map["d"];
+
+      THEN("it returns a reference to the mapped value")
+      {
+        REQUIRE(val == 20);
+      }
+    }
+
+    WHEN("accessing a non-existent element using the subscript operator")
+    {
+      auto & val = map["g"];
+
+      THEN("a new, default constructed element is inserted")
+      {
+        REQUIRE(val == 0);
+      }
+    }
   }
 }