1 files changed, 11 insertions, 23 deletions

diff --git a/native/jni/com_android_inputmethod_latin_BinaryDictionary.cpp b/native/jni/com_android_inputmethod_latin_BinaryDictionary.cpp
index 18b78c4df..28aaf2d1a 100644
--- a/native/jni/com_android_inputmethod_latin_BinaryDictionary.cpp
+++ b/native/jni/com_android_inputmethod_latin_BinaryDictionary.cpp
@@ -28,6 +28,7 @@
 #include "suggest/core/dictionary/property/unigram_property.h"
 #include "suggest/core/dictionary/property/word_property.h"
 #include "suggest/core/result/suggestion_results.h"
+#include "suggest/core/session/prev_words_info.h"
 #include "suggest/core/suggest_options.h"
 #include "suggest/policyimpl/dictionary/structure/dictionary_structure_with_buffer_policy_factory.h"
 #include "utils/char_utils.h"
@@ -190,7 +191,9 @@ static void latinime_BinaryDictionary_getSuggestions(JNIEnv *env, jclass clazz,
     ProximityInfo *pInfo = reinterpret_cast<ProximityInfo *>(proximityInfo);
     DicTraverseSession *traverseSession =
             reinterpret_cast<DicTraverseSession *>(dicTraverseSession);
-
+    if (!traverseSession) {
+        return;
+    }
     // Input values
     int xCoordinates[inputSize];
     int yCoordinates[inputSize];
@@ -245,15 +248,15 @@ static void latinime_BinaryDictionary_getSuggestions(JNIEnv *env, jclass clazz,
     float languageWeight;
     env->GetFloatArrayRegion(inOutLanguageWeight, 0, 1 /* len */, &languageWeight);
     SuggestionResults suggestionResults(MAX_RESULTS);
+    const PrevWordsInfo prevWordsInfo(prevWordCodePoints, prevWordCodePointsLength,
+            false /* isStartOfSentence */);
     if (givenSuggestOptions.isGesture() || inputSize > 0) {
         // TODO: Use SuggestionResults to return suggestions.
         dictionary->getSuggestions(pInfo, traverseSession, xCoordinates, yCoordinates,
-                times, pointerIds, inputCodePoints, inputSize, prevWordCodePoints,
-                prevWordCodePointsLength, &givenSuggestOptions, languageWeight,
-                &suggestionResults);
+                times, pointerIds, inputCodePoints, inputSize, &prevWordsInfo,
+                &givenSuggestOptions, languageWeight, &suggestionResults);
     } else {
-        dictionary->getPredictions(prevWordCodePoints, prevWordCodePointsLength,
-                &suggestionResults);
+        dictionary->getPredictions(&prevWordsInfo, &suggestionResults);
     }
     suggestionResults.outputSuggestions(env, outSuggestionCount, outCodePointsArray,
             outScoresArray, outSpaceIndicesArray, outTypesArray,
@@ -280,8 +283,8 @@ static jint latinime_BinaryDictionary_getBigramProbability(JNIEnv *env, jclass c
     int word1CodePoints[word1Length];
     env->GetIntArrayRegion(word0, 0, word0Length, word0CodePoints);
     env->GetIntArrayRegion(word1, 0, word1Length, word1CodePoints);
-    return dictionary->getBigramProbability(word0CodePoints, word0Length, word1CodePoints,
-            word1Length);
+    const PrevWordsInfo prevWordsInfo(word0CodePoints, word0Length, false /* isStartOfSentence */);
+    return dictionary->getBigramProbability(&prevWordsInfo, word1CodePoints, word1Length);
 }
 
 // Method to iterate all words in the dictionary for makedict.
@@ -467,16 +470,6 @@ static int latinime_BinaryDictionary_addMultipleDictionaryEntries(JNIEnv *env, j
     return languageModelParamCount;
 }
 
-static int latinime_BinaryDictionary_calculateProbabilityNative(JNIEnv *env, jclass clazz,
-        jlong dict, jint unigramProbability, jint bigramProbability) {
-    Dictionary *dictionary = reinterpret_cast<Dictionary *>(dict);
-    if (!dictionary) {
-        return NOT_A_PROBABILITY;
-    }
-    return dictionary->getDictionaryStructurePolicy()->getProbability(unigramProbability,
-            bigramProbability);
-}
-
 static jstring latinime_BinaryDictionary_getProperty(JNIEnv *env, jclass clazz, jlong dict,
         jstring query) {
     Dictionary *dictionary = reinterpret_cast<Dictionary *>(dict);
@@ -670,11 +663,6 @@ static const JNINativeMethod sMethods[] = {
         reinterpret_cast<void *>(latinime_BinaryDictionary_addMultipleDictionaryEntries)
     },
     {
-        const_cast<char *>("calculateProbabilityNative"),
-        const_cast<char *>("(JII)I"),
-        reinterpret_cast<void *>(latinime_BinaryDictionary_calculateProbabilityNative)
-    },
-    {
         const_cast<char *>("getPropertyNative"),
         const_cast<char *>("(JLjava/lang/String;)Ljava/lang/String;"),
         reinterpret_cast<void *>(latinime_BinaryDictionary_getProperty)