Make NgramProperty have NgramContext.

Bug: 14425059
Change-Id: I210acb816b122857dbbe1ee4dd6a35c5335bf2bf

12 files changed, 48 insertions, 43 deletions

diff --git a/native/jni/com_android_inputmethod_latin_BinaryDictionary.cpp b/native/jni/com_android_inputmethod_latin_BinaryDictionary.cpp
index f8dadb488..1020f6865 100644
--- a/native/jni/com_android_inputmethod_latin_BinaryDictionary.cpp
+++ b/native/jni/com_android_inputmethod_latin_BinaryDictionary.cpp
@@ -409,9 +409,10 @@ static bool latinime_BinaryDictionary_addNgramEntry(JNIEnv *env, jclass clazz, j
     int wordCodePoints[wordLength];
     env->GetIntArrayRegion(word, 0, wordLength, wordCodePoints);
     // Use 1 for count to indicate the ngram has inputted.
-    const NgramProperty ngramProperty(CodePointArrayView(wordCodePoints, wordLength).toVector(),
+    const NgramProperty ngramProperty(ngramContext,
+            CodePointArrayView(wordCodePoints, wordLength).toVector(),
             probability, HistoricalInfo(timestamp, 0 /* level */, 1 /* count */));
-    return dictionary->addNgramEntry(&ngramContext, &ngramProperty);
+    return dictionary->addNgramEntry(&ngramProperty);
 }
 
 static bool latinime_BinaryDictionary_removeNgramEntry(JNIEnv *env, jclass clazz, jlong dict,
@@ -526,12 +527,12 @@ static int latinime_BinaryDictionary_addMultipleDictionaryEntries(JNIEnv *env, j
         if (word0) {
             jint bigramProbability = env->GetIntField(languageModelParam, bigramProbabilityFieldId);
             // Use 1 for count to indicate the bigram has inputted.
-            const NgramProperty ngramProperty(
-                    CodePointArrayView(word1CodePoints, word1Length).toVector(),
-                    bigramProbability, HistoricalInfo(timestamp, 0 /* level */, 1 /* count */));
             const NgramContext ngramContext(word0CodePoints, word0Length,
                     false /* isBeginningOfSentence */);
-            dictionary->addNgramEntry(&ngramContext, &ngramProperty);
+            const NgramProperty ngramProperty(ngramContext,
+                    CodePointArrayView(word1CodePoints, word1Length).toVector(),
+                    bigramProbability, HistoricalInfo(timestamp, 0 /* level */, 1 /* count */));
+            dictionary->addNgramEntry(&ngramProperty);
         }
         if (dictionary->needsToRunGC(true /* mindsBlockByGC */)) {
             return i + 1;
@@ -641,11 +642,8 @@ static bool latinime_BinaryDictionary_migrateNative(JNIEnv *env, jclass clazz, j
                 return false;
             }
         }
-        const NgramContext ngramContext(wordCodePoints, wordCodePointCount,
-                wordProperty.getUnigramProperty()->representsBeginningOfSentence());
         for (const NgramProperty &ngramProperty : *wordProperty.getNgramProperties()) {
-            if (!dictionaryStructureWithBufferPolicy->addNgramEntry(&ngramContext,
-                    &ngramProperty)) {
+            if (!dictionaryStructureWithBufferPolicy->addNgramEntry(&ngramProperty)) {
                 LogUtils::logToJava(env, "Cannot add ngram to the new dict.");
                 return false;
             }
diff --git a/native/jni/src/suggest/core/dictionary/dictionary.cpp b/native/jni/src/suggest/core/dictionary/dictionary.cpp
index 697e99ffb..bfe17cc4c 100644
--- a/native/jni/src/suggest/core/dictionary/dictionary.cpp
+++ b/native/jni/src/suggest/core/dictionary/dictionary.cpp
@@ -140,10 +140,9 @@ bool Dictionary::removeUnigramEntry(const CodePointArrayView codePoints) {
     return mDictionaryStructureWithBufferPolicy->removeUnigramEntry(codePoints);
 }
 
-bool Dictionary::addNgramEntry(const NgramContext *const ngramContext,
-        const NgramProperty *const ngramProperty) {
+bool Dictionary::addNgramEntry(const NgramProperty *const ngramProperty) {
     TimeKeeper::setCurrentTime();
-    return mDictionaryStructureWithBufferPolicy->addNgramEntry(ngramContext, ngramProperty);
+    return mDictionaryStructureWithBufferPolicy->addNgramEntry(ngramProperty);
 }
 
 bool Dictionary::removeNgramEntry(const NgramContext *const ngramContext,
diff --git a/native/jni/src/suggest/core/dictionary/dictionary.h b/native/jni/src/suggest/core/dictionary/dictionary.h
index 843aec473..a5e986d15 100644
--- a/native/jni/src/suggest/core/dictionary/dictionary.h
+++ b/native/jni/src/suggest/core/dictionary/dictionary.h
@@ -85,8 +85,7 @@ class Dictionary {
 
     bool removeUnigramEntry(const CodePointArrayView codePoints);
 
-    bool addNgramEntry(const NgramContext *const ngramContext,
-            const NgramProperty *const ngramProperty);
+    bool addNgramEntry(const NgramProperty *const ngramProperty);
 
     bool removeNgramEntry(const NgramContext *const ngramContext,
             const CodePointArrayView codePoints);
diff --git a/native/jni/src/suggest/core/dictionary/property/ngram_property.h b/native/jni/src/suggest/core/dictionary/property/ngram_property.h
index 8709799f9..e67b4da31 100644
--- a/native/jni/src/suggest/core/dictionary/property/ngram_property.h
+++ b/native/jni/src/suggest/core/dictionary/property/ngram_property.h
@@ -21,15 +21,20 @@
 
 #include "defines.h"
 #include "suggest/core/dictionary/property/historical_info.h"
+#include "suggest/core/session/ngram_context.h"
 
 namespace latinime {
 
 class NgramProperty {
  public:
-    NgramProperty(const std::vector<int> &&targetCodePoints, const int probability,
-            const HistoricalInfo historicalInfo)
-            : mTargetCodePoints(std::move(targetCodePoints)), mProbability(probability),
-              mHistoricalInfo(historicalInfo) {}
+    NgramProperty(const NgramContext &ngramContext, const std::vector<int> &&targetCodePoints,
+            const int probability, const HistoricalInfo historicalInfo)
+            : mNgramContext(ngramContext), mTargetCodePoints(std::move(targetCodePoints)),
+              mProbability(probability), mHistoricalInfo(historicalInfo) {}
+
+    const NgramContext *getNgramContext() const {
+        return &mNgramContext;
+    }
 
     const std::vector<int> *getTargetCodePoints() const {
         return &mTargetCodePoints;
@@ -48,6 +53,7 @@ class NgramProperty {
     DISALLOW_DEFAULT_CONSTRUCTOR(NgramProperty);
     DISALLOW_ASSIGNMENT_OPERATOR(NgramProperty);
 
+    const NgramContext mNgramContext;
     const std::vector<int> mTargetCodePoints;
     const int mProbability;
     const HistoricalInfo mHistoricalInfo;
diff --git a/native/jni/src/suggest/core/dictionary/property/word_property.h b/native/jni/src/suggest/core/dictionary/property/word_property.h
index 0c23e8225..01b8987b5 100644
--- a/native/jni/src/suggest/core/dictionary/property/word_property.h
+++ b/native/jni/src/suggest/core/dictionary/property/word_property.h
@@ -34,9 +34,9 @@ class WordProperty {
             : mCodePoints(), mUnigramProperty(), mNgrams() {}
 
     WordProperty(const std::vector<int> &&codePoints, const UnigramProperty *const unigramProperty,
-            const std::vector<NgramProperty> *const bigrams)
+            const std::vector<NgramProperty> *const ngrams)
             : mCodePoints(std::move(codePoints)), mUnigramProperty(*unigramProperty),
-              mNgrams(*bigrams) {}
+              mNgrams(*ngrams) {}
 
     void outputProperties(JNIEnv *const env, jintArray outCodePoints, jbooleanArray outFlags,
             jintArray outProbabilityInfo, jobject outBigramTargets, jobject outBigramProbabilities,
diff --git a/native/jni/src/suggest/core/policy/dictionary_structure_with_buffer_policy.h b/native/jni/src/suggest/core/policy/dictionary_structure_with_buffer_policy.h
index ceda5c03f..33a0fbc19 100644
--- a/native/jni/src/suggest/core/policy/dictionary_structure_with_buffer_policy.h
+++ b/native/jni/src/suggest/core/policy/dictionary_structure_with_buffer_policy.h
@@ -40,7 +40,6 @@ class UnigramProperty;
  * This class abstracts the structure of dictionaries.
  * Implement this policy to support additional dictionaries.
  */
-// TODO: Use word id instead of terminal PtNode position.
 class DictionaryStructureWithBufferPolicy {
  public:
     typedef std::unique_ptr<DictionaryStructureWithBufferPolicy> StructurePolicyPtr;
@@ -81,8 +80,7 @@ class DictionaryStructureWithBufferPolicy {
     virtual bool removeUnigramEntry(const CodePointArrayView wordCodePoints) = 0;
 
     // Returns whether the update was success or not.
-    virtual bool addNgramEntry(const NgramContext *const ngramContext,
-            const NgramProperty *const ngramProperty) = 0;
+    virtual bool addNgramEntry(const NgramProperty *const ngramProperty) = 0;
 
     // Returns whether the update was success or not.
     virtual bool removeNgramEntry(const NgramContext *const ngramContext,
@@ -106,7 +104,6 @@ class DictionaryStructureWithBufferPolicy {
     virtual void getProperty(const char *const query, const int queryLength, char *const outResult,
             const int maxResultLength) = 0;
 
-    // Used for testing.
     virtual const WordProperty getWordProperty(const CodePointArrayView wordCodePoints) const = 0;
 
     // Method to iterate all words in the dictionary.
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.cpp
index d0dccc3be..66f750120 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.cpp
@@ -344,8 +344,7 @@ bool Ver4PatriciaTriePolicy::removeUnigramEntry(const CodePointArrayView wordCod
     return mNodeWriter.suppressUnigramEntry(&ptNodeParams);
 }
 
-bool Ver4PatriciaTriePolicy::addNgramEntry(const NgramContext *const ngramContext,
-        const NgramProperty *const ngramProperty) {
+bool Ver4PatriciaTriePolicy::addNgramEntry(const NgramProperty *const ngramProperty) {
     if (!mBuffers->isUpdatable()) {
         AKLOGI("Warning: addNgramEntry() is called for non-updatable dictionary.");
         return false;
@@ -355,6 +354,7 @@ bool Ver4PatriciaTriePolicy::addNgramEntry(const NgramContext *const ngramContex
                 mDictBuffer->getTailPosition());
         return false;
     }
+    const NgramContext *const ngramContext = ngramProperty->getNgramContext();
     if (!ngramContext->isValid()) {
         AKLOGE("Ngram context is not valid for adding n-gram entry to the dictionary.");
         return false;
@@ -463,9 +463,9 @@ bool Ver4PatriciaTriePolicy::updateEntriesForWordWithNgramContext(
     }
     const int probabilityForNgram = ngramContext->isNthPrevWordBeginningOfSentence(1 /* n */)
             ? NOT_A_PROBABILITY : probability;
-    const NgramProperty ngramProperty(wordCodePoints.toVector(), probabilityForNgram,
+    const NgramProperty ngramProperty(*ngramContext, wordCodePoints.toVector(), probabilityForNgram,
             historicalInfo);
-    if (!addNgramEntry(ngramContext, &ngramProperty)) {
+    if (!addNgramEntry(&ngramProperty)) {
         AKLOGE("Cannot update unigarm entry in updateEntriesForWordWithNgramContext().");
         return false;
     }
@@ -585,6 +585,8 @@ const WordProperty Ver4PatriciaTriePolicy::getWordProperty(
                             bigramEntry.getHistoricalInfo(), mHeaderPolicy) :
                     bigramEntry.getProbability();
             ngrams.emplace_back(
+                    NgramContext(wordCodePoints.data(), wordCodePoints.size(),
+                            ptNodeParams.representsBeginningOfSentence()),
                     CodePointArrayView(bigramWord1CodePoints, codePointCount).toVector(),
                     probability, *historicalInfo);
         }
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.h b/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.h
index 2cda0d3fa..0480876ed 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.h
@@ -113,8 +113,7 @@ class Ver4PatriciaTriePolicy : public DictionaryStructureWithBufferPolicy {
 
     bool removeUnigramEntry(const CodePointArrayView wordCodePoints);
 
-    bool addNgramEntry(const NgramContext *const ngramContext,
-            const NgramProperty *const ngramProperty);
+    bool addNgramEntry(const NgramProperty *const ngramProperty);
 
     bool removeNgramEntry(const NgramContext *const ngramContext,
             const CodePointArrayView wordCodePoints);
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.cpp
index b7f1199c5..b0d5c901e 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.cpp
@@ -451,6 +451,8 @@ const WordProperty PatriciaTriePolicy::getWordProperty(
                     bigramWord1CodePoints, &word1Probability);
             const int probability = getProbability(word1Probability, bigramsIt.getProbability());
             ngrams.emplace_back(
+                    NgramContext(wordCodePoints.data(), wordCodePoints.size(),
+                            ptNodeParams.representsBeginningOfSentence()),
                     CodePointArrayView(bigramWord1CodePoints, word1CodePointCount).toVector(),
                     probability, HistoricalInfo());
         }
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.h b/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.h
index b17681388..8933962ab 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.h
@@ -93,8 +93,7 @@ class PatriciaTriePolicy : public DictionaryStructureWithBufferPolicy {
         return false;
     }
 
-    bool addNgramEntry(const NgramContext *const ngramContext,
-            const NgramProperty *const ngramProperty) {
+    bool addNgramEntry(const NgramProperty *const ngramProperty) {
         // This method should not be called for non-updatable dictionary.
         AKLOGI("Warning: addNgramEntry() is called for non-updatable dictionary.");
         return false;
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 ead1bde50..094106746 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
@@ -264,8 +264,7 @@ bool Ver4PatriciaTriePolicy::removeUnigramEntry(const CodePointArrayView wordCod
     return true;
 }
 
-bool Ver4PatriciaTriePolicy::addNgramEntry(const NgramContext *const ngramContext,
-        const NgramProperty *const ngramProperty) {
+bool Ver4PatriciaTriePolicy::addNgramEntry(const NgramProperty *const ngramProperty) {
     if (!mBuffers->isUpdatable()) {
         AKLOGI("Warning: addNgramEntry() is called for non-updatable dictionary.");
         return false;
@@ -275,6 +274,7 @@ bool Ver4PatriciaTriePolicy::addNgramEntry(const NgramContext *const ngramContex
                 mDictBuffer->getTailPosition());
         return false;
     }
+    const NgramContext *const ngramContext = ngramProperty->getNgramContext();
     if (!ngramContext->isValid()) {
         AKLOGE("Ngram context is not valid for adding n-gram entry to the dictionary.");
         return false;
@@ -451,7 +451,8 @@ bool Ver4PatriciaTriePolicy::needsToRunGC(const bool mindsBlockByGC) const {
         // Needs to reduce dictionary size.
         return true;
     } else if (mHeaderPolicy->isDecayingDict()) {
-        return ForgettingCurveUtils::needsToDecay(mindsBlockByGC, mEntryCounters.getEntryCounts(), mHeaderPolicy);
+        return ForgettingCurveUtils::needsToDecay(mindsBlockByGC, mEntryCounters.getEntryCounts(),
+                mHeaderPolicy);
     }
     return false;
 }
@@ -501,12 +502,16 @@ const WordProperty Ver4PatriciaTriePolicy::getWordProperty(
             prevWordIds)) {
         const int codePointCount = getCodePointsAndReturnCodePointCount(entry.getWordId(),
                 MAX_WORD_LENGTH, bigramWord1CodePoints);
-        const ProbabilityEntry probabilityEntry = entry.getProbabilityEntry();
-        const HistoricalInfo *const historicalInfo = probabilityEntry.getHistoricalInfo();
-        const int probability = probabilityEntry.hasHistoricalInfo() ?
+        const ProbabilityEntry ngramProbabilityEntry = entry.getProbabilityEntry();
+        const HistoricalInfo *const historicalInfo = ngramProbabilityEntry.getHistoricalInfo();
+        const int probability = ngramProbabilityEntry.hasHistoricalInfo() ?
                 ForgettingCurveUtils::decodeProbability(historicalInfo, mHeaderPolicy) :
-                probabilityEntry.getProbability();
-        ngrams.emplace_back(CodePointArrayView(bigramWord1CodePoints, codePointCount).toVector(),
+                ngramProbabilityEntry.getProbability();
+        ngrams.emplace_back(
+                NgramContext(
+                        wordCodePoints.data(), wordCodePoints.size(),
+                        probabilityEntry.representsBeginningOfSentence()),
+                CodePointArrayView(bigramWord1CodePoints, codePointCount).toVector(),
                 probability, *historicalInfo);
     }
     // Fetch shortcut information.
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.h b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.h
index e3611cb32..13700b390 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.h
@@ -92,8 +92,7 @@ class Ver4PatriciaTriePolicy : public DictionaryStructureWithBufferPolicy {
 
     bool removeUnigramEntry(const CodePointArrayView wordCodePoints);
 
-    bool addNgramEntry(const NgramContext *const ngramContext,
-            const NgramProperty *const ngramProperty);
+    bool addNgramEntry(const NgramProperty *const ngramProperty);
 
     bool removeNgramEntry(const NgramContext *const ngramContext,
             const CodePointArrayView wordCodePoints);