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
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content.h b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content.h
index f181dfe..18f2e01 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content.h
@@ -76,6 +76,8 @@
     bool runGCInner(const TerminalPositionLookupTable::TerminalIdMap *const terminalIdMap,
             const TrieMap::TrieMapRange trieMapRange, const int nextLevelBitmapEntryIndex,
             int *const outNgramCount);
+
+    int getBitmapEntryIndex(const WordIdArrayView prevWordIds) const;
 };
 } // namespace latinime
 #endif /* LATINIME_LANGUAGE_MODEL_DICT_CONTENT_H */
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 a294ab8..3e5c401 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.h
@@ -169,6 +169,10 @@
         return mBuffer.isNearSizeLimit();
     }
 
+    int getRootBitmapEntryIndex() const {
+        return ROOT_BITMAP_ENTRY_INDEX;
+    }
+
     // Returns bitmapEntryIndex. Create the next level map if it doesn't exist.
     int getNextLevelBitmapEntryIndex(const int key) {
         return getNextLevelBitmapEntryIndex(key, ROOT_BITMAP_ENTRY_INDEX);
diff --git a/native/jni/src/utils/int_array_view.h b/native/jni/src/utils/int_array_view.h
index 4bc2487..2418b42 100644
--- a/native/jni/src/utils/int_array_view.h
+++ b/native/jni/src/utils/int_array_view.h
@@ -73,6 +73,14 @@
         return mPtr;
     }
 
+    AK_FORCE_INLINE const int *begin() const {
+        return mPtr;
+    }
+
+    AK_FORCE_INLINE const int *end() const {
+        return mPtr + mSize;
+    }
+
  private:
     DISALLOW_ASSIGNMENT_OPERATOR(IntArrayView);
 
diff --git a/native/jni/tests/utils/int_array_view_test.cpp b/native/jni/tests/utils/int_array_view_test.cpp
index 9aa8cdc..f799dd8 100644
--- a/native/jni/tests/utils/int_array_view_test.cpp
+++ b/native/jni/tests/utils/int_array_view_test.cpp
@@ -24,15 +24,24 @@
 namespace {
 
 TEST(MemoryViewTest, TestAccess) {
-    static const int DATA_SIZE = 10000;
-
     std::vector<int> intVector = {3, 2, 1, 0, -1, -2};
     IntArrayView intArrayView(intVector);
     EXPECT_EQ(intVector.size(), intArrayView.size());
-    for (int i = 0; i < DATA_SIZE; ++i) {
+    for (int i = 0; i < static_cast<int>(intVector.size()); ++i) {
         EXPECT_EQ(intVector[i], intArrayView[i]);
     }
 }
 
+TEST(MemoryViewTest, TestIteration) {
+    std::vector<int> intVector = {3, 2, 1, 0, -1, -2};
+    IntArrayView intArrayView(intVector);
+    std::set<int> intSet(intVector.begin(), intVector.end());
+    for (const int i : intArrayView) {
+        EXPECT_TRUE(intSet.count(i) > 0);
+        intSet.erase(i);
+    }
+    EXPECT_TRUE(intSet.empty());
+}
+
 }  // namespace
 }  // namespace latinime