Use enum to specify ngram type.

Change-Id: Ie28768ceadcd7a2d940c57eb30be7d4c364e509f

1 files changed, 20 insertions, 64 deletions

diff --git a/native/jni/src/suggest/policyimpl/dictionary/utils/entry_counters.h b/native/jni/src/suggest/policyimpl/dictionary/utils/entry_counters.h
index 73dc42a18..7269913e8 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/utils/entry_counters.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/utils/entry_counters.h
@@ -20,6 +20,7 @@
 #include <array>
 
 #include "defines.h"
+#include "utils/ngram_utils.h"
 
 namespace latinime {
 
@@ -28,34 +29,22 @@ class EntryCounts final {
  public:
     EntryCounts() : mEntryCounts({{0, 0, 0}}) {}
 
-    EntryCounts(const int unigramCount, const int bigramCount, const int trigramCount)
-            : mEntryCounts({{unigramCount, bigramCount, trigramCount}}) {}
-
     explicit EntryCounts(const std::array<int, MAX_PREV_WORD_COUNT_FOR_N_GRAM + 1> &counters)
             : mEntryCounts(counters) {}
 
-    int getUnigramCount() const {
-        return mEntryCounts[0];
-    }
-
-    int getBigramCount() const {
-        return mEntryCounts[1];
-    }
-
-    int getTrigramCount() const {
-        return mEntryCounts[2];
+    int getNgramCount(const NgramType ngramType) const {
+        return mEntryCounts[static_cast<int>(ngramType)];
     }
 
-    int getNgramCount(const size_t n) const {
-        if (n < 1 || n > mEntryCounts.size()) {
-            return 0;
-        }
-        return mEntryCounts[n - 1];
+    const std::array<int, MAX_PREV_WORD_COUNT_FOR_N_GRAM + 1> &getCountArray() const {
+        return mEntryCounts;
     }
 
  private:
     DISALLOW_ASSIGNMENT_OPERATOR(EntryCounts);
 
+    // Counts from Unigram (0-th element) to (MAX_PREV_WORD_COUNT_FOR_N_GRAM + 1)-gram
+    // (MAX_PREV_WORD_COUNT_FOR_N_GRAM-th element)
     const std::array<int, MAX_PREV_WORD_COUNT_FOR_N_GRAM + 1> mEntryCounts;
 };
 
@@ -65,68 +54,35 @@ class MutableEntryCounters final {
         mEntryCounters.fill(0);
     }
 
-    MutableEntryCounters(const int unigramCount, const int bigramCount, const int trigramCount)
-            : mEntryCounters({{unigramCount, bigramCount, trigramCount}}) {}
+    explicit MutableEntryCounters(
+            const std::array<int, MAX_PREV_WORD_COUNT_FOR_N_GRAM + 1> &counters)
+            : mEntryCounters(counters) {}
 
     const EntryCounts getEntryCounts() const {
         return EntryCounts(mEntryCounters);
     }
 
-    int getUnigramCount() const {
-        return mEntryCounters[0];
-    }
-
-    int getBigramCount() const {
-        return mEntryCounters[1];
-    }
-
-    int getTrigramCount() const {
-        return mEntryCounters[2];
-    }
-
-    void incrementUnigramCount() {
-        ++mEntryCounters[0];
-    }
-
-    void decrementUnigramCount() {
-        ASSERT(mEntryCounters[0] != 0);
-        --mEntryCounters[0];
-    }
-
-    void incrementBigramCount() {
-        ++mEntryCounters[1];
-    }
-
-    void decrementBigramCount() {
-        ASSERT(mEntryCounters[1] != 0);
-        --mEntryCounters[1];
+    void incrementNgramCount(const NgramType ngramType) {
+        ++mEntryCounters[static_cast<int>(ngramType)];
     }
 
-    void incrementNgramCount(const size_t n) {
-        if (n < 1 || n > mEntryCounters.size()) {
-            return;
-        }
-        ++mEntryCounters[n - 1];
+    void decrementNgramCount(const NgramType ngramType) {
+        --mEntryCounters[static_cast<int>(ngramType)];
     }
 
-    void decrementNgramCount(const size_t n) {
-        if (n < 1 || n > mEntryCounters.size()) {
-            return;
-        }
-        ASSERT(mEntryCounters[n - 1] != 0);
-        --mEntryCounters[n - 1];
+    int getNgramCount(const NgramType ngramType) const {
+        return mEntryCounters[static_cast<int>(ngramType)];
     }
 
-    void setNgramCount(const size_t n, const int count) {
-        if (n < 1 || n > mEntryCounters.size()) {
-            return;
-        }
-        mEntryCounters[n - 1] = count;
+    void setNgramCount(const NgramType ngramType, const int count) {
+        mEntryCounters[static_cast<int>(ngramType)] = count;
     }
 
  private:
     DISALLOW_COPY_AND_ASSIGN(MutableEntryCounters);
 
+    // Counters from Unigram (0-th element) to (MAX_PREV_WORD_COUNT_FOR_N_GRAM + 1)-gram
+    // (MAX_PREV_WORD_COUNT_FOR_N_GRAM-th element)
     std::array<int, MAX_PREV_WORD_COUNT_FOR_N_GRAM + 1> mEntryCounters;
 };
 } // namespace latinime