1 files changed, 37 insertions, 11 deletions
diff --git a/native/jni/com_android_inputmethod_latin_BinaryDictionary.cpp b/native/jni/com_android_inputmethod_latin_BinaryDictionary.cpp
index 586a306ec..c919ebd91 100644
--- a/native/jni/com_android_inputmethod_latin_BinaryDictionary.cpp
+++ b/native/jni/com_android_inputmethod_latin_BinaryDictionary.cpp
@@ -24,7 +24,7 @@
 #include "jni.h"
 #include "jni_common.h"
 #include "suggest/core/dictionary/dictionary.h"
-#include "suggest/core/dictionary/unigram_property.h"
+#include "suggest/core/dictionary/word_property.h"
 #include "suggest/core/suggest_options.h"
 #include "suggest/policyimpl/dictionary/structure/dictionary_structure_with_buffer_policy_factory.h"
 #include "suggest/policyimpl/dictionary/utils/dict_file_writing_utils.h"
@@ -260,19 +260,39 @@ static jint latinime_BinaryDictionary_getBigramProbability(JNIEnv *env, jclass c
             word1Length);
 }
 
-static void latinime_BinaryDictionary_getUnigramProperty(JNIEnv *env, jclass clazz,
+// Method to iterate all words in the dictionary for makedict.
+// If token is 0, this method newly starts iterating the dictionary. This method returns 0 when
+// the dictionary does not have a next word.
+static jint latinime_BinaryDictionary_getNextWord(JNIEnv *env, jclass clazz,
+        jlong dict, jint token, jintArray outCodePoints) {
+    Dictionary *dictionary = reinterpret_cast<Dictionary *>(dict);
+    if (!dictionary) return 0;
+    const jsize outCodePointsLength = env->GetArrayLength(outCodePoints);
+    if (outCodePointsLength != MAX_WORD_LENGTH) {
+        AKLOGE("Invalid outCodePointsLength: %d", outCodePointsLength);
+        ASSERT(false);
+        return 0;
+    }
+    int wordCodePoints[outCodePointsLength];
+    memset(wordCodePoints, 0, sizeof(wordCodePoints));
+    const int nextToken = dictionary->getNextWordAndNextToken(token, wordCodePoints);
+    env->SetIntArrayRegion(outCodePoints, 0, outCodePointsLength, wordCodePoints);
+    return nextToken;
+}
+
+static void latinime_BinaryDictionary_getWordProperty(JNIEnv *env, jclass clazz,
         jlong dict, jintArray word, jintArray outCodePoints, jbooleanArray outFlags,
-        jintArray outProbability, jintArray outHistoricalInfo, jobject outShortcutTargets,
-        jobject outShortcutProbabilities) {
+        jintArray outProbabilityInfo, jobject outBigramTargets, jobject outBigramProbabilityInfo,
+        jobject outShortcutTargets, jobject outShortcutProbabilities) {
     Dictionary *dictionary = reinterpret_cast<Dictionary *>(dict);
     if (!dictionary) return;
     const jsize wordLength = env->GetArrayLength(word);
     int wordCodePoints[wordLength];
     env->GetIntArrayRegion(word, 0, wordLength, wordCodePoints);
-    const UnigramProperty unigramProperty = dictionary->getUnigramProperty(
-            wordCodePoints, wordLength);
-    unigramProperty.outputProperties(env, outCodePoints, outFlags, outProbability,
-            outHistoricalInfo, outShortcutTargets, outShortcutProbabilities);
+    const WordProperty wordProperty = dictionary->getWordProperty(wordCodePoints, wordLength);
+    wordProperty.outputProperties(env, outCodePoints, outFlags, outProbabilityInfo,
+            outBigramTargets, outBigramProbabilityInfo, outShortcutTargets,
+            outShortcutProbabilities);
 }
 
 static jfloat latinime_BinaryDictionary_calcNormalizedScore(JNIEnv *env, jclass clazz,
@@ -521,9 +541,15 @@ static const JNINativeMethod sMethods[] = {
         reinterpret_cast<void *>(latinime_BinaryDictionary_getBigramProbability)
     },
     {
-        const_cast<char *>("getUnigramPropertyNative"),
-        const_cast<char *>("(J[I[I[Z[I[ILjava/util/ArrayList;Ljava/util/ArrayList;)V"),
-        reinterpret_cast<void *>(latinime_BinaryDictionary_getUnigramProperty)
+        const_cast<char *>("getWordPropertyNative"),
+        const_cast<char *>("(J[I[I[Z[ILjava/util/ArrayList;Ljava/util/ArrayList;"
+                "Ljava/util/ArrayList;Ljava/util/ArrayList;)V"),
+        reinterpret_cast<void *>(latinime_BinaryDictionary_getWordProperty)
+    },
+    {
+        const_cast<char *>("getNextWordNative"),
+        const_cast<char *>("(JI[I)I"),
+        reinterpret_cast<void *>(latinime_BinaryDictionary_getNextWord)
     },
     {
         const_cast<char *>("calcNormalizedScoreNative"),