kstd: implement erase_if for flat map

4 files changed, 58 insertions, 7 deletions

diff --git a/kernel/kernel/filesystem/devfs/filesystem.cpp b/kernel/kernel/filesystem/devfs/filesystem.cpp
index c8447a2d..cdc3bd1c 100644
--- a/kernel/kernel/filesystem/devfs/filesystem.cpp
+++ b/kernel/kernel/filesystem/devfs/filesystem.cpp
@@ -11,10 +11,10 @@
 
 #include <kapi/filesystem.hpp>
 
+#include <kstd/flat_map.hpp>
 #include <kstd/memory.hpp>
 #include <kstd/result.hpp>
 #include <kstd/system_error.hpp>
-#include <kstd/vector.hpp>
 
 #include <algorithm>
 #include <optional>
@@ -50,7 +50,7 @@ namespace kernel::filesystem::devfs
     std::ranges::for_each(device_number_registry::get().all(), [&](auto const & numbered) {
       if (auto device = numbered.device.lock())
       {
-        m_inodes.emplace_back(numbered.name, kstd::make_shared<device_inode>(device));
+        m_inodes.emplace(numbered.name, kstd::make_shared<device_inode>(device));
       }
     });
 
@@ -97,7 +97,7 @@ namespace kernel::filesystem::devfs
   {
     if (auto device = numbered.device.lock())
     {
-      m_inodes.emplace_back(numbered.name, kstd::make_shared<device_inode>(device));
+      m_inodes.emplace(numbered.name, kstd::make_shared<device_inode>(device));
     }
   }
 
diff --git a/kernel/kernel/filesystem/devfs/filesystem.hpp b/kernel/kernel/filesystem/devfs/filesystem.hpp
index 8d2199e3..986dc27f 100644
--- a/kernel/kernel/filesystem/devfs/filesystem.hpp
+++ b/kernel/kernel/filesystem/devfs/filesystem.hpp
@@ -14,11 +14,9 @@
 #include <kstd/result.hpp>
 #include <kstd/string.hpp>
 #include <kstd/system_error.hpp>
-#include <kstd/vector.hpp>
 
 #include <optional>
 #include <string_view>
-#include <utility>
 
 namespace kernel::filesystem::devfs
 {
@@ -51,8 +49,7 @@ namespace kernel::filesystem::devfs
     auto on_device_numbered(device_number_registry::entry const & numbered) -> void override;
     auto on_device_unnumbered(device_number_registry::entry const & numbered) -> void override;
 
-    // TODO: change to flat_map again once flat_map supports erase.
-    kstd::vector<std::pair<kstd::string, kstd::shared_ptr<device_inode>>> m_inodes{};
+    kstd::flat_map<kstd::string, kstd::shared_ptr<device_inode>> m_inodes{};
   };
 }  // namespace kernel::filesystem::devfs
 
diff --git a/libs/kstd/kstd/flat_map.hpp b/libs/kstd/kstd/flat_map.hpp
index f3cf827a..e140f7dd 100644
--- a/libs/kstd/kstd/flat_map.hpp
+++ b/libs/kstd/kstd/flat_map.hpp
@@ -1121,6 +1121,13 @@ namespace kstd
       key_compare m_comparator;
     };
 
+    template<typename EKeyType, typename EMappedType, typename EKeyCompare, typename EKeyContainerType,
+             typename EMappedContainerType, typename Predicate>
+    constexpr auto friend erase_if(
+        flat_map<EKeyType, EMappedType, EKeyCompare, EKeyContainerType, EMappedContainerType> & map,
+        Predicate predicate)
+        -> flat_map<EKeyType, EMappedType, EKeyCompare, EKeyContainerType, EMappedContainerType>::size_type;
+
     constexpr auto do_sort() -> void
     {
       auto zipped = std::views::zip(m_containers.keys, m_containers.values);
@@ -1161,6 +1168,22 @@ namespace kstd
       -> flat_map<typename KeyContainer::value_type, typename MappedContainer::value_type, Compare, KeyContainer,
                   MappedContainer>;
 
+  template<typename KeyType, typename MappedType, typename KeyCompare, typename KeyContainerType,
+           typename MappedContainerType, typename Predicate>
+  constexpr auto erase_if(flat_map<KeyType, MappedType, KeyCompare, KeyContainerType, MappedContainerType> & map,
+                          Predicate predicate)
+      -> flat_map<KeyType, MappedType, KeyCompare, KeyContainerType, MappedContainerType>::size_type
+  {
+    auto zipped = std::views::zip(map.m_containers.keys, map.m_containers.values);
+    auto removed = std::ranges::remove_if(zipped, predicate, [](auto const & e) {
+      return
+          typename flat_map<KeyType, MappedType, KeyCompare, KeyContainerType, MappedContainerType>::const_reference{e};
+    });
+    auto count = removed.size();
+    map.erase(map.end() - count, map.end());
+    return count;
+  }
+
 }  // namespace kstd
 
 #endif
\ No newline at end of file
diff --git a/libs/kstd/kstd/flat_map.tests.cpp b/libs/kstd/kstd/flat_map.tests.cpp
index 59b5fb99..1c16b5a3 100644
--- a/libs/kstd/kstd/flat_map.tests.cpp
+++ b/libs/kstd/kstd/flat_map.tests.cpp
@@ -388,6 +388,37 @@ SCENARIO("Flat Map modifiers", "[flat_map]")
         REQUIRE(map.size() == old_size);
       }
     }
+
+    WHEN("erasing a single element by predicate")
+    {
+      auto count = erase_if(map, [](auto const & v) { return v.first == '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 multiple elements by predicate")
+    {
+      auto count = erase_if(map, [](auto const & v) { return v.first == 'b' || v.first == 'f'; });
+
+      THEN("two elements are removed")
+      {
+        REQUIRE(count == 2);
+      }
+
+      THEN("the size is reduced")
+      {
+        REQUIRE(map.size() < old_size);
+      }
+    }
   }
 }