Support dumping ngram entries.

Bug: 14425059
Change-Id: Ib03a0c3d166ed6f1e60c67127b28006d55143b6b

4 files changed, 47 insertions, 23 deletions

diff --git a/native/jni/src/suggest/core/dictionary/property/word_property.cpp b/native/jni/src/suggest/core/dictionary/property/word_property.cpp
index a707f1ba2..019f0880f 100644
--- a/native/jni/src/suggest/core/dictionary/property/word_property.cpp
+++ b/native/jni/src/suggest/core/dictionary/property/word_property.cpp
@@ -22,8 +22,9 @@
 namespace latinime {
 
 void WordProperty::outputProperties(JNIEnv *const env, jintArray outCodePoints,
-        jbooleanArray outFlags, jintArray outProbabilityInfo, jobject outBigramTargets,
-        jobject outBigramProbabilities, jobject outShortcutTargets,
+        jbooleanArray outFlags, jintArray outProbabilityInfo,
+        jobject outNgramPrevWordsArray, jobject outNgramPrevWordIsBeginningOfSentenceArray,
+        jobject outNgramTargets, jobject outNgramProbabilities, jobject outShortcutTargets,
         jobject outShortcutProbabilities) const {
     JniDataUtils::outputCodePoints(env, outCodePoints, 0 /* start */,
             MAX_WORD_LENGTH /* maxLength */, mCodePoints.data(), mCodePoints.size(),
@@ -43,16 +44,39 @@ void WordProperty::outputProperties(JNIEnv *const env, jintArray outCodePoints,
     jclass arrayListClass = env->FindClass("java/util/ArrayList");
     jmethodID addMethodId = env->GetMethodID(arrayListClass, "add", "(Ljava/lang/Object;)Z");
 
-    // Output bigrams.
-    // TODO: Support n-gram
+    // Output ngrams.
+    jclass intArrayClass = env->FindClass("[I");
     for (const auto &ngramProperty : mNgrams) {
-        const std::vector<int> *const word1CodePoints = ngramProperty.getTargetCodePoints();
-        jintArray bigramWord1CodePointArray = env->NewIntArray(word1CodePoints->size());
-        JniDataUtils::outputCodePoints(env, bigramWord1CodePointArray, 0 /* start */,
-                word1CodePoints->size(), word1CodePoints->data(), word1CodePoints->size(),
-                false /* needsNullTermination */);
-        env->CallBooleanMethod(outBigramTargets, addMethodId, bigramWord1CodePointArray);
-        env->DeleteLocalRef(bigramWord1CodePointArray);
+        const NgramContext *const ngramContext = ngramProperty.getNgramContext();
+        jobjectArray prevWordWordCodePointsArray = env->NewObjectArray(
+                ngramContext->getPrevWordCount(), intArrayClass, nullptr);
+        jbooleanArray prevWordIsBeginningOfSentenceArray =
+                env->NewBooleanArray(ngramContext->getPrevWordCount());
+        for (size_t i = 0; i < ngramContext->getPrevWordCount(); ++i) {
+            const CodePointArrayView codePoints = ngramContext->getNthPrevWordCodePoints(i + 1);
+            jintArray prevWordCodePoints = env->NewIntArray(codePoints.size());
+            JniDataUtils::outputCodePoints(env, prevWordCodePoints, 0 /* start */,
+                    codePoints.size(), codePoints.data(), codePoints.size(),
+                    false /* needsNullTermination */);
+            env->SetObjectArrayElement(prevWordWordCodePointsArray, i, prevWordCodePoints);
+            env->DeleteLocalRef(prevWordCodePoints);
+            JniDataUtils::putBooleanToArray(env, prevWordIsBeginningOfSentenceArray, i,
+                    ngramContext->isNthPrevWordBeginningOfSentence(i + 1));
+        }
+        env->CallBooleanMethod(outNgramPrevWordsArray, addMethodId, prevWordWordCodePointsArray);
+        env->CallBooleanMethod(outNgramPrevWordIsBeginningOfSentenceArray, addMethodId,
+                prevWordIsBeginningOfSentenceArray);
+        env->DeleteLocalRef(prevWordWordCodePointsArray);
+        env->DeleteLocalRef(prevWordIsBeginningOfSentenceArray);
+
+        const std::vector<int> *const targetWordCodePoints = ngramProperty.getTargetCodePoints();
+        jintArray targetWordCodePointArray = env->NewIntArray(targetWordCodePoints->size());
+        JniDataUtils::outputCodePoints(env, targetWordCodePointArray, 0 /* start */,
+                targetWordCodePoints->size(), targetWordCodePoints->data(),
+                targetWordCodePoints->size(), false /* needsNullTermination */);
+        env->CallBooleanMethod(outNgramTargets, addMethodId, targetWordCodePointArray);
+        env->DeleteLocalRef(targetWordCodePointArray);
+
         const HistoricalInfo &ngramHistoricalInfo = ngramProperty.getHistoricalInfo();
         int bigramProbabilityInfo[] = {ngramProperty.getProbability(),
                 ngramHistoricalInfo.getTimestamp(), ngramHistoricalInfo.getLevel(),
@@ -60,7 +84,7 @@ void WordProperty::outputProperties(JNIEnv *const env, jintArray outCodePoints,
         jintArray bigramProbabilityInfoArray = env->NewIntArray(NELEMS(bigramProbabilityInfo));
         env->SetIntArrayRegion(bigramProbabilityInfoArray, 0 /* start */,
                 NELEMS(bigramProbabilityInfo), bigramProbabilityInfo);
-        env->CallBooleanMethod(outBigramProbabilities, addMethodId, bigramProbabilityInfoArray);
+        env->CallBooleanMethod(outNgramProbabilities, addMethodId, bigramProbabilityInfoArray);
         env->DeleteLocalRef(bigramProbabilityInfoArray);
     }
 
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 01b8987b5..b5314faaa 100644
--- a/native/jni/src/suggest/core/dictionary/property/word_property.h
+++ b/native/jni/src/suggest/core/dictionary/property/word_property.h
@@ -39,8 +39,10 @@ class WordProperty {
               mNgrams(*ngrams) {}
 
     void outputProperties(JNIEnv *const env, jintArray outCodePoints, jbooleanArray outFlags,
-            jintArray outProbabilityInfo, jobject outBigramTargets, jobject outBigramProbabilities,
-            jobject outShortcutTargets, jobject outShortcutProbabilities) const;
+            jintArray outProbabilityInfo, jobject outNgramPrevWordsArray,
+            jobject outNgramPrevWordIsBeginningOfSentenceArray, jobject outNgramTargets,
+            jobject outNgramProbabilities, jobject outShortcutTargets,
+            jobject outShortcutProbabilities) const;
 
     const UnigramProperty *getUnigramProperty() const {
         return &mUnigramProperty;
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 b96290437..509bd683b 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
@@ -90,8 +90,8 @@ const WordAttributes LanguageModelDictContent::getWordAttributes(const WordIdArr
         // TODO: Some flags in unigramProbabilityEntry should be overwritten by flags in
         // probabilityEntry.
         const ProbabilityEntry unigramProbabilityEntry = getProbabilityEntry(wordId);
-        return WordAttributes(probability, unigramProbabilityEntry.isNotAWord(),
-                unigramProbabilityEntry.isBlacklisted(),
+        return WordAttributes(probability, unigramProbabilityEntry.isBlacklisted(),
+                unigramProbabilityEntry.isNotAWord(),
                 unigramProbabilityEntry.isPossiblyOffensive());
     }
     // Cannot find the word.
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 193326d82..249d822b2 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
@@ -488,9 +488,6 @@ const WordProperty Ver4PatriciaTriePolicy::getWordProperty(
         AKLOGE("getWordProperty is called for invalid word.");
         return WordProperty();
     }
-    const int ptNodePos =
-            mBuffers->getTerminalPositionLookupTable()->getTerminalPtNodePosition(wordId);
-    const PtNodeParams ptNodeParams = mNodeReader.fetchPtNodeParamsInBufferFromPtNodePos(ptNodePos);
     const LanguageModelDictContent *const languageModelDictContent =
             mBuffers->getLanguageModelDictContent();
     // Fetch ngram information.
@@ -541,12 +538,13 @@ const WordProperty Ver4PatriciaTriePolicy::getWordProperty(
                     shortcutProbability);
         }
     }
-    const ProbabilityEntry probabilityEntry = languageModelDictContent->getProbabilityEntry(
-            ptNodeParams.getTerminalId());
+    const WordAttributes wordAttributes = languageModelDictContent->getWordAttributes(
+            WordIdArrayView(), wordId, mHeaderPolicy);
+    const ProbabilityEntry probabilityEntry = languageModelDictContent->getProbabilityEntry(wordId);
     const HistoricalInfo *const historicalInfo = probabilityEntry.getHistoricalInfo();
     const UnigramProperty unigramProperty(probabilityEntry.representsBeginningOfSentence(),
-            probabilityEntry.isNotAWord(), probabilityEntry.isBlacklisted(),
-            probabilityEntry.isPossiblyOffensive(), probabilityEntry.getProbability(),
+            wordAttributes.isNotAWord(), wordAttributes.isBlacklisted(),
+            wordAttributes.isPossiblyOffensive(), wordAttributes.getProbability(),
             *historicalInfo, std::move(shortcuts));
     return WordProperty(wordCodePoints.toVector(), &unigramProperty, &ngrams);
 }