Repository: s-yata/marisa-trie
Pull request #77: Fix Agent copy for predictive search
Author: login: glebm
Created: 2025-05-21T07:53:21Z | Merged: 2025-05-21T14:30:23Z

When an agent is used for predictive search, its key points to `state.key_buf_`.

Once we've copied the state, we need to also update `agent.key_.ptr_` to point to the new state's `key_buf_`.


Diff:
diff --git a/lib/marisa/agent.cc b/lib/marisa/agent.cc
index 83f506f..162d49d 100644
--- a/lib/marisa/agent.cc
+++ b/lib/marisa/agent.cc
@@ -4,8 +4,28 @@
 #include "marisa/agent.h"
 #include "marisa/grimoire/trie.h"
 #include "marisa/grimoire/trie/state.h"
+#include "marisa/key.h"
 
 namespace marisa {
+namespace {
+void UpdateAgentAfterCopyingState(const grimoire::trie::State &state,
+                                  Agent &agent) {
+  // Point the agent's key to the newly copied buffer if necessary.
+  switch (state.status_code()) {
+    case grimoire::trie::MARISA_READY_TO_PREDICTIVE_SEARCH:
+    case grimoire::trie::MARISA_END_OF_PREDICTIVE_SEARCH:
+      // In states corresponding to predictive_search, the agent's
+      // key points into the state key buffer. We need to repoint
+      // after copying the state.
+      agent.set_key(state.key_buf().begin(), state.key_buf().size());
+      break;
+    default:
+      // In other states, they key is either null, or points to the
+      // query, so we do not need to repoint it.
+      break;
+  }
+}
+}  // namespace
 
 Agent::Agent() : query_(), key_(), state_() {}
 
@@ -14,12 +34,25 @@ Agent::~Agent() {}
 Agent::Agent(const Agent &other)
     : query_(other.query_),
       key_(other.key_),
-      state_(other.has_state() ? new (std::nothrow) grimoire::trie::State(other.state()) : NULL) {}
+      state_(other.has_state() ? new (std::nothrow)
+                                     grimoire::trie::State(other.state())
+                               : nullptr) {
+  if (other.has_state()) {
+    MARISA_THROW_IF(state_ == nullptr, MARISA_MEMORY_ERROR);
+    UpdateAgentAfterCopyingState(*state_, *this);
+  }
+}
 
 Agent &Agent::operator=(const Agent &other) {
   query_ = other.query_;
   key_ = other.key_;
-  state_.reset(other.has_state() ? new (std::nothrow) grimoire::trie::State(other.state()) : NULL);
+  if (other.has_state()) {
+    state_.reset(new (std::nothrow) grimoire::trie::State(other.state()));
+    MARISA_THROW_IF(state_ == nullptr, MARISA_MEMORY_ERROR);
+    UpdateAgentAfterCopyingState(*state_, *this);
+  } else {
+    state_ = nullptr;
+  }
   return *this;
 }
 
diff --git a/tests/marisa-test.cc b/tests/marisa-test.cc
index 9490645..81b5579 100644
--- a/tests/marisa-test.cc
+++ b/tests/marisa-test.cc
@@ -4,6 +4,7 @@
 #include <ctime>
 #include <random>
 #include <sstream>
+#include <string>
 #include <utility>
 #include <vector>
 
@@ -252,6 +253,18 @@ void TestCommonPrefixSearch(const marisa::Trie &trie,
   }
 }
 
+void TestCommonPrefixSearchAgentCopy(const marisa::Trie &trie,
+    const marisa::Keyset &keyset) {
+  if (keyset.empty()) return;
+  marisa::Agent agent;
+  agent.set_query(keyset[0].ptr(), keyset[0].length());
+  ASSERT(trie.common_prefix_search(agent));
+  const std::string original_agent_key(agent.key().ptr(), agent.key().length());
+  marisa::Agent agent_copy = agent;
+  trie.common_prefix_search(agent);
+  ASSERT(std::string(agent_copy.key().ptr(), agent_copy.key().length()) == original_agent_key);
+}
+
 void TestPredictiveSearch(const marisa::Trie &trie,
     const marisa::Keyset &keyset) {
   marisa::Agent agent;
@@ -275,9 +288,11 @@ void TestPredictiveSearchAgentCopy(const marisa::Trie &trie,
 
     std::vector<marisa::Agent> agent_copies;
     std::vector<std::size_t> ids;
+    std::vector<std::string> keys;
     while (trie.predictive_search(agent)) {
       ASSERT(agent.key().id() > keyset[i].id());
       ids.push_back(agent.key().id());
+      keys.emplace_back(agent.key().ptr(), agent.key().length());
 
       // Tests copy constructor.
       agent_copies.push_back(agent);
@@ -290,9 +305,11 @@ void TestPredictiveSearchAgentCopy(const marisa::Trie &trie,
       agent_copy = agent_copies[i];
 
       ASSERT(agent_copy.key().id() == ids[i]);
+      ASSERT(std::string(agent_copy.key().ptr(), agent_copy.key().length()) == keys[i]);
       if (i + 1 < agent_copies.size()) {
         ASSERT(trie.predictive_search(agent_copy));
         ASSERT(agent_copy.key().id() == ids[i + 1]);
+        ASSERT(std::string(agent_copy.key().ptr(), agent_copy.key().length()) == keys[i + 1]);
       } else {
         ASSERT(!trie.predictive_search(agent_copy));
       }
@@ -345,6 +362,7 @@ void TestTrie(int num_tries, marisa::TailMode tail_mode,
 
   TestLookup(trie, keyset);
   TestCommonPrefixSearch(trie, keyset);
+  TestCommonPrefixSearchAgentCopy(trie, keyset);
   TestPredictiveSearch(trie, keyset);
   TestPredictiveSearchAgentCopy(trie, keyset);
   TestPredictiveSearchAgentMove(trie, keyset);


Reviews:
user: login: jmr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T09:57:59Z
user: login: jmr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T09:59:34Z
user: login: jmr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T10:00:55Z
user: login: jmr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T12:22:37Z
user: login: glebm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T12:23:21Z
user: login: glebm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T12:23:55Z
user: login: glebm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T12:26:18Z
user: login: glebm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T12:27:48Z

Review comments:
diff_hunk: @@ -252,6 +253,17 @@ void TestCommonPrefixSearch(const marisa::Trie &trie,
   }
 }
 
+void TestCommonPrefixSearchAgentCopy(const marisa::Trie &trie,
path: tests/marisa-test.cc
user: login: jmr
body: Where is this called?
created_at: 2025-05-21T09:57:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -4,8 +4,23 @@
 #include "marisa/agent.h"
 #include "marisa/grimoire/trie.h"
 #include "marisa/grimoire/trie/state.h"
+#include "marisa/key.h"
 
 namespace marisa {
+namespace {
+void UpdateAgentAfterCopyingState(const grimoire::trie::State &state,
+                                  Agent &agent) {
+  // Point the agent's key to the newly copied buffer if necessary.
path: lib/marisa/agent.cc
user: login: jmr
body: Can you explain a bit more? Why are these the only two status codes requiring a copy?
created_at: 2025-05-21T09:59:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -14,12 +29,25 @@ Agent::~Agent() {}
 Agent::Agent(const Agent &other)
     : query_(other.query_),
       key_(other.key_),
-      state_(other.has_state() ? new (std::nothrow) grimoire::trie::State(other.state()) : NULL) {}
+      state_(other.has_state() ? new (std::nothrow)
+                                     grimoire::trie::State(other.state())
+                               : nullptr) {
+  if (other.has_state()) {
+    MARISA_THROW_IF(state_ == nullptr, MARISA_MEMORY_ERROR);
path: lib/marisa/agent.cc
user: login: jmr
body: Maybe we can come back and use `std::optional` after we bump to C++17.
created_at: 2025-05-21T10:00:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -252,6 +253,17 @@ void TestCommonPrefixSearch(const marisa::Trie &trie,
   }
 }
 
+void TestCommonPrefixSearchAgentCopy(const marisa::Trie &trie,
path: tests/marisa-test.cc
user: login: jmr
body: And does this test fail without the code fix?
created_at: 2025-05-21T12:22:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -252,6 +253,17 @@ void TestCommonPrefixSearch(const marisa::Trie &trie,
   }
 }
 
+void TestCommonPrefixSearchAgentCopy(const marisa::Trie &trie,
path: tests/marisa-test.cc
user: login: glebm
body: It wasn't, fixed.
created_at: 2025-05-21T12:23:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -4,8 +4,23 @@
 #include "marisa/agent.h"
 #include "marisa/grimoire/trie.h"
 #include "marisa/grimoire/trie/state.h"
+#include "marisa/key.h"
 
 namespace marisa {
+namespace {
+void UpdateAgentAfterCopyingState(const grimoire::trie::State &state,
+                                  Agent &agent) {
+  // Point the agent's key to the newly copied buffer if necessary.
path: lib/marisa/agent.cc
user: login: glebm
body: I've now added more comments to the code explaining this
created_at: 2025-05-21T12:23:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -14,12 +29,25 @@ Agent::~Agent() {}
 Agent::Agent(const Agent &other)
     : query_(other.query_),
       key_(other.key_),
-      state_(other.has_state() ? new (std::nothrow) grimoire::trie::State(other.state()) : NULL) {}
+      state_(other.has_state() ? new (std::nothrow)
+                                     grimoire::trie::State(other.state())
+                               : nullptr) {
+  if (other.has_state()) {
+    MARISA_THROW_IF(state_ == nullptr, MARISA_MEMORY_ERROR);
path: lib/marisa/agent.cc
user: login: glebm
body: That would make sense. It will require exposing the `State` class to the end user though.
created_at: 2025-05-21T12:26:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -252,6 +253,17 @@ void TestCommonPrefixSearch(const marisa::Trie &trie,
   }
 }
 
+void TestCommonPrefixSearchAgentCopy(const marisa::Trie &trie,
path: tests/marisa-test.cc
user: login: glebm
body: > And does this test fail without the code fix?

This test doesn't because in the common prefix search case we don't need to do anything.
I've added it for completeness.

The new assertions in the predictive search test do fail with this PR.

created_at: 2025-05-21T12:27:48Z
author_association: CONTRIBUTOR

PR comments:
user: login: s-yata
created_at: 2025-05-21T08:24:42Z
body: I missed a problem.

The current Agent copy constructor and copy assignment can silently fail --- it rarely happens but it can.
`new (std::nothrow)` returns nullptr on failure.

I'm not sure if these can throw exceptons or not.
Is it bad-mannered?

author_association: OWNER
user: login: glebm
created_at: 2025-05-21T08:34:34Z
body: It's OK for copy-constructor/assignment to throw (but not for move constructor/assignment)
I've added `MARISA_THROW_IF`
author_association: CONTRIBUTOR
user: login: s-yata
created_at: 2025-05-21T14:40:08Z
body: Thanks!
author_association: OWNER