kstd: implement flat map erase

2 files changed, 173 insertions, 0 deletions

diff --git a/libs/kstd/kstd/flat_map.hpp b/libs/kstd/kstd/flat_map.hpp
index 5da34083..71e142f7 100644
--- a/libs/kstd/kstd/flat_map.hpp
+++ b/libs/kstd/kstd/flat_map.hpp
@@ -229,11 +229,30 @@ namespace kstd
           , m_index{index}
       {}
 
+      template<bool OtherIsConst>
+      constexpr flat_map_iterator(containers & containers, flat_map_iterator<OtherIsConst> const & other)
+        requires(!IsConst)
+          : m_containers{&containers}
+          , m_index{other.m_index}
+      {}
+
+      template<bool OtherIsConst>
+      constexpr flat_map_iterator(containers const & containers, flat_map_iterator<OtherIsConst> const & other)
+        requires(IsConst)
+          : m_containers{&containers}
+          , m_index{other.m_index}
+      {}
+
       [[nodiscard]] constexpr auto key_iterator() const noexcept
       {
         return m_containers->keys.begin() + m_index;
       }
 
+      [[nodiscard]] constexpr auto value_iterator() const noexcept
+      {
+        return m_containers->values.begin() + m_index;
+      }
+
       std::conditional_t<IsConst, containers const, containers> * m_containers;
       std::size_t m_index{std::numeric_limits<std::size_t>::max()};
     };
@@ -730,6 +749,62 @@ namespace kstd
       };
     }
 
+    //! Erase the element at the given position from this flat map.
+    //!
+    //! @param position The position of the element to remove.
+    //! @return An iterator to the element after the one that was removed.
+    constexpr auto erase(iterator position) -> iterator
+    {
+      return erase(static_cast<const_iterator>(position));
+    }
+
+    //! Erase the element at the given position from this flat map.
+    //!
+    //! @param position The position of the element to remove.
+    //! @return An iterator to the element after the one that was removed.
+    constexpr auto erase(const_iterator position) -> iterator
+    {
+      m_containers.keys.erase(position.key_iterator());
+      m_containers.values.erase(position.value_iterator());
+      return iterator{m_containers, position};
+    }
+
+    //! Erase the elements in a given range from this flat map.
+    //!
+    //! @param first The begin of the range to erase.
+    //! @param last The end of the range to erase.
+    //! @return An iterator to the element after the last removed one.
+    constexpr auto erase(const_iterator first, const_iterator last) -> iterator
+    {
+      auto key_iterator = m_containers.keys.erase(first.key_iterator(), last.key_iterator());
+      m_containers.values.erase(first.value_iterator(), last.value_iterator());
+      auto offset = std::distance(m_containers.keys.begin(), key_iterator);
+      return iterator{m_containers, offset};
+    }
+
+    //! Erase all elements with keys equivalent to a given one from this flat map.
+    //!
+    //! @param key The key of the elements to remove.
+    //! @return The number of removed elements.
+    constexpr auto erase(key_type const & key) -> size_type
+    {
+      return erase<key_type const &>(key);
+    }
+
+    //! Erase all elements with keys equivalent to a given one from this flat map.
+    //!
+    //! @param key The key of the elements to remove.
+    //! @return The number of removed elements.
+    template<typename K>
+      requires std::same_as<key_type, std::remove_cvref_t<K>> || bits::transparent<KeyCompare>
+    constexpr auto erase(K && key) -> size_type
+    {
+      auto [first, last] = equal_range(std::forward<K>(key));
+      auto count = std::ranges::distance(first, last);
+      erase(first, last);
+      return count;
+    }
+
     //! Swap the contents of this flat map with the one of a different one.
     //!
     //! @param other The flat map instance to swap with.
diff --git a/libs/kstd/kstd/flat_map.tests.cpp b/libs/kstd/kstd/flat_map.tests.cpp
index a53b73db..59b5fb99 100644
--- a/libs/kstd/kstd/flat_map.tests.cpp
+++ b/libs/kstd/kstd/flat_map.tests.cpp
@@ -291,6 +291,104 @@ SCENARIO("Flat Map modifiers", "[flat_map]")
       }
     }
   }
+
+  GIVEN("A populated Flat Map")
+  {
+    auto map = kstd::flat_map<char, int>{};
+    map.emplace('b', 10);
+    map.emplace('d', 20);
+    map.emplace('f', 30);
+
+    auto old_size = map.size();
+
+    WHEN("erasing an existing element")
+    {
+      auto it = map.erase(map.find('b'));
+
+      THEN("it returns an iterator to the next element")
+      {
+        REQUIRE(it == map.find('d'));
+      }
+
+      THEN("the element is removed")
+      {
+        REQUIRE_FALSE(map.contains('b'));
+      }
+
+      THEN("the size is reduced")
+      {
+        REQUIRE(map.size() < old_size);
+      }
+    }
+
+    WHEN("erasing multiple elements")
+    {
+      auto it = map.erase(map.find('b'), map.find('f'));
+
+      THEN("it returns an iterator to the next element")
+      {
+        REQUIRE(it == map.find('f'));
+      }
+
+      THEN("the elements are removed")
+      {
+        REQUIRE_FALSE(map.contains('b'));
+        REQUIRE_FALSE(map.contains('d'));
+        REQUIRE(map.contains('f'));
+      }
+
+      THEN("the size is reduced")
+      {
+        REQUIRE(map.size() < old_size);
+      }
+    }
+
+    WHEN("erasing all elements")
+    {
+      auto it = map.erase(map.begin(), map.end());
+
+      THEN("it returns the end iterator")
+      {
+        REQUIRE(it == map.end());
+      }
+
+      THEN("the size is 0")
+      {
+        REQUIRE(map.size() == 0);
+      }
+    }
+
+    WHEN("erasing an existing element by key")
+    {
+      auto count = map.erase('b');
+
+      THEN("one element is removed")
+      {
+        REQUIRE(count == 1);
+        REQUIRE_FALSE(map.contains('b'));
+      }
+
+      THEN("the size is reduced")
+      {
+        REQUIRE(map.size() < old_size);
+      }
+    }
+
+    WHEN("erasing a non-existent element by key")
+    {
+      auto count = map.erase('g');
+
+      THEN("no elements are removed")
+      {
+        REQUIRE(count == 0);
+      }
+
+      THEN("the size does not change")
+      {
+        REQUIRE(map.size() == old_size);
+      }
+    }
+  }
 }
 
 SCENARIO("Flat Map element access", "[flat_map]")