Merge "Remove entry from language model dict content."

3 files changed, 18 insertions, 2 deletions

diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content.cpp
index f6721c063..bbcea2ee0 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content.cpp
@@ -45,12 +45,15 @@ ProbabilityEntry LanguageModelDictContent::getNgramProbabilityEntry(
 }
 
 bool LanguageModelDictContent::setNgramProbabilityEntry(const WordIdArrayView prevWordIds,
-        const int terminalId, const ProbabilityEntry *const probabilityEntry) {
+        const int wordId, const ProbabilityEntry *const probabilityEntry) {
+    if (wordId == Ver4DictConstants::NOT_A_TERMINAL_ID) {
+        return false;
+    }
     const int bitmapEntryIndex = createAndGetBitmapEntryIndex(prevWordIds);
     if (bitmapEntryIndex == TrieMap::INVALID_INDEX) {
         return false;
     }
-    return mTrieMap.put(terminalId, probabilityEntry->encode(mHasHistoricalInfo), bitmapEntryIndex);
+    return mTrieMap.put(wordId, probabilityEntry->encode(mHasHistoricalInfo), bitmapEntryIndex);
 }
 
 bool LanguageModelDictContent::removeNgramProbabilityEntry(const WordIdArrayView prevWordIds,
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.cpp
index 2c848cb29..62e008b94 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.cpp
@@ -255,6 +255,14 @@ bool Ver4PatriciaTrieNodeWriter::addNgramEntry(const WordIdArrayView prevWordIds
 
 bool Ver4PatriciaTrieNodeWriter::removeNgramEntry(const WordIdArrayView prevWordIds,
         const int wordId) {
+    // TODO: Support n-gram.
+    LanguageModelDictContent *const languageModelDictContent =
+            mBuffers->getMutableLanguageModelDictContent();
+    if (!languageModelDictContent->removeNgramProbabilityEntry(prevWordIds.limit(1 /* maxSize */),
+            wordId)) {
+        // TODO: Uncomment.
+        // return false;
+    }
     // TODO: Remove.
     return mBigramPolicy->removeEntry(prevWordIds[0], wordId);
 }
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp
index d8f4595ff..04e3018da 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp
@@ -263,6 +263,11 @@ bool Ver4PatriciaTriePolicy::removeUnigramEntry(const int *const word, const int
         AKLOGE("Cannot remove unigram. ptNodePos: %d", ptNodePos);
         return false;
     }
+    if (!mBuffers->getMutableLanguageModelDictContent()->removeProbabilityEntry(
+            ptNodeParams.getTerminalId())) {
+        // TODO: Uncomment.
+        // return false;
+    }
     if (!ptNodeParams.representsNonWordInfo()) {
         mUnigramCount--;
     }