kstd:: add basic operator[] to flat map

2 files changed, 48 insertions, 11 deletions

diff --git a/libs/kstd/kstd/flat_map.hpp b/libs/kstd/kstd/flat_map.hpp
index dcf73392..1ed999ac 100644
--- a/libs/kstd/kstd/flat_map.hpp
+++ b/libs/kstd/kstd/flat_map.hpp
@@ -451,12 +451,11 @@ namespace kstd
 
     //! Get a reference to the mapped value associated with the given key.
     //!
-    //! @note This overload only participates in overload resolution if the key compare type is transparent.
     //! @warning This function will panic if the key is not found.
     //! @param x The key to look up.
     //! @return A reference to the mapped value.
     template<typename K>
-      requires requires(K const & k, key_type const & l) { typename key_compare::is_transparent; }
+      requires requires { typename key_compare::is_transparent; }
     [[nodiscard]] constexpr auto at(K const & x) -> mapped_type &
     {
       auto found = find(x);
@@ -469,12 +468,12 @@ namespace kstd
     }
 
     //! Get a reference to the mapped value associated with the given key.
-    //! @note This overload only participates in overload resolution if the key compare type is transparent.
+    //!
     //! @warning This function will panic if the key is not found.
     //! @param x The key to look up.
     //! @return A const reference to the mapped value.
     template<typename K>
-      requires requires(K const & k, key_type const & l) { typename key_compare::is_transparent; }
+      requires requires { typename key_compare::is_transparent; }
     [[nodiscard]] auto at(K const & x) const -> mapped_type const &
     {
       auto found = find(x);
@@ -486,6 +485,15 @@ namespace kstd
       os::panic("[kstd::flat_map] Key not found");
     }
 
+    //! 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 const & key) -> mapped_type &
+    {
+      auto result = try_emplace(key);
+      return result.first->second;
+    }
+
     //! Get an iterator to the first element.
     [[nodiscard]] auto begin() noexcept -> iterator
     {
@@ -587,17 +595,16 @@ namespace kstd
     {
       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);
 
       if (found != m_containers.keys.cend() && !m_comparator(value.first, *found) && !m_comparator(*found, value.first))
       {
-        auto offset = std::distance(m_containers.keys.begin(), found);
         return {
           iterator{m_containers, offset},
           false
         };
       }
 
-      auto offset = std::distance(m_containers.keys.begin(), found);
       auto key_iterator = m_containers.keys.begin() + offset;
       auto mapped_iterator = m_containers.values.begin() + offset;
 
@@ -620,16 +627,21 @@ namespace kstd
     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);
+
       if (found != m_containers.keys.cend() && !m_comparator(*found, key) && !m_comparator(key, *found))
       {
-        return {found, false};
+        return {
+          iterator{m_containers, offset},
+          false
+        };
       }
 
-      auto offset = std::distance(m_containers.keys.cbegin(), found);
-      auto insertion_point = m_containers.value.begin() + offset;
+      auto key_iterator = m_containers.keys.begin() + offset;
+      auto mapped_iterator = m_containers.values.begin() + offset;
 
-      auto inserted_key = m_containers.keys.emplace(key);
-      auto inserted_mapped = m_containers.values.emplace(std::forward<Args>(args)...);
+      auto inserted_key = m_containers.keys.emplace(key_iterator, key);
+      auto inserted_mapped = m_containers.values.emplace(key_iterator, std::forward<Args>(args)...);
 
       return {
         iterator{m_containers, inserted_key - m_containers.keys.begin()},
diff --git a/libs/kstd/kstd/flat_map.tests.cpp b/libs/kstd/kstd/flat_map.tests.cpp
index 6a1b6f2a..f0e0f813 100644
--- a/libs/kstd/kstd/flat_map.tests.cpp
+++ b/libs/kstd/kstd/flat_map.tests.cpp
@@ -217,6 +217,31 @@ SCENARIO("Flat Map modifiers", "[flat_map]")
       }
     }
 
+    WHEN("accessing a non-existant element using the subscript operator")
+    {
+      auto & element = map[10];
+
+      THEN("the element is default constructed")
+      {
+        REQUIRE(element == int{});
+      }
+
+      THEN("the element is inserted")
+      {
+        REQUIRE(map.contains(10));
+      }
+
+      AND_WHEN("assigning to that element")
+      {
+        element = 100;
+
+        THEN("the element is assigned")
+        {
+          REQUIRE(map.at(10) == 100);
+        }
+      }
+    }
+
     AND_GIVEN("a populated Flat Map")
     {
       auto other = kstd::flat_map<int, int>{};