1 files changed, 21 insertions, 16 deletions
diff --git a/native/jni/com_android_inputmethod_latin_DicTraverseSession.cpp b/native/jni/com_android_inputmethod_latin_DicTraverseSession.cpp
index 0f5c39642..947360e8e 100644
--- a/native/jni/com_android_inputmethod_latin_DicTraverseSession.cpp
+++ b/native/jni/com_android_inputmethod_latin_DicTraverseSession.cpp
@@ -17,37 +17,42 @@
 #define LOG_TAG "LatinIME: jni: Session"
 
 #include "com_android_inputmethod_latin_DicTraverseSession.h"
+#include "dic_traverse_wrapper.h"
 #include "jni.h"
 #include "jni_common.h"
 
 namespace latinime {
-void *(*DicTraverseWrapper::sDicTraverseSessionFactoryMethod)() = 0;
-void (*DicTraverseWrapper::sDicTraverseSessionReleaseMethod)(void *) = 0;
-void (*DicTraverseWrapper::sDicTraverseSessionInitMethod)(
-        JNIEnv *, void *, const jintArray, const jint) = 0;
-
-static jlong latinime_setDicTraverseSession(JNIEnv *env, jobject object,
-        jstring localejStr) {
-    void *traverseSession = DicTraverseWrapper::getDicTraverseSession();
+class Dictionary;
+static jlong latinime_setDicTraverseSession(JNIEnv *env, jobject object, jstring localeJStr) {
+    void *traverseSession = DicTraverseWrapper::getDicTraverseSession(env, localeJStr);
     return reinterpret_cast<jlong>(traverseSession);
 }
 
-static void latinime_initDicTraverseSession(JNIEnv *env, jlong traverseSession,
-        jintArray previousWord, jint previousWordLength) {
+static void latinime_initDicTraverseSession(JNIEnv *env, jobject object, jlong traverseSession,
+        jlong dictionary, jintArray previousWord, jint previousWordLength) {
     void *ts = reinterpret_cast<void*>(traverseSession);
-    DicTraverseWrapper::initDicTraverseSession(env, ts, previousWord, previousWordLength);
+    Dictionary *dict = reinterpret_cast<Dictionary*>(dictionary);
+    if (!previousWord) {
+        DicTraverseWrapper::initDicTraverseSession(ts, dict, 0, 0);
+        return;
+    }
+    int prevWord[previousWordLength];
+    env->GetIntArrayRegion(previousWord, 0, previousWordLength, prevWord);
+    DicTraverseWrapper::initDicTraverseSession(ts, dict, prevWord, previousWordLength);
 }
 
-static void latinime_DicTraverseSession_release(
-        JNIEnv *env, jobject object, jlong traverseSession) {
+static void latinime_releaseDicTraverseSession(JNIEnv *env, jobject object, jlong traverseSession) {
     void *ts = reinterpret_cast<void*>(traverseSession);
     DicTraverseWrapper::releaseDicTraverseSession(ts);
 }
 
 static JNINativeMethod sMethods[] = {
-    {"setDicTraverseSessionNative", "(Ljava/lang/String;)J", (void*)latinime_setDicTraverseSession},
-    {"initDicTraverseSessionNative", "(J[II)V", (void*)latinime_initDicTraverseSession},
-    {"releaseDicTraverseSessionNative", "(J)V", (void*)latinime_DicTraverseSession_release}
+    {"setDicTraverseSessionNative", "(Ljava/lang/String;)J",
+            reinterpret_cast<void *>(latinime_setDicTraverseSession)},
+    {"initDicTraverseSessionNative", "(JJ[II)V",
+            reinterpret_cast<void *>(latinime_initDicTraverseSession)},
+    {"releaseDicTraverseSessionNative", "(J)V",
+            reinterpret_cast<void *>(latinime_releaseDicTraverseSession)}
 };
 
 int register_DicTraverseSession(JNIEnv *env) {