1 files changed, 6 insertions, 10 deletions
diff --git a/native/jni/src/utils/jni_data_utils.h b/native/jni/src/utils/jni_data_utils.h
index cb82d3c3b..25cc41742 100644
--- a/native/jni/src/utils/jni_data_utils.h
+++ b/native/jni/src/utils/jni_data_utils.h
@@ -21,7 +21,7 @@
 
 #include "defines.h"
 #include "jni.h"
-#include "suggest/core/session/prev_words_info.h"
+#include "suggest/core/session/ngram_context.h"
 #include "suggest/core/policy/dictionary_header_structure_policy.h"
 #include "suggest/policyimpl/dictionary/header/header_read_write_utils.h"
 #include "utils/char_utils.h"
@@ -96,18 +96,14 @@ class JniDataUtils {
         }
     }
 
-    static PrevWordsInfo constructPrevWordsInfo(JNIEnv *env, jobjectArray prevWordCodePointArrays,
-            jbooleanArray isBeginningOfSentenceArray) {
+    static NgramContext constructNgramContext(JNIEnv *env, jobjectArray prevWordCodePointArrays,
+            jbooleanArray isBeginningOfSentenceArray, const size_t prevWordCount) {
         int prevWordCodePoints[MAX_PREV_WORD_COUNT_FOR_N_GRAM][MAX_WORD_LENGTH];
         int prevWordCodePointCount[MAX_PREV_WORD_COUNT_FOR_N_GRAM];
         bool isBeginningOfSentence[MAX_PREV_WORD_COUNT_FOR_N_GRAM];
-        jsize prevWordsCount = env->GetArrayLength(prevWordCodePointArrays);
-        for (size_t i = 0; i < NELEMS(prevWordCodePoints); ++i) {
+        for (size_t i = 0; i < prevWordCount; ++i) {
             prevWordCodePointCount[i] = 0;
             isBeginningOfSentence[i] = false;
-            if (prevWordsCount <= static_cast<int>(i)) {
-                continue;
-            }
             jintArray prevWord = (jintArray)env->GetObjectArrayElement(prevWordCodePointArrays, i);
             if (!prevWord) {
                 continue;
@@ -123,8 +119,8 @@ class JniDataUtils {
                     &isBeginningOfSentenceBoolean);
             isBeginningOfSentence[i] = isBeginningOfSentenceBoolean == JNI_TRUE;
         }
-        return PrevWordsInfo(prevWordCodePoints, prevWordCodePointCount, isBeginningOfSentence,
-                MAX_PREV_WORD_COUNT_FOR_N_GRAM);
+        return NgramContext(prevWordCodePoints, prevWordCodePointCount, isBeginningOfSentence,
+                prevWordCount);
     }
 
     static void putBooleanToArray(JNIEnv *env, jbooleanArray array, const int index,