[base] MemTrie::Erase().

base/base_tests/mem_trie_test.cpp

@@ -7,45 +7,108 @@
 #include <utility>
 #include <vector>
 
+using namespace std;
+
 namespace
 {
-using Trie = my::MemTrie<std::string, my::VectorValues<int>>;
-using Data = std::vector<std::pair<std::string, int>>;
+using Key = string;
+using Value = int;
+using Trie = my::MemTrie<Key, my::VectorValues<Value>>;
+using Data = vector<pair<Key, Value>>;
 
-void GetTrieContents(Trie const & trie, Data & data)
+Data GetTrieContents(Trie const & trie)
 {
-  data.clear();
-  trie.ForEachInTrie([&data](std::string const & k, int v) { data.emplace_back(k, v); });
-  std::sort(data.begin(), data.end());
+  Data data;
+  trie.ForEachInTrie([&data](string const & k, int v) { data.emplace_back(k, v); });
+  sort(data.begin(), data.end());
+  return data;
 }
 
-UNIT_TEST(MemTrie_Basic)
+class MemTrieTest
 {
-  Data data = {{"roger", 3}, {"amy", 1},   {"emma", 1}, {"ann", 1},
-               {"rob", 1},   {"roger", 2}, {"", 0},     {"roger", 1}};
-  Trie trie;
-  TEST_EQUAL(trie.GetNumNodes(), 1, ());
+public:
+  Data GetActualContents() const { return ::GetTrieContents(m_trie); }
 
-  for (auto const & p : data)
-    trie.Add(p.first, p.second);
-  TEST_EQUAL(trie.GetNumNodes(), 16, ());
+  Data GetExpectedContents() const { return m_data; }
 
-  std::sort(data.begin(), data.end());
+  size_t GetNumNodes() const { return m_trie.GetNumNodes(); }
 
-  Data contents;
-  GetTrieContents(trie, contents);
-  TEST_EQUAL(contents, data, ());
+  void Add(Key const & key, Value const & value)
+  {
+    m_trie.Add(key, value);
 
-  TEST_EQUAL(trie.GetNumNodes(), 16, ());
+    auto const kv = make_pair(key, value);
+    auto it = lower_bound(m_data.begin(), m_data.end(), kv);
+    m_data.insert(it, kv);
+  }
 
-  Trie newTrie(move(trie));
+  void Erase(Key const & key, Value const & value)
+  {
+    m_trie.Erase(key, value);
 
-  TEST_EQUAL(trie.GetNumNodes(), 1, ());
-  GetTrieContents(trie, contents);
-  TEST(contents.empty(), ());
+    auto const kv = make_pair(key, value);
+    auto it = lower_bound(m_data.begin(), m_data.end(), kv);
+    if (it != m_data.end() && *it == kv)
+      m_data.erase(it);
+  }
+
+protected:
+  Trie m_trie;
+  Data m_data;
+};
+
+UNIT_CLASS_TEST(MemTrieTest, Basic)
+{
+  TEST_EQUAL(GetNumNodes(), 1, ());
+
+  Data const data = {{"roger", 3}, {"amy", 1},   {"emma", 1}, {"ann", 1},
+                     {"rob", 1},   {"roger", 2}, {"", 0},     {"roger", 1}};
+  for (auto const & kv : data)
+    Add(kv.first, kv.second);
+  TEST_EQUAL(GetNumNodes(), 16, ());
+
+  TEST_EQUAL(GetExpectedContents(), GetActualContents(), ());
+  TEST_EQUAL(GetNumNodes(), 16, ());
+
+  Trie newTrie(move(m_trie));
+
+  TEST_EQUAL(m_trie.GetNumNodes(), 1, ());
+  TEST(GetTrieContents(m_trie).empty(), ());
 
   TEST_EQUAL(newTrie.GetNumNodes(), 16, ());
-  GetTrieContents(newTrie, contents);
-  TEST_EQUAL(contents, data, ());
+  TEST_EQUAL(GetTrieContents(newTrie), GetExpectedContents(), ());
+}
+
+UNIT_CLASS_TEST(MemTrieTest, KeysRemoval)
+{
+  TEST_EQUAL(GetNumNodes(), 1, ());
+
+  Data const data = {{"bobby", 1}, {"robby", 2}, {"rob", 3}, {"r", 4}, {"robert", 5}, {"bob", 6}};
+
+  for (auto const & kv : data)
+    Add(kv.first, kv.second);
+
+  TEST_EQUAL(GetNumNodes(), 14, ());
+  TEST_EQUAL(GetExpectedContents(), GetActualContents(), ());
+
+  Erase("r", 3);
+  TEST_EQUAL(GetNumNodes(), 14, ());
+  TEST_EQUAL(GetExpectedContents(), GetActualContents(), ());
+
+  Erase("r", 4);
+  TEST_EQUAL(GetNumNodes(), 14, ());
+  TEST_EQUAL(GetExpectedContents(), GetActualContents(), ());
+
+  Erase("robert", 5);
+  TEST_EQUAL(GetNumNodes(), 11, ());
+  TEST_EQUAL(GetExpectedContents(), GetActualContents(), ());
+
+  Erase("rob", 3);
+  TEST_EQUAL(GetNumNodes(), 11, ());
+  TEST_EQUAL(GetExpectedContents(), GetActualContents(), ());
+
+  Erase("robby", 2);
+  TEST_EQUAL(GetNumNodes(), 6, ());
+  TEST_EQUAL(GetExpectedContents(), GetActualContents(), ());
 }
 }  // namespace

base/mem_trie.hpp

@@ -1,9 +1,11 @@
 #pragma once
 
+#include "base/assert.hpp"
 #include "base/macros.hpp"
 #include "base/stl_add.hpp"
 
 #include <algorithm>
+#include <cstddef>
 #include <cstdint>
 #include <functional>
 #include <map>
@@ -46,6 +48,10 @@ public:
     return *node;
   }
 
+  void EraseSubtree(Char const & c) { m_subtrees.erase(c); }
+
+  bool Empty() const { return m_subtrees.empty(); }
+
   void Clear() { m_subtrees.clear(); }
 
 private:
@@ -61,12 +67,22 @@ struct VectorValues
     m_values.emplace_back(std::forward<V>(v));
   }
 
+  template <typename V>
+  void Erase(V const & v)
+  {
+    auto const it = find(m_values.begin(), m_values.end(), v);
+    if (it != m_values.end())
+      m_values.erase(it);
+  }
+
   template <typename ToDo>
   void ForEach(ToDo && toDo) const
   {
     std::for_each(m_values.begin(), m_values.end(), std::forward<ToDo>(toDo));
   }
 
+  bool Empty() const { return m_values.empty(); }
+
   void Clear() { m_values.clear(); }
 
   std::vector<Value> m_values;
@@ -138,6 +154,12 @@ public:
     cur->Add(std::forward<Value>(value));
   }
 
+  template <typename Value>
+  void Erase(String const & key, Value const & value)
+  {
+    return Erase(m_root, 0 /* level */, key, value);
+  }
+
   // Traverses all key-value pairs in the trie and calls |toDo| on each of them.
   template <typename ToDo>
   void ForEachInTrie(ToDo && toDo) const
@@ -189,12 +211,24 @@ private:
       return m_moves.GetOrCreateSubtree(c, created);
     }
 
+    Node * GetMove(Char const & c) const { return m_moves.GetSubtree(c); }
+
+    void EraseMove(Char const & c) { m_moves.EraseSubtree(c); }
+
     template <typename Value>
     void Add(Value && value)
     {
       m_values.Add(std::forward<Value>(value));
     }
 
+    template <typename Value>
+    void EraseValue(Value const & value)
+    {
+      m_values.Erase(value);
+    }
+
+    bool Empty() const { return m_moves.Empty() && m_values.Empty(); }
+
     void Clear()
     {
       m_moves.Clear();
@@ -212,13 +246,34 @@ private:
     auto const * cur = &m_root;
     for (auto const & c : key)
     {
-      cur = cur->m_moves.GetSubtree(c);
+      cur = cur->GetMove(c);
       if (!cur)
         break;
     }
     return cur;
   }
 
+  template <typename Value>
+  void Erase(Node & root, size_t level, String const & key, Value const & value)
+  {
+    if (level == key.size())
+    {
+      root.EraseValue(value);
+      return;
+    }
+
+    ASSERT_LESS(level, key.size(), ());
+    auto * child = root.GetMove(key[level]);
+    if (!child)
+      return;
+    Erase(*child, level + 1, key, value);
+    if (child->Empty())
+    {
+      root.EraseMove(key[level]);
+      --m_numNodes;
+    }
+  }
+
   // Calls |toDo| for each key-value pair in a |node| that is
   // reachable by |prefix| from the trie root.
   template <typename ToDo>