Merge "Use LanguageModelDictContent in iterateNgramEntries()."

1 files changed, 15 insertions, 5 deletions

diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp
index 2ea248e86..461397b2c 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp
@@ -158,11 +158,21 @@ void Ver4PatriciaTriePolicy::iterateNgramEntries(const int *const prevWordsPtNod
     if (!prevWordsPtNodePos) {
         return;
     }
-    const int bigramsPosition = getBigramsPositionOfPtNode(prevWordsPtNodePos[0]);
-    BinaryDictionaryBigramsIterator bigramsIt(&mBigramPolicy, bigramsPosition);
-    while (bigramsIt.hasNext()) {
-        bigramsIt.next();
-        listener->onVisitEntry(bigramsIt.getProbability(), bigramsIt.getBigramPos());
+    // TODO: Support n-gram.
+    const PtNodeParams ptNodeParams =
+            mNodeReader.fetchPtNodeParamsInBufferFromPtNodePos(prevWordsPtNodePos[0]);
+    const int prevWordId = ptNodeParams.getTerminalId();
+    const WordIdArrayView prevWordIds = WordIdArrayView::fromObject(&prevWordId);
+    const auto languageModelDictContent = mBuffers->getLanguageModelDictContent();
+    for (const auto entry : languageModelDictContent->getProbabilityEntries(prevWordIds)) {
+        const ProbabilityEntry &probabilityEntry = entry.getProbabilityEntry();
+        const int probability = probabilityEntry.hasHistoricalInfo() ?
+                ForgettingCurveUtils::decodeProbability(
+                        probabilityEntry.getHistoricalInfo(), mHeaderPolicy) :
+                probabilityEntry.getProbability();
+        const int ptNodePos = mBuffers->getTerminalPositionLookupTable()->getTerminalPtNodePosition(
+                entry.getWordId());
+        listener->onVisitEntry(probability, ptNodePos);
     }
 }