Add/Get n-gram probability entry in languageModelDictContent

Bug: 14425059
Change-Id: I7926c3812f89b9a71fe1873a5bc32f793f91b640
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 07e1051..5dc91ba 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
@@ -32,11 +32,11 @@
 
 ProbabilityEntry LanguageModelDictContent::getNgramProbabilityEntry(
         const WordIdArrayView prevWordIds, const int wordId) const {
-    if (!prevWordIds.empty()) {
-        // TODO: Read n-gram entry.
+    const int bitmapEntryIndex = getBitmapEntryIndex(prevWordIds);
+    if (bitmapEntryIndex == TrieMap::INVALID_INDEX) {
         return ProbabilityEntry();
     }
-    const TrieMap::Result result = mTrieMap.getRoot(wordId);
+    const TrieMap::Result result = mTrieMap.get(wordId, bitmapEntryIndex);
     if (!result.mIsValid) {
         // Not found.
         return ProbabilityEntry();
@@ -46,14 +46,13 @@
 
 bool LanguageModelDictContent::setNgramProbabilityEntry(const WordIdArrayView prevWordIds,
         const int terminalId, const ProbabilityEntry *const probabilityEntry) {
-    if (!prevWordIds.empty()) {
-        // TODO: Add n-gram entry.
+    const int bitmapEntryIndex = getBitmapEntryIndex(prevWordIds);
+    if (bitmapEntryIndex == TrieMap::INVALID_INDEX) {
         return false;
     }
-    return mTrieMap.putRoot(terminalId, probabilityEntry->encode(mHasHistoricalInfo));
+    return mTrieMap.put(terminalId, probabilityEntry->encode(mHasHistoricalInfo), bitmapEntryIndex);
 }
 
-
 bool LanguageModelDictContent::runGCInner(
         const TerminalPositionLookupTable::TerminalIdMap *const terminalIdMap,
         const TrieMap::TrieMapRange trieMapRange,
@@ -81,4 +80,16 @@
     return true;
 }
 
+int LanguageModelDictContent::getBitmapEntryIndex(const WordIdArrayView prevWordIds) const {
+    int bitmapEntryIndex = mTrieMap.getRootBitmapEntryIndex();
+    for (const int wordId : prevWordIds) {
+        const TrieMap::Result result = mTrieMap.get(wordId, bitmapEntryIndex);
+        if (!result.mIsValid) {
+            return TrieMap::INVALID_INDEX;
+        }
+        bitmapEntryIndex = result.mNextLevelBitmapEntryIndex;
+    }
+    return bitmapEntryIndex;
+}
+
 } // namespace latinime