Merge pull request #5263 from ygorshenin/fix-mem-trie

[base] MemTrie is switched to unique_ptrs.

base/base_tests/mem_trie_test.cpp

@@ -2,31 +2,50 @@
 
 #include "base/mem_trie.hpp"
 
-#include "std/algorithm.hpp"
-#include "std/string.hpp"
-#include "std/utility.hpp"
-#include "std/vector.hpp"
+#include <algorithm>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace
+{
+using Trie = my::MemTrie<string, int>;
+using Data = std::vector<std::pair<std::string, int>>;
+
+void GetTrieContents(Trie const & trie, Data & data)
+{
+  data.clear();
+  trie.ForEachInTrie([&data](std::string const & k, int v) { data.emplace_back(k, v); });
+  std::sort(data.begin(), data.end());
+}
 
 UNIT_TEST(MemTrie_Basic)
 {
-  vector<pair<string, int>> data = {{"roger", 3},
-                                    {"amy", 1},
-                                    {"emma", 1},
-                                    {"ann", 1},
-                                    {"rob", 1},
-                                    {"roger", 2},
-                                    {"", 0},
-                                    {"roger", 1}};
-  my::MemTrie<string, int> trie;
+  Data data = {{"roger", 3}, {"amy", 1},   {"emma", 1}, {"ann", 1},
+               {"rob", 1},   {"roger", 2}, {"", 0},     {"roger", 1}};
+  Trie trie;
+  TEST_EQUAL(trie.GetNumNodes(), 1, ());
+
   for (auto const & p : data)
     trie.Add(p.first, p.second);
+  TEST_EQUAL(trie.GetNumNodes(), 16, ());
 
-  vector<pair<string, int>> trie_data;
-  trie.ForEach([&trie_data](string const & k, int v)
-               {
-                 trie_data.emplace_back(k, v);
-               });
-  sort(data.begin(), data.end());
-  sort(trie_data.begin(), trie_data.end());
-  TEST_EQUAL(data, trie_data, ());
+  std::sort(data.begin(), data.end());
+
+  Data contents;
+  GetTrieContents(trie, contents);
+  TEST_EQUAL(contents, data, ());
+
+  TEST_EQUAL(trie.GetNumNodes(), 16, ());
+
+  Trie newTrie(move(trie));
+
+  TEST_EQUAL(trie.GetNumNodes(), 1, ());
+  GetTrieContents(trie, contents);
+  TEST(contents.empty(), ());
+
+  TEST_EQUAL(newTrie.GetNumNodes(), 16, ());
+  GetTrieContents(newTrie, contents);
+  TEST_EQUAL(contents, data, ());
 }
+}  // namespace

base/mem_trie.hpp

@@ -1,61 +1,72 @@
 #pragma once
 
 #include "base/macros.hpp"
+#include "base/stl_add.hpp"
 
+#include <cstdint>
 #include <map>
+#include <memory>
 #include <vector>
 
 namespace my
 {
 // This class is a simple in-memory trie which allows to add
 // key-value pairs and then traverse them in a sorted order.
-template <typename TString, typename TValue>
+template <typename String, typename Value>
 class MemTrie
 {
 public:
   MemTrie() = default;
+  MemTrie(MemTrie && rhs) { *this = std::move(rhs); }
 
-  MemTrie(MemTrie && other) : m_root(move(other.m_root))
+  MemTrie & operator=(MemTrie && rhs)
   {
-    m_numNodes = other.m_numNodes;
-    other.m_numNodes = 0;
-  }
-
-  MemTrie & operator=(MemTrie && other)
-  {
-    m_root = move(other.m_root);
-    m_numNodes = other.m_numNodes;
-    other.m_numNodes = 0;
+    m_root = std::move(rhs.m_root);
+    m_numNodes = rhs.m_numNodes;
+    rhs.m_numNodes = 1;
     return *this;
   }
 
   // Adds a key-value pair to the trie.
-  void Add(TString const & key, TValue const & value)
+  void Add(String const & key, Value const & value)
   {
-    Node * cur = &m_root;
+    auto * cur = &m_root;
     for (auto const & c : key)
     {
-      size_t numNewNodes;
-      cur = cur->GetMove(c, numNewNodes);
-      m_numNodes += numNewNodes;
+      bool created;
+      cur = &cur->GetMove(c, created);
+      if (created)
+        ++m_numNodes;
     }
     cur->AddValue(value);
   }
 
   // Traverses all key-value pairs in the trie and calls |toDo| on each of them.
   template <typename ToDo>
-  void ForEach(ToDo && toDo)
+  void ForEachInTrie(ToDo && toDo) const
   {
-    TString prefix;
-    ForEach(&m_root, prefix, std::forward<ToDo>(toDo));
+    String prefix;
+    ForEachInSubtree(m_root, prefix, std::forward<ToDo>(toDo));
   }
 
+  // Calls |toDo| for each key-value pair in a node that is reachable
+  // by |prefix| from the trie root. Does nothing if such node does
+  // not exist.
   template <typename ToDo>
-  void ForEachInSubtree(TString prefix, ToDo && toDo) const
+  void ForEachInNode(String const & prefix, ToDo && toDo) const
   {
-    Node const * node = MoveTo(prefix);
-    if (node)
-      ForEach(node, prefix, std::forward<ToDo>(toDo));
+    if (auto const * root = MoveTo(prefix))
+      ForEachInNode(*root, prefix, std::forward<ToDo>(toDo));
+  }
+
+  // Calls |toDo| for each key-value pair in a subtree that is
+  // reachable by |prefix| from the trie root. Does nothing if such
+  // subtree does not exist.
+  template <typename ToDo>
+  void ForEachInSubtree(String prefix, ToDo && toDo) const
+  {
+    if (auto const * root = MoveTo(prefix))
+      ForEachInSubtree(*root, prefix, std::forward<ToDo>(toDo));
   }
 
   size_t GetNumNodes() const { return m_numNodes; }
@@ -63,72 +74,76 @@ public:
 private:
   struct Node
   {
-    using TChar = typename TString::value_type;
+    using Char = typename String::value_type;
 
     Node() = default;
+    Node(Node && /* rhs */) = default;
 
-    Node(Node && other) = default;
+    Node & operator=(Node && /* rhs */) = default;
 
-    ~Node()
+    Node & GetMove(Char const & c, bool & created)
     {
-      for (auto const & move : m_moves)
-        delete move.second;
-    }
-
-    Node & operator=(Node && other) = default;
-
-    Node * GetMove(TChar const & c, size_t & numNewNodes)
-    {
-      numNewNodes = 0;
-      Node *& node = m_moves[c];
+      auto & node = m_moves[c];
       if (!node)
       {
-        node = new Node();
-        ++numNewNodes;
+        node = my::make_unique<Node>();
+        created = true;
       }
-      return node;
+      else
+      {
+        created = false;
+      }
+      return *node;
     }
 
-    void AddValue(TValue const & value) { m_values.push_back(value); }
+    void AddValue(Value const & value) { m_values.push_back(value); }
 
-    std::map<TChar, Node *> m_moves;
-    std::vector<TValue> m_values;
+    std::map<Char, std::unique_ptr<Node>> m_moves;
+    std::vector<Value> m_values;
 
     DISALLOW_COPY(Node);
   };
 
-  Node const * MoveTo(TString const & key) const
+  Node const * MoveTo(String const & key) const
   {
-    Node const * cur = &m_root;
+    auto const * cur = &m_root;
     for (auto const & c : key)
     {
       auto const it = cur->m_moves.find(c);
       if (it == cur->m_moves.end())
         return nullptr;
-      cur = it->second;
+      cur = it->second.get();
     }
     return cur;
   }
 
+  // Calls |toDo| for each key-value pair in a |node| that is
+  // reachable by |prefix| from the trie root.
   template <typename ToDo>
-  void ForEach(Node const * root, TString & prefix, ToDo && toDo) const
+  void ForEachInNode(Node const & node, String const & prefix, ToDo && toDo) const
   {
-    if (!root->m_values.empty())
-    {
-      for (auto const & value : root->m_values)
-        toDo(prefix, value);
-    }
+    for (auto const & value : node.m_values)
+      toDo(prefix, value);
+  }
 
-    for (auto const & move : root->m_moves)
+  // Calls |toDo| for each key-value pair in subtree where |node| is a
+  // root of the subtree. |prefix| is a path from the trie root to the
+  // |node|.
+  template <typename ToDo>
+  void ForEachInSubtree(Node const & node, String & prefix, ToDo && toDo) const
+  {
+    ForEachInNode(node, prefix, toDo);
+
+    for (auto const & move : node.m_moves)
     {
       prefix.push_back(move.first);
-      ForEach(move.second, prefix, toDo);
+      ForEachInSubtree(*move.second, prefix, toDo);
       prefix.pop_back();
     }
   }
 
   Node m_root;
-  size_t m_numNodes = 0;
+  size_t m_numNodes = 1;
 
   DISALLOW_COPY(MemTrie);
 };  // class MemTrie

base/stl_add.hpp

@@ -14,8 +14,7 @@ std::unique_ptr<T> make_unique(Args &&... args)
 {
   return std::unique_ptr<T>(new T(std::forward<Args>(args)...));
 }
-}
-
+}  // namespace my
 
 template <class ContainerT> class BackInsertFunctor
 {