1 files changed, 14 insertions, 8 deletions
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 7449cd02b..a96719533 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
@@ -31,6 +31,7 @@
 #include "suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_reader.h"
 #include "suggest/policyimpl/dictionary/utils/forgetting_curve_utils.h"
 #include "suggest/policyimpl/dictionary/utils/probability_utils.h"
+#include "utils/ngram_utils.h"
 
 namespace latinime {
 
@@ -215,7 +216,7 @@ bool Ver4PatriciaTriePolicy::addUnigramEntry(const CodePointArrayView wordCodePo
     if (mUpdatingHelper.addUnigramWord(&readingHelper, codePointArrayView, unigramProperty,
             &addedNewUnigram)) {
         if (addedNewUnigram && !unigramProperty->representsBeginningOfSentence()) {
-            mEntryCounters.incrementUnigramCount();
+            mEntryCounters.incrementNgramCount(NgramType::Unigram);
         }
         if (unigramProperty->getShortcuts().size() > 0) {
             // Add shortcut target.
@@ -263,7 +264,7 @@ bool Ver4PatriciaTriePolicy::removeUnigramEntry(const CodePointArrayView wordCod
         return false;
     }
     if (!ptNodeParams.representsNonWordInfo()) {
-        mEntryCounters.decrementUnigramCount();
+        mEntryCounters.decrementNgramCount(NgramType::Unigram);
     }
     return true;
 }
@@ -321,7 +322,8 @@ bool Ver4PatriciaTriePolicy::addNgramEntry(const NgramProperty *const ngramPrope
     bool addedNewEntry = false;
     if (mNodeWriter.addNgramEntry(prevWordIds, wordId, ngramProperty, &addedNewEntry)) {
         if (addedNewEntry) {
-            mEntryCounters.incrementNgramCount(prevWordIds.size() + 1);
+            mEntryCounters.incrementNgramCount(
+                    NgramUtils::getNgramTypeFromWordCount(prevWordIds.size() + 1));
         }
         return true;
     } else {
@@ -359,7 +361,8 @@ bool Ver4PatriciaTriePolicy::removeNgramEntry(const NgramContext *const ngramCon
         return false;
     }
     if (mNodeWriter.removeNgramEntry(prevWordIds, wordId)) {
-        mEntryCounters.decrementNgramCount(prevWordIds.size());
+        mEntryCounters.decrementNgramCount(
+                NgramUtils::getNgramTypeFromWordCount(prevWordIds.size() + 1));
         return true;
     } else {
         return false;
@@ -477,20 +480,23 @@ void Ver4PatriciaTriePolicy::getProperty(const char *const query, const int quer
         char *const outResult, const int maxResultLength) {
     const int compareLength = queryLength + 1 /* terminator */;
     if (strncmp(query, UNIGRAM_COUNT_QUERY, compareLength) == 0) {
-        snprintf(outResult, maxResultLength, "%d", mEntryCounters.getUnigramCount());
+        snprintf(outResult, maxResultLength, "%d",
+                mEntryCounters.getNgramCount(NgramType::Unigram));
     } else if (strncmp(query, BIGRAM_COUNT_QUERY, compareLength) == 0) {
-        snprintf(outResult, maxResultLength, "%d", mEntryCounters.getBigramCount());
+        snprintf(outResult, maxResultLength, "%d", mEntryCounters.getNgramCount(NgramType::Bigram));
     } else if (strncmp(query, MAX_UNIGRAM_COUNT_QUERY, compareLength) == 0) {
         snprintf(outResult, maxResultLength, "%d",
                 mHeaderPolicy->isDecayingDict() ?
                         ForgettingCurveUtils::getEntryCountHardLimit(
-                                mHeaderPolicy->getMaxUnigramCount()) :
+                                mHeaderPolicy->getMaxNgramCounts().getNgramCount(
+                                        NgramType::Unigram)) :
                         static_cast<int>(Ver4DictConstants::MAX_DICTIONARY_SIZE));
     } else if (strncmp(query, MAX_BIGRAM_COUNT_QUERY, compareLength) == 0) {
         snprintf(outResult, maxResultLength, "%d",
                 mHeaderPolicy->isDecayingDict() ?
                         ForgettingCurveUtils::getEntryCountHardLimit(
-                                mHeaderPolicy->getMaxBigramCount()) :
+                                mHeaderPolicy->getMaxNgramCounts().getNgramCount(
+                                        NgramType::Bigram)) :
                         static_cast<int>(Ver4DictConstants::MAX_DICTIONARY_SIZE));
     }
 }