3 files changed, 36 insertions, 5 deletions
diff --git a/dictionary/jni/com_android_inputmethod_latin_BinaryDictionary.cpp b/dictionary/jni/com_android_inputmethod_latin_BinaryDictionary.cpp
index 3076085e4..3dda062fe 100644
--- a/dictionary/jni/com_android_inputmethod_latin_BinaryDictionary.cpp
+++ b/dictionary/jni/com_android_inputmethod_latin_BinaryDictionary.cpp
@@ -89,7 +89,7 @@ static jint latinime_BinaryDictionary_open
 static int latinime_BinaryDictionary_getSuggestions(
         JNIEnv *env, jobject object, jint dict, jintArray inputArray, jint arraySize,
         jcharArray outputArray, jintArray frequencyArray, jint maxWordLength, jint maxWords,
-        jint maxAlternatives, jint skipPos)
+        jint maxAlternatives, jint skipPos, jintArray nextLettersArray, jint nextLettersSize)
 {
     Dictionary *dictionary = (Dictionary*) dict;
     if (dictionary == NULL)
@@ -98,13 +98,16 @@ static int latinime_BinaryDictionary_getSuggestions(
     int *frequencies = env->GetIntArrayElements(frequencyArray, NULL);
     int *inputCodes = env->GetIntArrayElements(inputArray, NULL);
     jchar *outputChars = env->GetCharArrayElements(outputArray, NULL);
+    int *nextLetters = nextLettersArray != NULL ? env->GetIntArrayElements(nextLettersArray, NULL)
+            : NULL;
 
     int count = dictionary->getSuggestions(inputCodes, arraySize, (unsigned short*) outputChars, frequencies,
-            maxWordLength, maxWords, maxAlternatives, skipPos);
+            maxWordLength, maxWords, maxAlternatives, skipPos, nextLetters, nextLettersSize);
     
     env->ReleaseIntArrayElements(frequencyArray, frequencies, 0);
     env->ReleaseIntArrayElements(inputArray, inputCodes, JNI_ABORT);
     env->ReleaseCharArrayElements(outputArray, outputChars, 0);
+    env->ReleaseIntArrayElements(nextLettersArray, nextLetters, 0);
     
     return count;
 }
@@ -136,7 +139,7 @@ static JNINativeMethod gMethods[] = {
     {"openNative",           "(Landroid/content/res/AssetManager;Ljava/lang/String;II)I",
                                           (void*)latinime_BinaryDictionary_open},
     {"closeNative",          "(I)V",            (void*)latinime_BinaryDictionary_close},
-    {"getSuggestionsNative", "(I[II[C[IIIII)I",  (void*)latinime_BinaryDictionary_getSuggestions},
+    {"getSuggestionsNative", "(I[II[C[IIIII[II)I",  (void*)latinime_BinaryDictionary_getSuggestions},
     {"isValidWordNative",    "(I[CI)Z",         (void*)latinime_BinaryDictionary_isValidWord}
 };
 
diff --git a/dictionary/src/dictionary.cpp b/dictionary/src/dictionary.cpp
index 306aff527..6e6f44182 100644
--- a/dictionary/src/dictionary.cpp
+++ b/dictionary/src/dictionary.cpp
@@ -49,7 +49,8 @@ Dictionary::~Dictionary()
 }
 
 int Dictionary::getSuggestions(int *codes, int codesSize, unsigned short *outWords, int *frequencies,
-        int maxWordLength, int maxWords, int maxAlternatives, int skipPos)
+        int maxWordLength, int maxWords, int maxAlternatives, int skipPos,
+        int *nextLetters, int nextLettersSize)
 {
     int suggWords;
     mFrequencies = frequencies;
@@ -61,6 +62,8 @@ int Dictionary::getSuggestions(int *codes, int codesSize, unsigned short *outWor
     mMaxWords = maxWords;
     mSkipPos = skipPos;
     mMaxEditDistance = mInputLength < 5 ? 2 : mInputLength / 2;
+    mNextLettersFrequencies = nextLetters;
+    mNextLettersSize = nextLettersSize;
 
     getWordsRec(0, 0, mInputLength * 3, false, 1, 0, 0);
 
@@ -68,9 +71,27 @@ int Dictionary::getSuggestions(int *codes, int codesSize, unsigned short *outWor
     suggWords = 0;
     while (suggWords < mMaxWords && mFrequencies[suggWords] > 0) suggWords++;
     if (DEBUG_DICT) LOGI("Returning %d words", suggWords);
+
+    if (DEBUG_DICT) {
+        LOGI("Next letters: ");
+        for (int k = 0; k < nextLettersSize; k++) {
+            if (mNextLettersFrequencies[k] > 0) {
+                LOGI("%c = %d,", k, mNextLettersFrequencies[k]);
+            }
+        }
+        LOGI("\n");
+    }
     return suggWords;
 }
 
+void
+Dictionary::registerNextLetter(unsigned short c)
+{
+    if (c < mNextLettersSize) {
+        mNextLettersFrequencies[c]++;
+    }
+}
+
 unsigned short
 Dictionary::getChar(int *pos)
 {
@@ -210,6 +231,9 @@ Dictionary::getWordsRec(int pos, int depth, int maxDepth, bool completion, int s
             mWord[depth] = c;
             if (terminal) {
                 addWord(mWord, depth + 1, freq * snr);
+                if (depth >= mInputLength && mSkipPos < 0) {
+                    registerNextLetter(mWord[mInputLength]);
+                }
             }
             if (childrenAddress != 0) {
                 getWordsRec(childrenAddress, depth + 1, maxDepth,
diff --git a/dictionary/src/dictionary.h b/dictionary/src/dictionary.h
index a12c035c8..3749f3d88 100644
--- a/dictionary/src/dictionary.h
+++ b/dictionary/src/dictionary.h
@@ -32,7 +32,8 @@ class Dictionary {
 public:
     Dictionary(void *dict, int typedLetterMultipler, int fullWordMultiplier);
     int getSuggestions(int *codes, int codesSize, unsigned short *outWords, int *frequencies,
-        int maxWordLength, int maxWords, int maxAlternatives, int skipPos);
+            int maxWordLength, int maxWords, int maxAlternatives, int skipPos,
+            int *nextLetters, int nextLettersSize);
     bool isValidWord(unsigned short *word, int length);
     void setAsset(void *asset) { mAsset = asset; }
     void *getAsset() { return mAsset; }
@@ -53,6 +54,7 @@ private:
     void getWordsRec(int pos, int depth, int maxDepth, bool completion, int frequency,
             int inputIndex, int diffs);
     bool isValidWordRec(int pos, unsigned short *word, int offset, int length);
+    void registerNextLetter(unsigned short c);
 
     unsigned char *mDict;
     void *mAsset;
@@ -70,6 +72,8 @@ private:
 
     int mFullWordMultiplier;
     int mTypedLetterMultiplier;
+    int *mNextLettersFrequencies;
+    int mNextLettersSize;
 };
 
 // ----------------------------------------------------------------------------