Add BinaryDictionary.getBigramProbabilityNative().

Bug: 10197478
Change-Id: Id16e3ee22a17679c46cc7a0199a9e28065177aa3

4 files changed, 12 insertions, 9 deletions

diff --git a/native/jni/src/suggest/core/dictionary/bigram_dictionary.cpp b/native/jni/src/suggest/core/dictionary/bigram_dictionary.cpp
index 425b07624..5ba71c168 100644
--- a/native/jni/src/suggest/core/dictionary/bigram_dictionary.cpp
+++ b/native/jni/src/suggest/core/dictionary/bigram_dictionary.cpp
@@ -150,24 +150,26 @@ int BigramDictionary::getBigramListPositionForWord(const int *prevWord, const in
     return mDictionaryStructurePolicy->getBigramsPositionOfNode(pos);
 }
 
-bool BigramDictionary::isValidBigram(const int *word0, int length0, const int *word1,
+int BigramDictionary::getBigramProbability(const int *word0, int length0, const int *word1,
         int length1) const {
     int pos = getBigramListPositionForWord(word0, length0, false /* forceLowerCaseSearch */);
     // getBigramListPositionForWord returns 0 if this word isn't in the dictionary or has no bigrams
-    if (NOT_A_DICT_POS == pos) return false;
+    if (NOT_A_DICT_POS == pos) return NOT_A_PROBABILITY;
     int nextWordPos = mDictionaryStructurePolicy->getTerminalNodePositionOfWord(word1, length1,
             false /* forceLowerCaseSearch */);
-    if (NOT_A_DICT_POS == nextWordPos) return false;
+    if (NOT_A_DICT_POS == nextWordPos) return NOT_A_PROBABILITY;
 
     BinaryDictionaryBigramsIterator bigramsIt(
             mDictionaryStructurePolicy->getBigramsStructurePolicy(), pos);
     while (bigramsIt.hasNext()) {
         bigramsIt.next();
         if (bigramsIt.getBigramPos() == nextWordPos) {
-            return true;
+            return mDictionaryStructurePolicy->getProbability(
+                    mDictionaryStructurePolicy->getUnigramProbabilityOfPtNode(nextWordPos),
+                    bigramsIt.getProbability());
         }
     }
-    return false;
+    return NOT_A_PROBABILITY;
 }
 
 // TODO: Move functions related to bigram to here
diff --git a/native/jni/src/suggest/core/dictionary/bigram_dictionary.h b/native/jni/src/suggest/core/dictionary/bigram_dictionary.h
index 99b964c49..8af7ee75d 100644
--- a/native/jni/src/suggest/core/dictionary/bigram_dictionary.h
+++ b/native/jni/src/suggest/core/dictionary/bigram_dictionary.h
@@ -29,7 +29,7 @@ class BigramDictionary {
 
     int getPredictions(const int *word, int length, int *outBigramCodePoints,
             int *outBigramProbability, int *outputTypes) const;
-    bool isValidBigram(const int *word1, int length1, const int *word2, int length2) const;
+    int getBigramProbability(const int *word1, int length1, const int *word2, int length2) const;
     ~BigramDictionary();
 
  private:
diff --git a/native/jni/src/suggest/core/dictionary/dictionary.cpp b/native/jni/src/suggest/core/dictionary/dictionary.cpp
index 033572201..ec1b63a12 100644
--- a/native/jni/src/suggest/core/dictionary/dictionary.cpp
+++ b/native/jni/src/suggest/core/dictionary/dictionary.cpp
@@ -93,8 +93,9 @@ int Dictionary::getProbability(const int *word, int length) const {
     return getDictionaryStructurePolicy()->getUnigramProbabilityOfPtNode(pos);
 }
 
-bool Dictionary::isValidBigram(const int *word0, int length0, const int *word1, int length1) const {
-    return mBigramDictionary->isValidBigram(word0, length0, word1, length1);
+int Dictionary::getBigramProbability(const int *word0, int length0, const int *word1,
+        int length1) const {
+    return mBigramDictionary->getBigramProbability(word0, length0, word1, length1);
 }
 
 void Dictionary::addUnigramWord(const int *const word, const int length, const int probability) {
diff --git a/native/jni/src/suggest/core/dictionary/dictionary.h b/native/jni/src/suggest/core/dictionary/dictionary.h
index 06e84bbfe..974447468 100644
--- a/native/jni/src/suggest/core/dictionary/dictionary.h
+++ b/native/jni/src/suggest/core/dictionary/dictionary.h
@@ -67,7 +67,7 @@ class Dictionary {
 
     int getProbability(const int *word, int length) const;
 
-    bool isValidBigram(const int *word0, int length0, const int *word1, int length1) const;
+    int getBigramProbability(const int *word0, int length0, const int *word1, int length1) const;
 
     void addUnigramWord(const int *const word, const int length, const int probability);