Add TrieMap::remove().

Bug: 14425059
Change-Id: Ifb9a873eab79efdd6bee8e0da73c0cf37727350e

3 files changed, 103 insertions, 4 deletions

diff --git a/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.cpp b/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.cpp
index 407b8efd0..e630aba9a 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.cpp
@@ -26,6 +26,7 @@ const int TrieMap::FIELD1_SIZE = 3;
 const int TrieMap::ENTRY_SIZE = FIELD0_SIZE + FIELD1_SIZE;
 const uint32_t TrieMap::VALUE_FLAG = 0x400000;
 const uint32_t TrieMap::VALUE_MASK = 0x3FFFFF;
+const uint32_t TrieMap::INVALID_VALUE_IN_KEY_VALUE_ENTRY = VALUE_MASK;
 const uint32_t TrieMap::TERMINAL_LINK_FLAG = 0x800000;
 const uint32_t TrieMap::TERMINAL_LINK_MASK = 0x7FFFFF;
 const int TrieMap::NUM_OF_BITS_USED_FOR_ONE_LEVEL = 5;
@@ -34,6 +35,7 @@ const int TrieMap::MAX_NUM_OF_ENTRIES_IN_ONE_LEVEL = 1 << NUM_OF_BITS_USED_FOR_O
 const int TrieMap::ROOT_BITMAP_ENTRY_INDEX = 0;
 const int TrieMap::ROOT_BITMAP_ENTRY_POS = MAX_NUM_OF_ENTRIES_IN_ONE_LEVEL * FIELD0_SIZE;
 const TrieMap::Entry TrieMap::EMPTY_BITMAP_ENTRY = TrieMap::Entry(0, 0);
+const int TrieMap::TERMINAL_LINKED_ENTRY_COUNT = 2; // Value entry and bitmap entry.
 const uint64_t TrieMap::MAX_VALUE =
         (static_cast<uint64_t>(1) << ((FIELD0_SIZE + FIELD1_SIZE) * CHAR_BIT)) - 1;
 const int TrieMap::MAX_BUFFER_SIZE = TERMINAL_LINK_MASK * ENTRY_SIZE;
@@ -76,7 +78,7 @@ int TrieMap::getNextLevelBitmapEntryIndex(const int key, const int bitmapEntryIn
         return terminalEntry.getValueEntryIndex() + 1;
     }
     // Create a value entry and a bitmap entry.
-    const int valueEntryIndex = allocateTable(2 /* entryCount */);
+    const int valueEntryIndex = allocateTable(TERMINAL_LINKED_ENTRY_COUNT);
     if (!writeEntry(Entry(0, terminalEntry.getValue()), valueEntryIndex)) {
         return INVALID_INDEX;
     }
@@ -108,6 +110,31 @@ bool TrieMap::save(FILE *const file) const {
     return DictFileWritingUtils::writeBufferToFileTail(file, &mBuffer);
 }
 
+bool TrieMap::remove(const int key, const int bitmapEntryIndex) {
+    const Entry bitmapEntry = readEntry(bitmapEntryIndex);
+    const uint32_t unsignedKey = static_cast<uint32_t>(key);
+    const int terminalEntryIndex = getTerminalEntryIndex(
+            unsignedKey, getBitShuffledKey(unsignedKey), bitmapEntry, 0 /* level */);
+    if (terminalEntryIndex == INVALID_INDEX) {
+        // Not found.
+        return false;
+    }
+    const Entry terminalEntry = readEntry(terminalEntryIndex);
+    if (!writeField1(VALUE_FLAG ^ INVALID_VALUE_IN_KEY_VALUE_ENTRY , terminalEntryIndex)) {
+        return false;
+    }
+    if (terminalEntry.hasTerminalLink()) {
+        const Entry nextLevelBitmapEntry = readEntry(terminalEntry.getValueEntryIndex() + 1);
+        if (!freeTable(terminalEntry.getValueEntryIndex(), TERMINAL_LINKED_ENTRY_COUNT)) {
+            return false;
+        }
+        if (!removeInner(nextLevelBitmapEntry)){
+            return false;
+        }
+    }
+    return true;
+}
+
 /**
  * Iterate next entry in a certain level.
  *
@@ -129,7 +156,7 @@ const TrieMap::Result TrieMap::iterateNext(std::vector<TableIterationState> *con
             if (entry.isBitmapEntry()) {
                 // Move to child.
                 iterationState->emplace_back(popCount(entry.getBitmap()), entry.getTableIndex());
-            } else {
+            } else if (entry.isValidTerminalEntry()) {
                 if (outKey) {
                     *outKey = entry.getKey();
                 }
@@ -162,12 +189,12 @@ uint32_t TrieMap::getBitShuffledKey(const uint32_t key) const {
 }
 
 bool TrieMap::writeValue(const uint64_t value, const int terminalEntryIndex) {
-    if (value <= VALUE_MASK) {
+    if (value < VALUE_MASK) {
         // Write value into the terminal entry.
         return writeField1(value | VALUE_FLAG, terminalEntryIndex);
     }
     // Create value entry and write value.
-    const int valueEntryIndex = allocateTable(2 /* entryCount */);
+    const int valueEntryIndex = allocateTable(TERMINAL_LINKED_ENTRY_COUNT);
     if (!writeEntry(Entry(value >> (FIELD1_SIZE * CHAR_BIT), value), valueEntryIndex)) {
         return false;
     }
@@ -227,6 +254,9 @@ int TrieMap::getTerminalEntryIndex(const uint32_t key, const uint32_t hashedKey,
         // Move to the next level.
         return getTerminalEntryIndex(key, hashedKey, entry, level + 1);
     }
+    if (!entry.isValidTerminalEntry()) {
+        return INVALID_INDEX;
+    }
     if (entry.getKey() == key) {
         // Terminal entry is found.
         return entryIndex;
@@ -287,6 +317,10 @@ bool TrieMap::putInternal(const uint32_t key, const uint64_t value, const uint32
         // Bitmap entry is found. Go to the next level.
         return putInternal(key, value, hashedKey, entryIndex, entry, level + 1);
     }
+    if (!entry.isValidTerminalEntry()) {
+        // Overwrite invalid terminal entry.
+        return writeTerminalEntry(key, value, entryIndex);
+    }
     if (entry.getKey() == key) {
         // Terminal entry for the key is found. Update the value.
         return updateValue(entry, value, entryIndex);
@@ -384,4 +418,33 @@ bool TrieMap::addNewEntryByExpandingTable(const uint32_t key, const uint64_t val
     return true;
 }
 
+bool TrieMap::removeInner(const Entry &bitmapEntry) {
+    const int tableSize = popCount(bitmapEntry.getBitmap());
+    for (int i = 0; i < tableSize; ++i) {
+        const int entryIndex = bitmapEntry.getTableIndex() + i;
+        const Entry entry = readEntry(entryIndex);
+        if (entry.isBitmapEntry()) {
+            // Delete next bitmap entry recursively.
+            if (!removeInner(entry)) {
+                return false;
+            }
+        } else {
+            // Invalidate terminal entry just in case.
+            if (!writeField1(VALUE_FLAG ^ INVALID_VALUE_IN_KEY_VALUE_ENTRY , entryIndex)) {
+                return false;
+            }
+            if (entry.hasTerminalLink()) {
+                const Entry nextLevelBitmapEntry = readEntry(entry.getValueEntryIndex() + 1);
+                if (!freeTable(entry.getValueEntryIndex(), TERMINAL_LINKED_ENTRY_COUNT)) {
+                    return false;
+                }
+                if (!removeInner(nextLevelBitmapEntry)) {
+                    return false;
+                }
+            }
+        }
+    }
+    return freeTable(bitmapEntry.getTableIndex(), tableSize);
+}
+
 }  // namespace latinime
diff --git a/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.h b/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.h
index 3e5c4010c..6d91790b2 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.h
@@ -202,6 +202,8 @@ class TrieMap {
 
     bool save(FILE *const file) const;
 
+    bool remove(const int key, const int bitmapEntryIndex);
+
  private:
     DISALLOW_COPY_AND_ASSIGN(TrieMap);
 
@@ -245,6 +247,11 @@ class TrieMap {
         }
 
         // For terminal entry.
+        AK_FORCE_INLINE bool isValidTerminalEntry() const {
+            return hasTerminalLink() || ((mData1 & VALUE_MASK) != INVALID_VALUE_IN_KEY_VALUE_ENTRY);
+        }
+
+        // For terminal entry.
         AK_FORCE_INLINE uint32_t getValueEntryIndex() const {
             return mData1 & TERMINAL_LINK_MASK;
         }
@@ -272,6 +279,7 @@ class TrieMap {
     static const int ENTRY_SIZE;
     static const uint32_t VALUE_FLAG;
     static const uint32_t VALUE_MASK;
+    static const uint32_t INVALID_VALUE_IN_KEY_VALUE_ENTRY;
     static const uint32_t TERMINAL_LINK_FLAG;
     static const uint32_t TERMINAL_LINK_MASK;
     static const int NUM_OF_BITS_USED_FOR_ONE_LEVEL;
@@ -280,6 +288,7 @@ class TrieMap {
     static const int ROOT_BITMAP_ENTRY_INDEX;
     static const int ROOT_BITMAP_ENTRY_POS;
     static const Entry EMPTY_BITMAP_ENTRY;
+    static const int TERMINAL_LINKED_ENTRY_COUNT;
     static const int MAX_BUFFER_SIZE;
 
     uint32_t getBitShuffledKey(const uint32_t key) const;
@@ -378,6 +387,8 @@ class TrieMap {
     AK_FORCE_INLINE int getTailEntryIndex() const {
         return (mBuffer.getTailPosition() - ROOT_BITMAP_ENTRY_POS) / ENTRY_SIZE;
     }
+
+    bool removeInner(const Entry &bitmapEntry);
 };
 
 } // namespace latinime
diff --git a/native/jni/tests/suggest/policyimpl/dictionary/utils/trie_map_test.cpp b/native/jni/tests/suggest/policyimpl/dictionary/utils/trie_map_test.cpp
index df778b6cf..8c8e8838a 100644
--- a/native/jni/tests/suggest/policyimpl/dictionary/utils/trie_map_test.cpp
+++ b/native/jni/tests/suggest/policyimpl/dictionary/utils/trie_map_test.cpp
@@ -47,6 +47,31 @@ TEST(TrieMapTest, TestSetAndGet) {
     EXPECT_EQ(0xFFFFFFFFFull, trieMap.getRoot(0).mValue);
 }
 
+TEST(TrieMapTest, TestRemove) {
+    TrieMap trieMap;
+    trieMap.putRoot(10, 10);
+    EXPECT_EQ(10ull, trieMap.getRoot(10).mValue);
+    EXPECT_TRUE(trieMap.remove(10, trieMap.getRootBitmapEntryIndex()));
+    EXPECT_FALSE(trieMap.getRoot(10).mIsValid);
+    for (const auto &element : trieMap.getEntriesInRootLevel()) {
+        EXPECT_TRUE(false);
+    }
+    EXPECT_TRUE(trieMap.putRoot(10, 0x3FFFFF));
+    EXPECT_FALSE(trieMap.remove(11, trieMap.getRootBitmapEntryIndex()))
+            << "Should fail if the key does not exist.";
+    EXPECT_EQ(0x3FFFFFull, trieMap.getRoot(10).mValue);
+    trieMap.putRoot(12, 11);
+    const int nextLevel = trieMap.getNextLevelBitmapEntryIndex(10);
+    trieMap.put(10, 10, nextLevel);
+    EXPECT_EQ(0x3FFFFFull, trieMap.getRoot(10).mValue);
+    EXPECT_EQ(10ull, trieMap.get(10, nextLevel).mValue);
+    EXPECT_TRUE(trieMap.remove(10, trieMap.getRootBitmapEntryIndex()));
+    const TrieMap::Result result = trieMap.getRoot(10);
+    EXPECT_FALSE(result.mIsValid);
+    EXPECT_EQ(TrieMap::INVALID_INDEX, result.mNextLevelBitmapEntryIndex);
+    EXPECT_EQ(11ull, trieMap.getRoot(12).mValue);
+}
+
 TEST(TrieMapTest, TestSetAndGetLarge) {
     static const int ELEMENT_COUNT = 200000;
     TrieMap trieMap;