3 files changed, 24 insertions, 7 deletions
diff --git a/native/jni/src/utils/byte_array_view.h b/native/jni/src/utils/byte_array_view.h
index 10d7ae278..2b778af6f 100644
--- a/native/jni/src/utils/byte_array_view.h
+++ b/native/jni/src/utils/byte_array_view.h
@@ -42,6 +42,13 @@ class ReadOnlyByteArrayView {
         return mPtr;
     }
 
+    AK_FORCE_INLINE const ReadOnlyByteArrayView skip(const size_t n) const {
+        if (mSize <= n) {
+            return ReadOnlyByteArrayView();
+        }
+        return ReadOnlyByteArrayView(mPtr + n, mSize - n);
+    }
+
  private:
     DISALLOW_ASSIGNMENT_OPERATOR(ReadOnlyByteArrayView);
 
diff --git a/native/jni/src/utils/int_array_view.h b/native/jni/src/utils/int_array_view.h
index cc5f328ba..f3a8589ca 100644
--- a/native/jni/src/utils/int_array_view.h
+++ b/native/jni/src/utils/int_array_view.h
@@ -115,6 +115,20 @@ class IntArrayView {
         memmove(buffer->data() + offset, mPtr, sizeof(int) * mSize);
     }
 
+    AK_FORCE_INLINE int firstOrDefault(const int defaultValue) const {
+        if (empty()) {
+            return defaultValue;
+        }
+        return mPtr[0];
+    }
+
+    AK_FORCE_INLINE int lastOrDefault(const int defaultValue) const {
+        if (empty()) {
+            return defaultValue;
+        }
+        return mPtr[mSize - 1];
+    }
+
  private:
     DISALLOW_ASSIGNMENT_OPERATOR(IntArrayView);
 
diff --git a/native/jni/src/utils/jni_data_utils.h b/native/jni/src/utils/jni_data_utils.h
index cb82d3c3b..235a03bba 100644
--- a/native/jni/src/utils/jni_data_utils.h
+++ b/native/jni/src/utils/jni_data_utils.h
@@ -97,17 +97,13 @@ class JniDataUtils {
     }
 
     static PrevWordsInfo constructPrevWordsInfo(JNIEnv *env, jobjectArray prevWordCodePointArrays,
-            jbooleanArray isBeginningOfSentenceArray) {
+            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;
@@ -124,7 +120,7 @@ class JniDataUtils {
             isBeginningOfSentence[i] = isBeginningOfSentenceBoolean == JNI_TRUE;
         }
         return PrevWordsInfo(prevWordCodePoints, prevWordCodePointCount, isBeginningOfSentence,
-                MAX_PREV_WORD_COUNT_FOR_N_GRAM);
+                prevWordCount);
     }
 
     static void putBooleanToArray(JNIEnv *env, jbooleanArray array, const int index,