Introduce SuggestionResults and use it for predictions.

Bug: 8187060
Bug: 13333066
Change-Id: I1ead897024508b7e40fbd93af7d14bfe74b93826

7 files changed, 231 insertions, 80 deletions

diff --git a/native/jni/src/suggest/core/dictionary/bigram_dictionary.cpp b/native/jni/src/suggest/core/dictionary/bigram_dictionary.cpp
index 0859df423..f793363a8 100644
--- a/native/jni/src/suggest/core/dictionary/bigram_dictionary.cpp
+++ b/native/jni/src/suggest/core/dictionary/bigram_dictionary.cpp
@@ -25,6 +25,7 @@
 #include "suggest/core/dictionary/binary_dictionary_bigrams_iterator.h"
 #include "suggest/core/dictionary/dictionary.h"
 #include "suggest/core/policy/dictionary_structure_with_buffer_policy.h"
+#include "suggest/core/result/suggestion_results.h"
 #include "utils/char_utils.h"
 
 namespace latinime {
@@ -40,71 +41,13 @@ BigramDictionary::BigramDictionary(
 BigramDictionary::~BigramDictionary() {
 }
 
-void BigramDictionary::addWordBigram(int *word, int length, int probability, int *bigramProbability,
-        int *bigramCodePoints, int *outputTypes) const {
-    if (length >= MAX_WORD_LENGTH) {
-        length = MAX_WORD_LENGTH - 1;
-    }
-    word[length] = 0;
-    if (DEBUG_DICT_FULL) {
-#ifdef FLAG_DBG
-        char s[length + 1];
-        for (int i = 0; i <= length; i++) s[i] = static_cast<char>(word[i]);
-        AKLOGI("Bigram: Found word = %s, freq = %d :", s, probability);
-#endif
-    }
-
-    // Find the right insertion point
-    int insertAt = 0;
-    while (insertAt < MAX_RESULTS) {
-        if (probability > bigramProbability[insertAt] || (bigramProbability[insertAt] == probability
-                && length < CharUtils::getCodePointCount(MAX_WORD_LENGTH,
-                        bigramCodePoints + insertAt * MAX_WORD_LENGTH))) {
-            break;
-        }
-        insertAt++;
-    }
-    if (DEBUG_DICT_FULL) {
-        AKLOGI("Bigram: InsertAt -> %d MAX_RESULTS: %d", insertAt, MAX_RESULTS);
-    }
-    if (insertAt >= MAX_RESULTS) {
-        return;
-    }
-    // Shift result buffers to insert the new entry.
-    memmove(bigramProbability + (insertAt + 1), bigramProbability + insertAt,
-            (MAX_RESULTS - insertAt - 1) * sizeof(bigramProbability[0]));
-    memmove(outputTypes + (insertAt + 1), outputTypes + insertAt,
-            (MAX_RESULTS - insertAt - 1) * sizeof(outputTypes[0]));
-    memmove(bigramCodePoints + (insertAt + 1) * MAX_WORD_LENGTH,
-            bigramCodePoints + insertAt * MAX_WORD_LENGTH,
-            (MAX_RESULTS - insertAt - 1) * sizeof(bigramCodePoints[0]) * MAX_WORD_LENGTH);
-    // Put the result.
-    bigramProbability[insertAt] = probability;
-    outputTypes[insertAt] = Dictionary::KIND_PREDICTION;
-    int *dest = bigramCodePoints + insertAt * MAX_WORD_LENGTH;
-    while (length--) {
-        *dest++ = *word++;
-    }
-    *dest = 0; // NULL terminate
-    if (DEBUG_DICT_FULL) {
-        AKLOGI("Bigram: Added word at %d", insertAt);
-    }
-}
-
 /* Parameters :
  * prevWord: the word before, the one for which we need to look up bigrams.
  * prevWordLength: its length.
- * outBigramCodePoints: an array for output, at the same format as outwords for getSuggestions.
- * outBigramProbability: an array to output frequencies.
- * outputTypes: an array to output types.
- * This method returns the number of bigrams this word has, for backward compatibility.
+ * outSuggestionResults: SuggestionResults to put the predictions.
  */
-int BigramDictionary::getPredictions(const int *prevWord, const int prevWordLength,
-        int *const outBigramCodePoints, int *const outBigramProbability,
-        int *const outputTypes) const {
-    // TODO: remove unused arguments, and refrain from storing stuff in members of this class
-    // TODO: have "in" arguments before "out" ones, and make out args explicit in the name
-
+void BigramDictionary::getPredictions(const int *prevWord, const int prevWordLength,
+        SuggestionResults *const outSuggestionResults) const {
     int pos = getBigramListPositionForWord(prevWord, prevWordLength,
             false /* forceLowerCaseSearch */);
     // getBigramListPositionForWord returns 0 if this word isn't in the dictionary or has no bigrams
@@ -114,11 +57,10 @@ int BigramDictionary::getPredictions(const int *prevWord, const int prevWordLeng
                 true /* forceLowerCaseSearch */);
     }
     // If still no bigrams, we really don't have them!
-    if (NOT_A_DICT_POS == pos) return 0;
+    if (NOT_A_DICT_POS == pos) return;
 
-    int bigramCount = 0;
     int unigramProbability = 0;
-    int bigramBuffer[MAX_WORD_LENGTH];
+    int bigramCodePoints[MAX_WORD_LENGTH];
     BinaryDictionaryBigramsIterator bigramsIt(
             mDictionaryStructurePolicy->getBigramsStructurePolicy(), pos);
     while (bigramsIt.hasNext()) {
@@ -128,7 +70,7 @@ int BigramDictionary::getPredictions(const int *prevWord, const int prevWordLeng
         }
         const int codePointCount = mDictionaryStructurePolicy->
                 getCodePointsAndProbabilityAndReturnCodePointCount(bigramsIt.getBigramPos(),
-                        MAX_WORD_LENGTH, bigramBuffer, &unigramProbability);
+                        MAX_WORD_LENGTH, bigramCodePoints, &unigramProbability);
         if (codePointCount <= 0) {
             continue;
         }
@@ -139,11 +81,8 @@ int BigramDictionary::getPredictions(const int *prevWord, const int prevWordLeng
         // here, but it can't get too bad.
         const int probability = mDictionaryStructurePolicy->getProbability(
                 unigramProbability, bigramsIt.getProbability());
-        addWordBigram(bigramBuffer, codePointCount, probability, outBigramProbability,
-                outBigramCodePoints, outputTypes);
-        ++bigramCount;
+        outSuggestionResults->addPrediction(bigramCodePoints, codePointCount, probability);
     }
-    return std::min(bigramCount, MAX_RESULTS);
 }
 
 // Returns a pointer to the start of the bigram list.
diff --git a/native/jni/src/suggest/core/dictionary/bigram_dictionary.h b/native/jni/src/suggest/core/dictionary/bigram_dictionary.h
index 8af7ee75d..12aaf20d3 100644
--- a/native/jni/src/suggest/core/dictionary/bigram_dictionary.h
+++ b/native/jni/src/suggest/core/dictionary/bigram_dictionary.h
@@ -22,21 +22,20 @@
 namespace latinime {
 
 class DictionaryStructureWithBufferPolicy;
+class SuggestionResults;
 
 class BigramDictionary {
  public:
     BigramDictionary(const DictionaryStructureWithBufferPolicy *const dictionaryStructurePolicy);
 
-    int getPredictions(const int *word, int length, int *outBigramCodePoints,
-            int *outBigramProbability, int *outputTypes) const;
+    void getPredictions(const int *word, int length,
+            SuggestionResults *const outSuggestionResults) const;
     int getBigramProbability(const int *word1, int length1, const int *word2, int length2) const;
     ~BigramDictionary();
 
  private:
     DISALLOW_IMPLICIT_CONSTRUCTORS(BigramDictionary);
 
-    void addWordBigram(int *word, int length, int probability, int *bigramProbability,
-            int *bigramCodePoints, int *outputTypes) const;
     int getBigramListPositionForWord(const int *prevWord, const int prevWordLength,
             const bool forceLowerCaseSearch) const;
 
diff --git a/native/jni/src/suggest/core/dictionary/dictionary.cpp b/native/jni/src/suggest/core/dictionary/dictionary.cpp
index 59a8a5500..ffa96e167 100644
--- a/native/jni/src/suggest/core/dictionary/dictionary.cpp
+++ b/native/jni/src/suggest/core/dictionary/dictionary.cpp
@@ -74,12 +74,11 @@ int Dictionary::getSuggestions(ProximityInfo *proximityInfo, DicTraverseSession
     }
 }
 
-int Dictionary::getBigrams(const int *word, int length, int *outWords, int *outputScores,
-        int *outputTypes) const {
+void Dictionary::getPredictions(const int *word, int length,
+        SuggestionResults *const outSuggestionResults) const {
     TimeKeeper::setCurrentTime();
-    if (length <= 0) return 0;
-    return mBigramDictionary->getPredictions(word, length, outWords, outputScores,
-            outputTypes);
+    if (length <= 0) return;
+    mBigramDictionary->getPredictions(word, length, outSuggestionResults);
 }
 
 int Dictionary::getProbability(const int *word, int length) const {
diff --git a/native/jni/src/suggest/core/dictionary/dictionary.h b/native/jni/src/suggest/core/dictionary/dictionary.h
index a7f19c9e6..2dea9fff8 100644
--- a/native/jni/src/suggest/core/dictionary/dictionary.h
+++ b/native/jni/src/suggest/core/dictionary/dictionary.h
@@ -33,6 +33,7 @@ namespace latinime {
 class DictionaryStructureWithBufferPolicy;
 class DicTraverseSession;
 class ProximityInfo;
+class SuggestionResults;
 class SuggestOptions;
 class WordProperty;
 
@@ -67,8 +68,8 @@ class Dictionary {
             const SuggestOptions *const suggestOptions, int *outWords, int *outputScores,
             int *spaceIndices, int *outputTypes, int *outputAutoCommitFirstWordConfidence) const;
 
-    int getBigrams(const int *word, int length, int *outWords, int *outputScores,
-            int *outputTypes) const;
+    void getPredictions(const int *word, int length,
+            SuggestionResults *const outSuggestionResults) const;
 
     int getProbability(const int *word, int length) const;
 
diff --git a/native/jni/src/suggest/core/result/suggested_word.h b/native/jni/src/suggest/core/result/suggested_word.h
new file mode 100644
index 000000000..48b29d6a6
--- /dev/null
+++ b/native/jni/src/suggest/core/result/suggested_word.h
@@ -0,0 +1,83 @@
+/*
+ * Copyright (C) 2014 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef LATINIME_SUGGESTED_WORD_H
+#define LATINIME_SUGGESTED_WORD_H
+
+#include <vector>
+
+#include "defines.h"
+#include "suggest/core/dictionary/dictionary.h"
+
+namespace latinime {
+
+class SuggestedWord {
+ public:
+    class Comparator {
+     public:
+        bool operator()(const SuggestedWord &left, const SuggestedWord &right) {
+            if (left.getScore() != right.getScore()) {
+                return left.getScore() < right.getScore();
+            }
+            return left.getCodePointCount() > right.getCodePointCount();
+        }
+
+     private:
+        DISALLOW_ASSIGNMENT_OPERATOR(Comparator);
+    };
+
+    SuggestedWord(const int *const codePoints, const int codePointCount,
+            const int score, const int type, const int indexToPartialCommit,
+            const int autoCommitFirstWordConfidence)
+            : mCodePoints(codePoints, codePoints + codePointCount), mScore(score),
+              mType(type), mIndexToPartialCommit(indexToPartialCommit),
+              mAutoCommitFirstWordConfidence(autoCommitFirstWordConfidence) {}
+
+    const int *getCodePoint() const {
+        return &mCodePoints.at(0);
+    }
+
+    int getCodePointCount() const {
+        return mCodePoints.size();
+    }
+
+    int getScore() const {
+        return mScore;
+    }
+
+    int getType() const {
+        return mType;
+    }
+
+    int getIndexToPartialCommit() const {
+        return mIndexToPartialCommit;
+    }
+
+    int getAutoCommitFirstWordConfidence() const {
+        return mAutoCommitFirstWordConfidence;
+    }
+
+ private:
+    DISALLOW_DEFAULT_CONSTRUCTOR(SuggestedWord);
+
+    std::vector<int> mCodePoints;
+    int mScore;
+    int mType;
+    int mIndexToPartialCommit;
+    int mAutoCommitFirstWordConfidence;
+};
+} // namespace latinime
+#endif /* LATINIME_SUGGESTED_WORD_H */
diff --git a/native/jni/src/suggest/core/result/suggestion_results.cpp b/native/jni/src/suggest/core/result/suggestion_results.cpp
new file mode 100644
index 000000000..2be757d83
--- /dev/null
+++ b/native/jni/src/suggest/core/result/suggestion_results.cpp
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2014 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "suggest/core/result/suggestion_results.h"
+
+namespace latinime {
+
+void SuggestionResults::outputSuggestions(JNIEnv *env, jintArray outSuggestionCount,
+        jintArray outputCodePointsArray, jintArray outScoresArray, jintArray outSpaceIndicesArray,
+        jintArray outTypesArray, jintArray outAutoCommitFirstWordConfidenceArray) {
+    int outputIndex = 0;
+    while (!mSuggestedWords.empty()) {
+        const SuggestedWord &suggestedWord = mSuggestedWords.top();
+        suggestedWord.getCodePointCount();
+        const int start = outputIndex * MAX_WORD_LENGTH;
+        env->SetIntArrayRegion(outputCodePointsArray, start, suggestedWord.getCodePointCount(),
+                suggestedWord.getCodePoint());
+        if (suggestedWord.getCodePointCount() < MAX_WORD_LENGTH) {
+            const int terminal = 0;
+            env->SetIntArrayRegion(outputCodePointsArray, start + suggestedWord.getCodePointCount(),
+                    1 /* len */, &terminal);
+        }
+        const int score = suggestedWord.getScore();
+        env->SetIntArrayRegion(outScoresArray, outputIndex, 1 /* len */, &score);
+        const int indexToPartialCommit = suggestedWord.getIndexToPartialCommit();
+        env->SetIntArrayRegion(outSpaceIndicesArray, outputIndex, 1 /* len */,
+                &indexToPartialCommit);
+        const int type = suggestedWord.getType();
+        env->SetIntArrayRegion(outTypesArray, outputIndex, 1 /* len */, &type);
+        if (mSuggestedWords.size() == 1) {
+            const int autoCommitFirstWordConfidence =
+                    suggestedWord.getAutoCommitFirstWordConfidence();
+            env->SetIntArrayRegion(outAutoCommitFirstWordConfidenceArray, 0 /* start */,
+                    1 /* len */, &autoCommitFirstWordConfidence);
+        }
+        ++outputIndex;
+        mSuggestedWords.pop();
+    }
+    env->SetIntArrayRegion(outSuggestionCount, 0 /* start */, 1 /* len */, &outputIndex);
+}
+
+void SuggestionResults::addPrediction(const int *const codePoints, const int codePointCount,
+        const int probability) {
+    if (codePointCount <= 0 || codePointCount > MAX_WORD_LENGTH
+            || probability == NOT_A_PROBABILITY) {
+        // Invalid word.
+        return;
+    }
+    // Use probability as a score of the word.
+    const int score = probability;
+    if (getSuggestionCount() >= mMaxSuggestionCount) {
+        const SuggestedWord &mWorstSuggestion = mSuggestedWords.top();
+        if (score > mWorstSuggestion.getScore() || (score == mWorstSuggestion.getScore()
+                && codePointCount < mWorstSuggestion.getCodePointCount())) {
+            mSuggestedWords.pop();
+        } else {
+            return;
+        }
+    }
+    mSuggestedWords.push(SuggestedWord(codePoints, codePointCount, score,
+            Dictionary::KIND_PREDICTION, NOT_AN_INDEX, NOT_A_FIRST_WORD_CONFIDENCE));
+}
+
+} // namespace latinime
diff --git a/native/jni/src/suggest/core/result/suggestion_results.h b/native/jni/src/suggest/core/result/suggestion_results.h
new file mode 100644
index 000000000..0b841ca19
--- /dev/null
+++ b/native/jni/src/suggest/core/result/suggestion_results.h
@@ -0,0 +1,53 @@
+/*
+ * Copyright (C) 2013 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef LATINIME_SUGGESTION_RESULTS_H
+#define LATINIME_SUGGESTION_RESULTS_H
+
+#include <queue>
+#include <vector>
+
+#include "defines.h"
+#include "jni.h"
+#include "suggest/core/result/suggested_word.h"
+
+namespace latinime {
+
+class SuggestionResults {
+ public:
+    explicit SuggestionResults(const int maxSuggestionCount)
+            : mMaxSuggestionCount(maxSuggestionCount), mSuggestedWords() {}
+
+    // Returns suggestion count.
+    void outputSuggestions(JNIEnv *env, jintArray outSuggestionCount, jintArray outCodePointsArray,
+            jintArray outScoresArray, jintArray outSpaceIndicesArray, jintArray outTypesArray,
+            jintArray outAutoCommitFirstWordConfidenceArray);
+
+    void addPrediction(const int *const codePoints, const int codePointCount, const int score);
+
+    int getSuggestionCount() const {
+        return mSuggestedWords.size();
+    }
+
+ private:
+    DISALLOW_IMPLICIT_CONSTRUCTORS(SuggestionResults);
+
+    const int mMaxSuggestionCount;
+    std::priority_queue<
+            SuggestedWord, std::vector<SuggestedWord>, SuggestedWord::Comparator> mSuggestedWords;
+};
+} // namespace latinime
+#endif // LATINIME_SUGGESTION_RESULTS_H