Sync from upstream.

Descriptions:
======================================================================
Fix the bug in PostingListAccessor found by the Icing Monkey test
======================================================================
Add the logic to handle fatal errors from IcingDynamicTrie to avoid crashing
======================================================================
Clear out the dead code IcingDynamicTrie::Compact
======================================================================
[MemoryMappedFile][RemapV2][1/x] Add factory method for
MemoryMappedFile
======================================================================
[MemoryMappedFile][RemapV2][2/x] Create GrowAndRemapIfNecessary and change factory method
======================================================================
[MemoryMappedFile][RemapV2][3/x] Migrate FileBackedVector to use GrowAndRemapIfNecessary
======================================================================
Add JNI latency for query latency stats breakdown
======================================================================

Bug: 247671531
Bug: 247929909
Bug: 253282365

Change-Id: Ic3b88d0f044edacfe2dfeb08fa381b2186c731cb

3 files changed, 111 insertions, 201 deletions

diff --git a/icing/legacy/index/icing-dynamic-trie.cc b/icing/legacy/index/icing-dynamic-trie.cc
index 38aaac0..378b666 100644
--- a/icing/legacy/index/icing-dynamic-trie.cc
+++ b/icing/legacy/index/icing-dynamic-trie.cc
@@ -75,6 +75,7 @@
 #include <memory>
 #include <utility>
 
+#include "icing/absl_ports/canonical_errors.h"
 #include "icing/legacy/core/icing-packed-pod.h"
 #include "icing/legacy/core/icing-string-util.h"
 #include "icing/legacy/core/icing-timer.h"
@@ -86,6 +87,7 @@
 #include "icing/util/i18n-utils.h"
 #include "icing/util/logging.h"
 #include "icing/util/math-util.h"
+#include "icing/util/status-macros.h"
 
 using std::inplace_merge;
 using std::lower_bound;
@@ -308,7 +310,7 @@ class IcingDynamicTrie::IcingDynamicTrieStorage {
   // REQUIRES: nodes_left() > 0.
   Node *AllocNode();
   // REQUIRES: nexts_left() >= kMaxNextArraySize.
-  Next *AllocNextArray(int size);
+  libtextclassifier3::StatusOr<Next *> AllocNextArray(int size);
   void FreeNextArray(Next *next, int log2_size);
   // REQUIRES: suffixes_left() >= strlen(suffix) + 1 + value_size()
   uint32_t MakeSuffix(const char *suffix, const void *value,
@@ -727,10 +729,11 @@ IcingDynamicTrie::Node *IcingDynamicTrie::IcingDynamicTrieStorage::AllocNode() {
   return GetMutableNode(hdr_.hdr.num_nodes() - 1);
 }
 
-IcingDynamicTrie::Next *
+libtextclassifier3::StatusOr<IcingDynamicTrie::Next *>
 IcingDynamicTrie::IcingDynamicTrieStorage::AllocNextArray(int size) {
   if (size > kMaxNextArraySize) {
-    ICING_LOG(FATAL) << "Array size exceeds the max 'next' array size";
+    return absl_ports::InternalError(
+        "Array size exceeds the max 'next' array size");
   }
 
   if (nexts_left() < static_cast<uint32_t>(kMaxNextArraySize)) {
@@ -1298,50 +1301,6 @@ void IcingDynamicTrie::OnSleep() {
   UpdateCrc();
 }
 
-IcingDynamicTrie::NewValueMap::~NewValueMap() {}
-
-bool IcingDynamicTrie::Compact(
-    const NewValueMap &old_tvi_to_new_value, IcingDynamicTrie *out,
-    std::unordered_map<uint32_t, uint32_t> *old_to_new_tvi) const {
-  if (old_to_new_tvi == nullptr) {
-    ICING_LOG(ERROR) << "TVI is null";
-  }
-
-  if (!is_initialized()) {
-    ICING_LOG(FATAL) << "DynamicTrie not initialized";
-  }
-
-  PropertyReadersAll prop_readers(*this);
-
-  old_to_new_tvi->clear();
-  old_to_new_tvi->rehash(size() * 2);
-
-  for (Iterator it_all(*this, ""); it_all.IsValid(); it_all.Advance()) {
-    uint32_t value_index = it_all.GetValueIndex();
-    const void *new_value = old_tvi_to_new_value.GetNewValue(value_index);
-    if (!new_value) continue;
-
-    uint32_t new_value_index;
-    if (!out->Insert(it_all.GetKey(), new_value, &new_value_index, false)) {
-      return false;
-    }
-
-    old_to_new_tvi->insert({value_index, new_value_index});
-
-    // Copy properties.
-    for (size_t i = 0; i < prop_readers.size(); i++) {
-      if (prop_readers.HasProperty(i, value_index)) {
-        if (!out->SetProperty(new_value_index, i)) {
-          // Ouch. We need to bail.
-          return false;
-        }
-      }
-    }
-  }
-
-  return true;
-}
-
 uint32_t IcingDynamicTrie::size() const {
   if (!is_initialized()) {
     ICING_LOG(FATAL) << "DynamicTrie not initialized";
@@ -1554,9 +1513,11 @@ bool IcingDynamicTrie::SortNextArray(const Node *node) {
   return true;
 }
 
-bool IcingDynamicTrie::Insert(const char *key, const void *value,
-                              uint32_t *value_index, bool replace,
-                              bool *pnew_key) {
+libtextclassifier3::Status IcingDynamicTrie::Insert(const char *key,
+                                                    const void *value,
+                                                    uint32_t *value_index,
+                                                    bool replace,
+                                                    bool *pnew_key) {
   if (!is_initialized()) {
     ICING_LOG(FATAL) << "DynamicTrie not initialized";
   }
@@ -1572,8 +1533,7 @@ bool IcingDynamicTrie::Insert(const char *key, const void *value,
   if (!(storage_->nodes_left() >= 2 + key_len + 1 &&
         storage_->nexts_left() >= 2 + key_len + 1 + kMaxNextArraySize &&
         storage_->suffixes_left() >= key_len + 1 + value_size())) {
-    // No more space left.
-    return false;
+    return absl_ports::ResourceExhaustedError("No more space left");
   }
 
   uint32_t best_node_index;
@@ -1615,7 +1575,7 @@ bool IcingDynamicTrie::Insert(const char *key, const void *value,
             storage_->GetSuffixIndex(prev_suffix_cur + 1), value_size());
         memcpy(mutable_prev_suffix_cur, value, value_size());
       }
-      return true;
+      return libtextclassifier3::Status::OK;
     }
 
     if (*prev_suffix_cur == *key_cur) {
@@ -1629,7 +1589,7 @@ bool IcingDynamicTrie::Insert(const char *key, const void *value,
     int common_len = prev_suffix_cur - prev_suffix;
     for (int i = 0; i < common_len; i++) {
       // Create a single-branch child node.
-      Next *split_next = storage_->AllocNextArray(1);
+      ICING_ASSIGN_OR_RETURN(Next * split_next, storage_->AllocNextArray(1));
       split_node->set_next_index(storage_->GetNextArrayIndex(split_next));
       split_node->set_is_leaf(false);
       split_node->set_log2_num_children(0);
@@ -1641,7 +1601,7 @@ bool IcingDynamicTrie::Insert(const char *key, const void *value,
     }
 
     // Fill a split.
-    Next *split_next = storage_->AllocNextArray(2);
+    ICING_ASSIGN_OR_RETURN(Next * split_next, storage_->AllocNextArray(2));
     split_node->set_next_index(storage_->GetNextArrayIndex(split_next));
     split_node->set_is_leaf(false);
     split_node->set_log2_num_children(1);
@@ -1700,7 +1660,7 @@ bool IcingDynamicTrie::Insert(const char *key, const void *value,
     Next *new_next = cur_next;
     if (next_len == (next_array_buffer_size)) {
       // Allocate a new, larger, array.
-      new_next = storage_->AllocNextArray(next_len + 1);
+      ICING_ASSIGN_OR_RETURN(new_next, storage_->AllocNextArray(next_len + 1));
       memcpy(new_next, cur_next, sizeof(Next) * next_len);
     }
 
@@ -1721,7 +1681,8 @@ bool IcingDynamicTrie::Insert(const char *key, const void *value,
 
       // 8 == log2(256)
       if (log2_num_children >= 8) {
-        ICING_LOG(FATAL) << "Number of children exceeds the max allowed size";
+        return absl_ports::InternalError(
+            "Number of children exceeds the max allowed size");
       }
 
       mutable_best_node->set_log2_num_children(log2_num_children + 1);
@@ -1735,7 +1696,7 @@ bool IcingDynamicTrie::Insert(const char *key, const void *value,
   storage_->inc_num_keys();
 
   if (pnew_key) *pnew_key = true;
-  return true;
+  return libtextclassifier3::Status::OK;
 }
 
 const void *IcingDynamicTrie::GetValueAtIndex(uint32_t value_index) const {
diff --git a/icing/legacy/index/icing-dynamic-trie.h b/icing/legacy/index/icing-dynamic-trie.h
index e55bfc2..18748d7 100644
--- a/icing/legacy/index/icing-dynamic-trie.h
+++ b/icing/legacy/index/icing-dynamic-trie.h
@@ -41,6 +41,8 @@
 #include <unordered_map>
 #include <vector>
 
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/legacy/core/icing-compat.h"
 #include "icing/legacy/core/icing-packed-pod.h"
 #include "icing/legacy/index/icing-filesystem.h"
@@ -312,23 +314,6 @@ class IcingDynamicTrie : public IIcingStorage {
   // Potentially about to get nuked.
   void OnSleep() override;
 
-  // Compact trie into out for value indices present in old_tvi_to_new_value.
-  class NewValueMap {
-   public:
-    virtual ~NewValueMap();
-
-    // Returns the new value we want to assign to the entry at old
-    // value index. We don't take ownership of the pointer.
-    virtual const void *GetNewValue(uint32_t old_value_index) const = 0;
-  };
-  // Compacts this trie. This drops all deleted keys, drops all keys for which
-  // old_tvi_to_new_value returns nullptr, updates values to be the values
-  // returned by old_tvi_to_new_value, rewrites tvis, and saves the results into
-  // the trie given in 'out'. 'old_to_new_tvi' is be populated with a mapping of
-  // old value_index to new value_index.
-  bool Compact(const NewValueMap &old_tvi_to_new_value, IcingDynamicTrie *out,
-               std::unordered_map<uint32_t, uint32_t> *old_to_new_tvi) const;
-
   // Insert value at key. If key already exists and replace == true,
   // replaces old value with value. We take a copy of value.
   //
@@ -336,18 +321,22 @@ class IcingDynamicTrie : public IIcingStorage {
   // value_index. This can then be used with SetValueAtIndex
   // below. value_index is not valid past a Clear/Read/Write.
   //
-  // Returns false if there is no space left in the trie.
-  //
   // REQUIRES: value a buffer of size value_size()
-  bool Insert(const char *key, const void *value) {
+  //
+  // Returns:
+  //   OK on success
+  //   RESOURCE_EXHAUSTED if no disk space is available
+  //   INTERNAL_ERROR if there are inconsistencies in the dynamic trie.
+  libtextclassifier3::Status Insert(const char *key, const void *value) {
     return Insert(key, value, nullptr, true, nullptr);
   }
-  bool Insert(const char *key, const void *value, uint32_t *value_index,
-              bool replace) {
+  libtextclassifier3::Status Insert(const char *key, const void *value,
+                                    uint32_t *value_index, bool replace) {
     return Insert(key, value, value_index, replace, nullptr);
   }
-  bool Insert(const char *key, const void *value, uint32_t *value_index,
-              bool replace, bool *pnew_key);
+  libtextclassifier3::Status Insert(const char *key, const void *value,
+                                    uint32_t *value_index, bool replace,
+                                    bool *pnew_key);
 
   // Get a value returned by Insert value_index. This points to the
   // value in the trie. The pointer is immutable and always valid
diff --git a/icing/legacy/index/icing-dynamic-trie_test.cc b/icing/legacy/index/icing-dynamic-trie_test.cc
index 850fcdc..dd63784 100644
--- a/icing/legacy/index/icing-dynamic-trie_test.cc
+++ b/icing/legacy/index/icing-dynamic-trie_test.cc
@@ -28,6 +28,7 @@
 #include "gtest/gtest.h"
 #include "icing/legacy/core/icing-string-util.h"
 #include "icing/legacy/index/icing-filesystem.h"
+#include "icing/testing/common-matchers.h"
 #include "icing/testing/random-string.h"
 #include "icing/testing/tmp-directory.h"
 #include "icing/util/logging.h"
@@ -48,17 +49,6 @@ constexpr uint32_t kNumKeys = ABSL_ARRAYSIZE(kKeys);
 
 class IcingDynamicTrieTest : public ::testing::Test {
  protected:
-  class NewValueMap : public IcingDynamicTrie::NewValueMap {
-   public:
-    const void* GetNewValue(uint32_t old_value_index) const override {
-      const_cast<NewValueMap*>(this)->buf_ = old_value_index;
-      return &buf_;
-    }
-
-   private:
-    uint32_t buf_;
-  };
-
   // Output trie stats to stderr.
   static void StatsDump(const IcingDynamicTrie& trie) {
     IcingDynamicTrie::Stats stats;
@@ -71,8 +61,7 @@ class IcingDynamicTrieTest : public ::testing::Test {
     for (uint32_t i = 0; i < kNumKeys; i++) {
       key.clear();
       IcingStringUtil::SStringAppendF(&key, 0, "%u+%010u", i % 2, i);
-      bool inserted = trie->Insert(key.c_str(), &i);
-      ASSERT_TRUE(inserted);
+      ASSERT_THAT(trie->Insert(key.c_str(), &i), IsOk());
     }
   }
 
@@ -138,7 +127,7 @@ TEST_F(IcingDynamicTrieTest, Simple) {
   ASSERT_TRUE(trie.Init());
 
   for (uint32_t i = 0; i < kNumKeys; i++) {
-    ASSERT_TRUE(trie.Insert(kKeys[i].data(), &i));
+    ASSERT_THAT(trie.Insert(kKeys[i].data(), &i), IsOk());
 
     uint32_t val;
     bool found = trie.Find(kKeys[i].data(), &val);
@@ -179,7 +168,7 @@ TEST_F(IcingDynamicTrieTest, Iterator) {
   ASSERT_TRUE(trie.Init());
 
   for (uint32_t i = 0; i < kNumKeys; i++) {
-    ASSERT_TRUE(trie.Insert(kKeys[i].data(), &i));
+    ASSERT_THAT(trie.Insert(kKeys[i].data(), &i), IsOk());
   }
 
   // Should get the entire trie.
@@ -289,7 +278,7 @@ TEST_F(IcingDynamicTrieTest, IteratorReverse) {
   ASSERT_TRUE(trie.Init());
 
   for (uint32_t i = 0; i < kNumKeys; i++) {
-    ASSERT_TRUE(trie.Insert(kKeys[i].data(), &i));
+    ASSERT_THAT(trie.Insert(kKeys[i].data(), &i), IsOk());
   }
 
   // Should get the entire trie.
@@ -405,7 +394,7 @@ TEST_F(IcingDynamicTrieTest, IteratorLoadTest) {
   // Randomly generate 1024 terms.
   for (int i = 0; i < 1024; ++i) {
     std::string term = RandomString("abcdefg", 5, &random) + std::to_string(i);
-    ASSERT_TRUE(trie.Insert(term.c_str(), &i));
+    ASSERT_THAT(trie.Insert(term.c_str(), &i), IsOk());
     exp_key_values.push_back(std::make_pair(term, i));
   }
   // Lexicographically sort the expected keys.
@@ -448,7 +437,7 @@ TEST_F(IcingDynamicTrieTest, Persistence) {
     ASSERT_TRUE(trie.Init());
 
     for (uint32_t i = 0; i < kCommonEnglishWordArrayLen; i++) {
-      ASSERT_TRUE(trie.Insert(kCommonEnglishWords[i].data(), &i));
+      ASSERT_THAT(trie.Insert(kCommonEnglishWords[i].data(), &i), IsOk());
     }
     // Explicitly omit sync.
 
@@ -462,7 +451,7 @@ TEST_F(IcingDynamicTrieTest, Persistence) {
     EXPECT_EQ(0U, trie.size());
 
     for (uint32_t i = 0; i < kCommonEnglishWordArrayLen; i++) {
-      ASSERT_TRUE(trie.Insert(kCommonEnglishWords[i].data(), &i));
+      ASSERT_THAT(trie.Insert(kCommonEnglishWords[i].data(), &i), IsOk());
     }
     trie.Sync();
 
@@ -510,7 +499,7 @@ TEST_F(IcingDynamicTrieTest, PersistenceShared) {
 
     uint32_t next_reopen = kCommonEnglishWordArrayLen / 16;
     for (uint32_t i = 0; i < kCommonEnglishWordArrayLen; i++) {
-      ASSERT_TRUE(trie.Insert(kCommonEnglishWords[i].data(), &i));
+      ASSERT_THAT(trie.Insert(kCommonEnglishWords[i].data(), &i), IsOk());
 
       if (i == next_reopen) {
         ASSERT_NE(0u, trie.UpdateCrc());
@@ -574,7 +563,7 @@ TEST_F(IcingDynamicTrieTest, Sync) {
     ASSERT_TRUE(trie.Init());
 
     for (uint32_t i = 0; i < kNumKeys; i++) {
-      ASSERT_TRUE(trie.Insert(kKeys[i].data(), &i));
+      ASSERT_THAT(trie.Insert(kKeys[i].data(), &i), IsOk());
 
       uint32_t val;
       bool found = trie.Find(kKeys[i].data(), &val);
@@ -632,7 +621,7 @@ TEST_F(IcingDynamicTrieTest, LimitsSmall) {
   ASSERT_LT(3U, kNumKeys);
 
   for (uint32_t i = 0; i < 3; i++) {
-    ASSERT_TRUE(trie.Insert(kKeys[i].data(), &i)) << i;
+    ASSERT_THAT(trie.Insert(kKeys[i].data(), &i), IsOk()) << i;
 
     uint32_t val;
     bool found = trie.Find(kKeys[i].data(), &val);
@@ -641,7 +630,8 @@ TEST_F(IcingDynamicTrieTest, LimitsSmall) {
   }
 
   uint32_t val = 3;
-  EXPECT_FALSE(trie.Insert(kKeys[3].data(), &val));
+  EXPECT_THAT(trie.Insert(kKeys[3].data(), &val),
+              StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
 
   StatsDump(trie);
   PrintTrie(trie);
@@ -667,7 +657,7 @@ TEST_F(IcingDynamicTrieTest, DISABLEDFingerprintedKeys) {
     IcingStringUtil::SStringAppendF(
         &key, 1000, "content://gmail-ls/account/conversation/%u/message/%u", i,
         10 * i);
-    ASSERT_TRUE(trie.Insert(key.c_str(), &i));
+    ASSERT_THAT(trie.Insert(key.c_str(), &i), IsOk());
 
     // Now compute a fingerprint.
     uint64_t fpkey = tc3farmhash::Fingerprint64(key);
@@ -679,8 +669,8 @@ TEST_F(IcingDynamicTrieTest, DISABLEDFingerprintedKeys) {
       fpkey /= 255;
     }
     fpkey_base255[8] = '\0';
-    ASSERT_TRUE(
-        triefp.Insert(reinterpret_cast<const char*>(fpkey_base255), &i));
+    ASSERT_THAT(triefp.Insert(reinterpret_cast<const char*>(fpkey_base255), &i),
+                IsOk());
 
     // Sync periodically to gauge write locality.
     if ((i + 1) % (kNumKeys / 10) == 0) {
@@ -830,42 +820,6 @@ TEST_F(IcingDynamicTrieTest, ClearSingleProperty) {
   }
 }
 
-TEST_F(IcingDynamicTrieTest, Compact) {
-  IcingFilesystem filesystem;
-  IcingDynamicTrie trie(trie_files_prefix_, IcingDynamicTrie::RuntimeOptions(),
-                        &filesystem);
-  ASSERT_TRUE(trie.CreateIfNotExist(IcingDynamicTrie::Options()));
-  ASSERT_TRUE(trie.Init());
-
-  for (uint32_t i = 0; i < kNumKeys; i++) {
-    ASSERT_TRUE(trie.Insert(kKeys[i].data(), &i));
-
-    uint32_t val;
-    bool found = trie.Find(kKeys[i].data(), &val);
-    EXPECT_TRUE(found) << kKeys[i];
-    if (found) EXPECT_EQ(i, val) << kKeys[i] << " " << val;
-  }
-
-  EXPECT_EQ(trie.size(), kNumKeys);
-
-  StatsDump(trie);
-  PrintTrie(trie);
-
-  IcingDynamicTrie trie2(trie_files_prefix_ + "-2",
-                         IcingDynamicTrie::RuntimeOptions(), &filesystem);
-  ASSERT_TRUE(trie2.CreateIfNotExist(IcingDynamicTrie::Options()));
-  ASSERT_TRUE(trie2.Init());
-
-  std::unordered_map<uint32_t, uint32_t> old_to_new_tvi;
-  trie.Compact(NewValueMap(), &trie2, &old_to_new_tvi);
-  for (uint32_t i = 0; i < kNumKeys; i++) {
-    uint32_t val;
-    bool found = trie2.Find(kKeys[i].data(), &val);
-    EXPECT_TRUE(found) << kKeys[i];
-    EXPECT_TRUE(old_to_new_tvi.find(val) != old_to_new_tvi.end());
-  }
-}
-
 TEST_F(IcingDynamicTrieTest, DeletionShouldWorkWhenRootIsLeaf) {
   IcingFilesystem filesystem;
   IcingDynamicTrie trie(trie_files_prefix_, IcingDynamicTrie::RuntimeOptions(),
@@ -875,7 +829,7 @@ TEST_F(IcingDynamicTrieTest, DeletionShouldWorkWhenRootIsLeaf) {
 
   // Inserts a key, the root is a leaf.
   uint32_t value = 1;
-  ASSERT_TRUE(trie.Insert("foo", &value));
+  ASSERT_THAT(trie.Insert("foo", &value), IsOk());
   ASSERT_TRUE(trie.Find("foo", &value));
 
   // Deletes the key.
@@ -899,8 +853,8 @@ TEST_F(IcingDynamicTrieTest, DeletionShouldWorkWhenLastCharIsLeaf) {
   //        / \
   //     null  r
   uint32_t value = 1;
-  ASSERT_TRUE(trie.Insert("bar", &value));
-  ASSERT_TRUE(trie.Insert("ba", &value));
+  ASSERT_THAT(trie.Insert("bar", &value), IsOk());
+  ASSERT_THAT(trie.Insert("ba", &value), IsOk());
   ASSERT_TRUE(trie.Find("bar", &value));
   ASSERT_TRUE(trie.Find("ba", &value));
 
@@ -926,8 +880,8 @@ TEST_F(IcingDynamicTrieTest, DeletionShouldWorkWithTerminationNode) {
   //        / \
   //     null  r
   uint32_t value = 1;
-  ASSERT_TRUE(trie.Insert("bar", &value));
-  ASSERT_TRUE(trie.Insert("ba", &value));
+  ASSERT_THAT(trie.Insert("bar", &value), IsOk());
+  ASSERT_THAT(trie.Insert("ba", &value), IsOk());
   ASSERT_TRUE(trie.Find("bar", &value));
   ASSERT_TRUE(trie.Find("ba", &value));
 
@@ -951,10 +905,10 @@ TEST_F(IcingDynamicTrieTest, DeletionShouldWorkWithMultipleNexts) {
   //      / | | \
   //     a  b c  d
   uint32_t value = 1;
-  ASSERT_TRUE(trie.Insert("ba", &value));
-  ASSERT_TRUE(trie.Insert("bb", &value));
-  ASSERT_TRUE(trie.Insert("bc", &value));
-  ASSERT_TRUE(trie.Insert("bd", &value));
+  ASSERT_THAT(trie.Insert("ba", &value), IsOk());
+  ASSERT_THAT(trie.Insert("bb", &value), IsOk());
+  ASSERT_THAT(trie.Insert("bc", &value), IsOk());
+  ASSERT_THAT(trie.Insert("bd", &value), IsOk());
   ASSERT_TRUE(trie.Find("ba", &value));
   ASSERT_TRUE(trie.Find("bb", &value));
   ASSERT_TRUE(trie.Find("bc", &value));
@@ -990,9 +944,9 @@ TEST_F(IcingDynamicTrieTest, DeletionShouldWorkWithMultipleTrieBranches) {
   //     |   |
   //     r   e
   uint32_t value = 1;
-  ASSERT_TRUE(trie.Insert("batter", &value));
-  ASSERT_TRUE(trie.Insert("battle", &value));
-  ASSERT_TRUE(trie.Insert("bar", &value));
+  ASSERT_THAT(trie.Insert("batter", &value), IsOk());
+  ASSERT_THAT(trie.Insert("battle", &value), IsOk());
+  ASSERT_THAT(trie.Insert("bar", &value), IsOk());
   ASSERT_TRUE(trie.Find("batter", &value));
   ASSERT_TRUE(trie.Find("battle", &value));
   ASSERT_TRUE(trie.Find("bar", &value));
@@ -1013,17 +967,17 @@ TEST_F(IcingDynamicTrieTest, InsertionShouldWorkAfterDeletion) {
 
   // Inserts some keys.
   uint32_t value = 1;
-  ASSERT_TRUE(trie.Insert("bar", &value));
-  ASSERT_TRUE(trie.Insert("bed", &value));
-  ASSERT_TRUE(trie.Insert("foo", &value));
+  ASSERT_THAT(trie.Insert("bar", &value), IsOk());
+  ASSERT_THAT(trie.Insert("bed", &value), IsOk());
+  ASSERT_THAT(trie.Insert("foo", &value), IsOk());
 
   // Deletes a key
   ASSERT_TRUE(trie.Delete("bed"));
   ASSERT_FALSE(trie.Find("bed", &value));
 
   // Inserts after deletion
-  EXPECT_TRUE(trie.Insert("bed", &value));
-  EXPECT_TRUE(trie.Insert("bedroom", &value));
+  ASSERT_THAT(trie.Insert("bed", &value), IsOk());
+  ASSERT_THAT(trie.Insert("bedroom", &value), IsOk());
   EXPECT_TRUE(trie.Find("bed", &value));
   EXPECT_TRUE(trie.Find("bedroom", &value));
 }
@@ -1037,9 +991,9 @@ TEST_F(IcingDynamicTrieTest, IteratorShouldWorkAfterDeletion) {
 
   // Inserts some keys.
   uint32_t value = 1;
-  ASSERT_TRUE(trie.Insert("bar", &value));
-  ASSERT_TRUE(trie.Insert("bed", &value));
-  ASSERT_TRUE(trie.Insert("foo", &value));
+  ASSERT_THAT(trie.Insert("bar", &value), IsOk());
+  ASSERT_THAT(trie.Insert("bed", &value), IsOk());
+  ASSERT_THAT(trie.Insert("foo", &value), IsOk());
 
   // Deletes a key
   ASSERT_TRUE(trie.Delete("bed"));
@@ -1070,8 +1024,8 @@ TEST_F(IcingDynamicTrieTest, DeletingNonExistingKeyShouldReturnTrue) {
 
   // Inserts some keys.
   uint32_t value = 1;
-  ASSERT_TRUE(trie.Insert("bar", &value));
-  ASSERT_TRUE(trie.Insert("bed", &value));
+  ASSERT_THAT(trie.Insert("bar", &value), IsOk());
+  ASSERT_THAT(trie.Insert("bed", &value), IsOk());
 
   // "ba" and bedroom are not keys in the trie.
   EXPECT_TRUE(trie.Delete("ba"));
@@ -1091,10 +1045,10 @@ TEST_F(IcingDynamicTrieTest, DeletionResortsFullNextArray) {
 
   uint32_t value = 1;
   // 'f' -> [ 'a', 'j', 'o', 'u' ]
-  ASSERT_TRUE(trie.Insert("foul", &value));
-  ASSERT_TRUE(trie.Insert("far", &value));
-  ASSERT_TRUE(trie.Insert("fudge", &value));
-  ASSERT_TRUE(trie.Insert("fjord", &value));
+  ASSERT_THAT(trie.Insert("foul", &value), IsOk());
+  ASSERT_THAT(trie.Insert("far", &value), IsOk());
+  ASSERT_THAT(trie.Insert("fudge", &value), IsOk());
+  ASSERT_THAT(trie.Insert("fjord", &value), IsOk());
 
   // Delete the third child
   EXPECT_TRUE(trie.Delete("foul"));
@@ -1116,9 +1070,9 @@ TEST_F(IcingDynamicTrieTest, DeletionResortsPartiallyFilledNextArray) {
 
   uint32_t value = 1;
   // 'f' -> [ 'a', 'o', 'u', 0xFF ]
-  ASSERT_TRUE(trie.Insert("foul", &value));
-  ASSERT_TRUE(trie.Insert("far", &value));
-  ASSERT_TRUE(trie.Insert("fudge", &value));
+  ASSERT_THAT(trie.Insert("foul", &value), IsOk());
+  ASSERT_THAT(trie.Insert("far", &value), IsOk());
+  ASSERT_THAT(trie.Insert("fudge", &value), IsOk());
 
   // Delete the second child
   EXPECT_TRUE(trie.Delete("foul"));
@@ -1145,7 +1099,7 @@ TEST_F(IcingDynamicTrieTest, DeletionLoadTest) {
   // Randomly generate 2048 terms.
   for (int i = 0; i < 2048; ++i) {
     terms.push_back(RandomString("abcdefg", 5, &random));
-    ASSERT_TRUE(trie.Insert(terms.back().c_str(), &value));
+    ASSERT_THAT(trie.Insert(terms.back().c_str(), &value), IsOk());
   }
 
   // Randomly delete 1024 terms.
@@ -1167,7 +1121,7 @@ TEST_F(IcingDynamicTrieTest, DeletionLoadTest) {
   // Check that we can still insert terms after delete.
   for (int i = 0; i < 2048; ++i) {
     std::string term = RandomString("abcdefg", 5, &random);
-    ASSERT_TRUE(trie.Insert(term.c_str(), &value));
+    ASSERT_THAT(trie.Insert(term.c_str(), &value), IsOk());
     exp_remaining.insert(term);
   }
   remaining.clear();
@@ -1207,7 +1161,7 @@ TEST_F(IcingDynamicTrieTest, TrieShouldRespectLimits) {
     // Inserts all the test words before the last one.
     uint32_t value = 0;
     for (size_t i = 0; i < kCommonEnglishWordArrayLen - 1; ++i) {
-      ASSERT_TRUE(trie.Insert(kCommonEnglishWords[i].data(), &value));
+      ASSERT_THAT(trie.Insert(kCommonEnglishWords[i].data(), &value), IsOk());
     }
 
     IcingDynamicTrieHeader header;
@@ -1245,12 +1199,14 @@ TEST_F(IcingDynamicTrieTest, TrieShouldRespectLimits) {
     // Inserts all the test words before the last one.
     uint32_t value = 0;
     for (size_t i = 0; i < kCommonEnglishWordArrayLen - 1; ++i) {
-      ASSERT_TRUE(trie.Insert(kCommonEnglishWords[i].data(), &value));
+      ASSERT_THAT(trie.Insert(kCommonEnglishWords[i].data(), &value), IsOk());
     }
 
     // Fails to insert the last word because no enough nodes left.
-    EXPECT_FALSE(trie.Insert(
-        kCommonEnglishWords[kCommonEnglishWordArrayLen - 1].data(), &value));
+    EXPECT_THAT(
+        trie.Insert(kCommonEnglishWords[kCommonEnglishWordArrayLen - 1].data(),
+                    &value),
+        StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
   }
 
   // Test a trie with just enough number of nexts.
@@ -1266,12 +1222,14 @@ TEST_F(IcingDynamicTrieTest, TrieShouldRespectLimits) {
     // Inserts all the test words before the last one.
     uint32_t value = 0;
     for (size_t i = 0; i < kCommonEnglishWordArrayLen - 1; ++i) {
-      ASSERT_TRUE(trie.Insert(kCommonEnglishWords[i].data(), &value));
+      ASSERT_THAT(trie.Insert(kCommonEnglishWords[i].data(), &value), IsOk());
     }
 
     // Fails to insert the last word because no enough nexts left.
-    EXPECT_FALSE(trie.Insert(
-        kCommonEnglishWords[kCommonEnglishWordArrayLen - 1].data(), &value));
+    EXPECT_THAT(
+        trie.Insert(kCommonEnglishWords[kCommonEnglishWordArrayLen - 1].data(),
+                    &value),
+        StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
   }
 
   // Test a trie with just enough suffixes size.
@@ -1287,12 +1245,14 @@ TEST_F(IcingDynamicTrieTest, TrieShouldRespectLimits) {
     // Inserts all the test words before the last one.
     uint32_t value = 0;
     for (size_t i = 0; i < kCommonEnglishWordArrayLen - 1; ++i) {
-      ASSERT_TRUE(trie.Insert(kCommonEnglishWords[i].data(), &value));
+      ASSERT_THAT(trie.Insert(kCommonEnglishWords[i].data(), &value), IsOk());
     }
 
     // Fails to insert the last word because no enough space for more suffixes.
-    EXPECT_FALSE(trie.Insert(
-        kCommonEnglishWords[kCommonEnglishWordArrayLen - 1].data(), &value));
+    EXPECT_THAT(
+        trie.Insert(kCommonEnglishWords[kCommonEnglishWordArrayLen - 1].data(),
+                    &value),
+        StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
   }
 }
 
@@ -1358,27 +1318,27 @@ TEST_F(IcingDynamicTrieTest, IsBranchingTermShouldWorkForExistingTerms) {
 
   uint32_t value = 1;
 
-  ASSERT_TRUE(trie.Insert("", &value));
+  ASSERT_THAT(trie.Insert("", &value), IsOk());
   EXPECT_FALSE(trie.IsBranchingTerm(""));
 
-  ASSERT_TRUE(trie.Insert("ab", &value));
+  ASSERT_THAT(trie.Insert("ab", &value), IsOk());
   EXPECT_FALSE(trie.IsBranchingTerm(""));
   EXPECT_FALSE(trie.IsBranchingTerm("ab"));
 
-  ASSERT_TRUE(trie.Insert("ac", &value));
+  ASSERT_THAT(trie.Insert("ac", &value), IsOk());
   // "" is a prefix of "ab" and "ac", but it is not a branching term.
   EXPECT_FALSE(trie.IsBranchingTerm(""));
   EXPECT_FALSE(trie.IsBranchingTerm("ab"));
   EXPECT_FALSE(trie.IsBranchingTerm("ac"));
 
-  ASSERT_TRUE(trie.Insert("ba", &value));
+  ASSERT_THAT(trie.Insert("ba", &value), IsOk());
   // "" now branches to "ba"
   EXPECT_TRUE(trie.IsBranchingTerm(""));
   EXPECT_FALSE(trie.IsBranchingTerm("ab"));
   EXPECT_FALSE(trie.IsBranchingTerm("ac"));
   EXPECT_FALSE(trie.IsBranchingTerm("ba"));
 
-  ASSERT_TRUE(trie.Insert("a", &value));
+  ASSERT_THAT(trie.Insert("a", &value), IsOk());
   EXPECT_TRUE(trie.IsBranchingTerm(""));
   // "a" branches to "ab" and "ac"
   EXPECT_TRUE(trie.IsBranchingTerm("a"));
@@ -1386,8 +1346,8 @@ TEST_F(IcingDynamicTrieTest, IsBranchingTermShouldWorkForExistingTerms) {
   EXPECT_FALSE(trie.IsBranchingTerm("ac"));
   EXPECT_FALSE(trie.IsBranchingTerm("ba"));
 
-  ASSERT_TRUE(trie.Insert("abc", &value));
-  ASSERT_TRUE(trie.Insert("acd", &value));
+  ASSERT_THAT(trie.Insert("abc", &value), IsOk());
+  ASSERT_THAT(trie.Insert("acd", &value), IsOk());
   EXPECT_TRUE(trie.IsBranchingTerm(""));
   EXPECT_TRUE(trie.IsBranchingTerm("a"));
   // "ab" is a prefix of "abc", but it is not a branching term.
@@ -1398,7 +1358,7 @@ TEST_F(IcingDynamicTrieTest, IsBranchingTermShouldWorkForExistingTerms) {
   EXPECT_FALSE(trie.IsBranchingTerm("abc"));
   EXPECT_FALSE(trie.IsBranchingTerm("acd"));
 
-  ASSERT_TRUE(trie.Insert("abcd", &value));
+  ASSERT_THAT(trie.Insert("abcd", &value), IsOk());
   EXPECT_TRUE(trie.IsBranchingTerm(""));
   EXPECT_TRUE(trie.IsBranchingTerm("a"));
   // "ab" is a prefix of "abc" and "abcd", but it is not a branching term.
@@ -1410,7 +1370,7 @@ TEST_F(IcingDynamicTrieTest, IsBranchingTermShouldWorkForExistingTerms) {
   EXPECT_FALSE(trie.IsBranchingTerm("acd"));
   EXPECT_FALSE(trie.IsBranchingTerm("abcd"));
 
-  ASSERT_TRUE(trie.Insert("abd", &value));
+  ASSERT_THAT(trie.Insert("abd", &value), IsOk());
   EXPECT_TRUE(trie.IsBranchingTerm(""));
   EXPECT_TRUE(trie.IsBranchingTerm("a"));
   // "ab" branches to "abc" and "abd"
@@ -1437,46 +1397,46 @@ TEST_F(IcingDynamicTrieTest, IsBranchingTermShouldWorkForNonExistingTerms) {
   EXPECT_FALSE(trie.IsBranchingTerm("ab"));
   EXPECT_FALSE(trie.IsBranchingTerm("abc"));
 
-  ASSERT_TRUE(trie.Insert("aa", &value));
+  ASSERT_THAT(trie.Insert("aa", &value), IsOk());
   EXPECT_FALSE(trie.IsBranchingTerm(""));
   EXPECT_FALSE(trie.IsBranchingTerm("a"));
   EXPECT_FALSE(trie.IsBranchingTerm("ab"));
   EXPECT_FALSE(trie.IsBranchingTerm("abc"));
 
-  ASSERT_TRUE(trie.Insert("ac", &value));
+  ASSERT_THAT(trie.Insert("ac", &value), IsOk());
   EXPECT_FALSE(trie.IsBranchingTerm(""));
   // "a" does not exist in the trie, but now it branches to "aa" and "ac".
   EXPECT_TRUE(trie.IsBranchingTerm("a"));
   EXPECT_FALSE(trie.IsBranchingTerm("ab"));
   EXPECT_FALSE(trie.IsBranchingTerm("abc"));
 
-  ASSERT_TRUE(trie.Insert("ad", &value));
+  ASSERT_THAT(trie.Insert("ad", &value), IsOk());
   EXPECT_FALSE(trie.IsBranchingTerm(""));
   EXPECT_TRUE(trie.IsBranchingTerm("a"));
   EXPECT_FALSE(trie.IsBranchingTerm("ab"));
   EXPECT_FALSE(trie.IsBranchingTerm("abc"));
 
-  ASSERT_TRUE(trie.Insert("abcd", &value));
+  ASSERT_THAT(trie.Insert("abcd", &value), IsOk());
   EXPECT_FALSE(trie.IsBranchingTerm(""));
   EXPECT_TRUE(trie.IsBranchingTerm("a"));
   EXPECT_FALSE(trie.IsBranchingTerm("ab"));
   EXPECT_FALSE(trie.IsBranchingTerm("abc"));
 
-  ASSERT_TRUE(trie.Insert("abd", &value));
+  ASSERT_THAT(trie.Insert("abd", &value), IsOk());
   EXPECT_FALSE(trie.IsBranchingTerm(""));
   EXPECT_TRUE(trie.IsBranchingTerm("a"));
   // "ab" does not exist in the trie, but now it branches to "abcd" and "abd".
   EXPECT_TRUE(trie.IsBranchingTerm("ab"));
   EXPECT_FALSE(trie.IsBranchingTerm("abc"));
 
-  ASSERT_TRUE(trie.Insert("abce", &value));
+  ASSERT_THAT(trie.Insert("abce", &value), IsOk());
   EXPECT_FALSE(trie.IsBranchingTerm(""));
   EXPECT_TRUE(trie.IsBranchingTerm("a"));
   EXPECT_TRUE(trie.IsBranchingTerm("ab"));
   // "abc" does not exist in the trie, but now it branches to "abcd" and "abce".
   EXPECT_TRUE(trie.IsBranchingTerm("abc"));
 
-  ASSERT_TRUE(trie.Insert("abc_suffix", &value));
+  ASSERT_THAT(trie.Insert("abc_suffix", &value), IsOk());
   EXPECT_FALSE(trie.IsBranchingTerm(""));
   EXPECT_TRUE(trie.IsBranchingTerm("a"));
   EXPECT_TRUE(trie.IsBranchingTerm("ab"));