diff --git a/java/src/com/android/inputmethod/keyboard/KeyboardLayoutSet.java b/java/src/com/android/inputmethod/keyboard/KeyboardLayoutSet.java
index 870ac86..d6d0b21 100644
--- a/java/src/com/android/inputmethod/keyboard/KeyboardLayoutSet.java
+++ b/java/src/com/android/inputmethod/keyboard/KeyboardLayoutSet.java
@@ -113,7 +113,7 @@
         boolean mIsSpellChecker;
         int mKeyboardWidth;
         int mKeyboardHeight;
-        int mScriptId;
+        int mScriptId = ScriptUtils.SCRIPT_LATIN;
         // Sparse array of KeyboardLayoutSet element parameters indexed by element's id.
         final SparseArray<ElementParams> mKeyboardLayoutSetElementIdToParamsMap =
                 new SparseArray<>();
diff --git a/native/jni/src/suggest/core/dictionary/dictionary.cpp b/native/jni/src/suggest/core/dictionary/dictionary.cpp
index 92f5c17..d625739 100644
--- a/native/jni/src/suggest/core/dictionary/dictionary.cpp
+++ b/native/jni/src/suggest/core/dictionary/dictionary.cpp
@@ -92,7 +92,11 @@
     TimeKeeper::setCurrentTime();
     NgramListenerForPrediction listener(prevWordsInfo, outSuggestionResults,
             mDictionaryStructureWithBufferPolicy.get());
-    mDictionaryStructureWithBufferPolicy->iterateNgramEntries(prevWordsInfo, &listener);
+    int prevWordsPtNodePos[MAX_PREV_WORD_COUNT_FOR_N_GRAM];
+    prevWordsInfo->getPrevWordsTerminalPtNodePos(
+            mDictionaryStructureWithBufferPolicy.get(), prevWordsPtNodePos,
+            true /* tryLowerCaseSearch */);
+    mDictionaryStructureWithBufferPolicy->iterateNgramEntries(prevWordsPtNodePos, &listener);
 }
 
 int Dictionary::getProbability(const int *word, int length) const {
@@ -111,7 +115,15 @@
     int nextWordPos = mDictionaryStructureWithBufferPolicy->getTerminalPtNodePositionOfWord(word,
             length, false /* forceLowerCaseSearch */);
     if (NOT_A_DICT_POS == nextWordPos) return NOT_A_PROBABILITY;
-    return getDictionaryStructurePolicy()->getProbabilityOfPtNode(prevWordsInfo, nextWordPos);
+    if (!prevWordsInfo) {
+        return getDictionaryStructurePolicy()->getProbabilityOfPtNode(
+                nullptr /* prevWordsPtNodePos */, nextWordPos);
+    }
+    int prevWordsPtNodePos[MAX_PREV_WORD_COUNT_FOR_N_GRAM];
+    prevWordsInfo->getPrevWordsTerminalPtNodePos(
+            mDictionaryStructureWithBufferPolicy.get(), prevWordsPtNodePos,
+            true /* tryLowerCaseSearch */);
+    return getDictionaryStructurePolicy()->getProbabilityOfPtNode(prevWordsPtNodePos, nextWordPos);
 }
 
 bool Dictionary::addUnigramEntry(const int *const word, const int length,
diff --git a/native/jni/src/suggest/core/policy/dictionary_structure_with_buffer_policy.h b/native/jni/src/suggest/core/policy/dictionary_structure_with_buffer_policy.h
index 81e38f7..7e3bf3f 100644
--- a/native/jni/src/suggest/core/policy/dictionary_structure_with_buffer_policy.h
+++ b/native/jni/src/suggest/core/policy/dictionary_structure_with_buffer_policy.h
@@ -59,10 +59,10 @@
     virtual int getProbability(const int unigramProbability,
             const int bigramProbability) const = 0;
 
-    virtual int getProbabilityOfPtNode(const PrevWordsInfo *const prevWordsInfo,
+    virtual int getProbabilityOfPtNode(const int *const prevWordsPtNodePos,
             const int nodePos) const = 0;
 
-    virtual void iterateNgramEntries(const PrevWordsInfo *const prevWordsInfo,
+    virtual void iterateNgramEntries(const int *const prevWordsPtNodePos,
             NgramListener *const listener) const = 0;
 
     virtual int getShortcutPositionOfPtNode(const int nodePos) const = 0;
diff --git a/native/jni/src/suggest/core/session/prev_words_info.h b/native/jni/src/suggest/core/session/prev_words_info.h
index 76276f5..e44e876 100644
--- a/native/jni/src/suggest/core/session/prev_words_info.h
+++ b/native/jni/src/suggest/core/session/prev_words_info.h
@@ -90,13 +90,6 @@
         }
     }
 
-    BinaryDictionaryBigramsIterator getBigramsIteratorForPrediction(
-            const DictionaryStructureWithBufferPolicy *const dictStructurePolicy) const {
-        return getBigramsIteratorForWordWithTryingLowerCaseSearch(
-                dictStructurePolicy, mPrevWordCodePoints[0], mPrevWordCodePointCount[0],
-                mIsBeginningOfSentence[0]);
-    }
-
     // n is 1-indexed.
     const int *getNthPrevWordCodePoints(const int n) const {
         if (n <= 0 || n > MAX_PREV_WORD_COUNT_FOR_N_GRAM) {
@@ -154,46 +147,6 @@
                 codePoints, codePointCount, true /* forceLowerCaseSearch */);
     }
 
-    static BinaryDictionaryBigramsIterator getBigramsIteratorForWordWithTryingLowerCaseSearch(
-            const DictionaryStructureWithBufferPolicy *const dictStructurePolicy,
-            const int *const wordCodePoints, const int wordCodePointCount,
-            const bool isBeginningOfSentence) {
-        if (!dictStructurePolicy || !wordCodePoints || wordCodePointCount > MAX_WORD_LENGTH) {
-            return BinaryDictionaryBigramsIterator();
-        }
-        int codePoints[MAX_WORD_LENGTH];
-        int codePointCount = wordCodePointCount;
-        memmove(codePoints, wordCodePoints, sizeof(int) * codePointCount);
-        if (isBeginningOfSentence) {
-            codePointCount = CharUtils::attachBeginningOfSentenceMarker(codePoints,
-                    codePointCount, MAX_WORD_LENGTH);
-            if (codePointCount <= 0) {
-                return BinaryDictionaryBigramsIterator();
-            }
-        }
-        BinaryDictionaryBigramsIterator bigramsIt = getBigramsIteratorForWord(dictStructurePolicy,
-                codePoints, codePointCount, false /* forceLowerCaseSearch */);
-        // getBigramsIteratorForWord returns an empty iterator if this word isn't in the dictionary
-        // or has no bigrams.
-        if (bigramsIt.hasNext()) {
-            return bigramsIt;
-        }
-        // If no bigrams for this exact word, search again in lower case.
-        return getBigramsIteratorForWord(dictStructurePolicy, codePoints,
-                codePointCount, true /* forceLowerCaseSearch */);
-    }
-
-    static BinaryDictionaryBigramsIterator getBigramsIteratorForWord(
-            const DictionaryStructureWithBufferPolicy *const dictStructurePolicy,
-            const int *wordCodePoints, const int wordCodePointCount,
-            const bool forceLowerCaseSearch) {
-        if (!wordCodePoints || wordCodePointCount <= 0) return BinaryDictionaryBigramsIterator();
-        const int terminalPtNodePos = dictStructurePolicy->getTerminalPtNodePositionOfWord(
-                wordCodePoints, wordCodePointCount, forceLowerCaseSearch);
-        if (NOT_A_DICT_POS == terminalPtNodePos) return BinaryDictionaryBigramsIterator();
-        return dictStructurePolicy->getBigramsIteratorOfPtNode(terminalPtNodePos);
-    }
-
     void clear() {
         for (size_t i = 0; i < NELEMS(mPrevWordCodePoints); ++i) {
             mPrevWordCodePointCount[i] = 0;
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.cpp
index 4b834a0..994c425 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.cpp
@@ -132,7 +132,7 @@
     }
 }
 
-int Ver4PatriciaTriePolicy::getProbabilityOfPtNode(const PrevWordsInfo *const prevWordsInfo,
+int Ver4PatriciaTriePolicy::getProbabilityOfPtNode(const int *const prevWordsPtNodePos,
         const int ptNodePos) const {
     if (ptNodePos == NOT_A_DICT_POS) {
         return NOT_A_PROBABILITY;
@@ -141,9 +141,9 @@
     if (ptNodeParams.isDeleted() || ptNodeParams.isBlacklisted() || ptNodeParams.isNotAWord()) {
         return NOT_A_PROBABILITY;
     }
-    if (prevWordsInfo) {
+    if (prevWordsPtNodePos) {
         BinaryDictionaryBigramsIterator bigramsIt =
-                prevWordsInfo->getBigramsIteratorForPrediction(this /* dictStructurePolicy */);
+                getBigramsIteratorOfPtNode(prevWordsPtNodePos[0]);
         while (bigramsIt.hasNext()) {
             bigramsIt.next();
             if (bigramsIt.getBigramPos() == ptNodePos
@@ -156,10 +156,12 @@
     return getProbability(ptNodeParams.getProbability(), NOT_A_PROBABILITY);
 }
 
-void Ver4PatriciaTriePolicy::iterateNgramEntries(const PrevWordsInfo *const prevWordsInfo,
+void Ver4PatriciaTriePolicy::iterateNgramEntries(const int *const prevWordsPtNodePos,
         NgramListener *const listener) const {
-    BinaryDictionaryBigramsIterator bigramsIt = prevWordsInfo->getBigramsIteratorForPrediction(
-            this /* dictStructurePolicy */);
+    if (!prevWordsPtNodePos) {
+        return;
+    }
+    BinaryDictionaryBigramsIterator bigramsIt = getBigramsIteratorOfPtNode(prevWordsPtNodePos[0]);
     while (bigramsIt.hasNext()) {
         bigramsIt.next();
         listener->onVisitEntry(bigramsIt.getProbability(), bigramsIt.getBigramPos());
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.h b/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.h
index e61c060..ff69de7 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.h
@@ -90,10 +90,9 @@
 
     int getProbability(const int unigramProbability, const int bigramProbability) const;
 
-    int getProbabilityOfPtNode(const PrevWordsInfo *const prevWordsInfo,
-            const int ptNodePos) const;
+    int getProbabilityOfPtNode(const int *const prevWordsPtNodePos, const int ptNodePos) const;
 
-    void iterateNgramEntries(const PrevWordsInfo *const prevWordsInfo,
+    void iterateNgramEntries(const int *const prevWordsPtNodePos,
             NgramListener *const listener) const;
 
     int getShortcutPositionOfPtNode(const int ptNodePos) const;
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.cpp
index 6f02ff3..53415ae 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.cpp
@@ -297,7 +297,7 @@
     }
 }
 
-int PatriciaTriePolicy::getProbabilityOfPtNode(const PrevWordsInfo *const prevWordsInfo,
+int PatriciaTriePolicy::getProbabilityOfPtNode(const int *const prevWordsPtNodePos,
         const int ptNodePos) const {
     if (ptNodePos == NOT_A_DICT_POS) {
         return NOT_A_PROBABILITY;
@@ -310,9 +310,9 @@
         // for shortcuts).
         return NOT_A_PROBABILITY;
     }
-    if (prevWordsInfo) {
+    if (prevWordsPtNodePos) {
         BinaryDictionaryBigramsIterator bigramsIt =
-                prevWordsInfo->getBigramsIteratorForPrediction(this /* dictStructurePolicy */);
+                getBigramsIteratorOfPtNode(prevWordsPtNodePos[0]);
         while (bigramsIt.hasNext()) {
             bigramsIt.next();
             if (bigramsIt.getBigramPos() == ptNodePos
@@ -325,10 +325,12 @@
     return getProbability(ptNodeParams.getProbability(), NOT_A_PROBABILITY);
 }
 
-void PatriciaTriePolicy::iterateNgramEntries(const PrevWordsInfo *const prevWordsInfo,
+void PatriciaTriePolicy::iterateNgramEntries(const int *const prevWordsPtNodePos,
         NgramListener *const listener) const {
-    BinaryDictionaryBigramsIterator bigramsIt = prevWordsInfo->getBigramsIteratorForPrediction(
-            this /* dictStructurePolicy */);
+    if (!prevWordsPtNodePos) {
+        return;
+    }
+    BinaryDictionaryBigramsIterator bigramsIt = getBigramsIteratorOfPtNode(prevWordsPtNodePos[0]);
     while (bigramsIt.hasNext()) {
         bigramsIt.next();
         listener->onVisitEntry(bigramsIt.getProbability(), bigramsIt.getBigramPos());
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.h b/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.h
index a3b2220..07cb72b 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.h
@@ -63,9 +63,9 @@
 
     int getProbability(const int unigramProbability, const int bigramProbability) const;
 
-    int getProbabilityOfPtNode(const PrevWordsInfo *const prevWordsInfo, const int ptNodePos) const;
+    int getProbabilityOfPtNode(const int *const prevWordsPtNodePos, const int ptNodePos) const;
 
-    void iterateNgramEntries(const PrevWordsInfo *const prevWordsInfo,
+    void iterateNgramEntries(const int *const prevWordsPtNodePos,
             NgramListener *const listener) const;
 
     int getShortcutPositionOfPtNode(const int ptNodePos) const;
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 23bbbbd..22f7e11 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
@@ -122,7 +122,7 @@
     }
 }
 
-int Ver4PatriciaTriePolicy::getProbabilityOfPtNode(const PrevWordsInfo *const prevWordsInfo,
+int Ver4PatriciaTriePolicy::getProbabilityOfPtNode(const int *const prevWordsPtNodePos,
         const int ptNodePos) const {
     if (ptNodePos == NOT_A_DICT_POS) {
         return NOT_A_PROBABILITY;
@@ -131,9 +131,9 @@
     if (ptNodeParams.isDeleted() || ptNodeParams.isBlacklisted() || ptNodeParams.isNotAWord()) {
         return NOT_A_PROBABILITY;
     }
-    if (prevWordsInfo) {
+    if (prevWordsPtNodePos) {
         BinaryDictionaryBigramsIterator bigramsIt =
-                prevWordsInfo->getBigramsIteratorForPrediction(this /* dictStructurePolicy */);
+                getBigramsIteratorOfPtNode(prevWordsPtNodePos[0]);
         while (bigramsIt.hasNext()) {
             bigramsIt.next();
             if (bigramsIt.getBigramPos() == ptNodePos
@@ -146,10 +146,12 @@
     return getProbability(ptNodeParams.getProbability(), NOT_A_PROBABILITY);
 }
 
-void Ver4PatriciaTriePolicy::iterateNgramEntries(const PrevWordsInfo *const prevWordsInfo,
+void Ver4PatriciaTriePolicy::iterateNgramEntries(const int *const prevWordsPtNodePos,
         NgramListener *const listener) const {
-    BinaryDictionaryBigramsIterator bigramsIt = prevWordsInfo->getBigramsIteratorForPrediction(
-            this /* dictStructurePolicy */);
+    if (!prevWordsPtNodePos) {
+        return;
+    }
+    BinaryDictionaryBigramsIterator bigramsIt = getBigramsIteratorOfPtNode(prevWordsPtNodePos[0]);
     while (bigramsIt.hasNext()) {
         bigramsIt.next();
         listener->onVisitEntry(bigramsIt.getProbability(), bigramsIt.getBigramPos());
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.h b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.h
index 1838454..c5b6a80 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.h
@@ -72,9 +72,9 @@
 
     int getProbability(const int unigramProbability, const int bigramProbability) const;
 
-    int getProbabilityOfPtNode(const PrevWordsInfo *const prevWordsInfo, const int ptNodePos) const;
+    int getProbabilityOfPtNode(const int *const prevWordsPtNodePos, const int ptNodePos) const;
 
-    void iterateNgramEntries(const PrevWordsInfo *const prevWordsInfo,
+    void iterateNgramEntries(const int *const prevWordsPtNodePos,
             NgramListener *const listener) const;
 
     int getShortcutPositionOfPtNode(const int ptNodePos) const;
diff --git a/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.cpp b/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.cpp
index a7d86f9..c700476 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.cpp
@@ -98,6 +98,43 @@
     return putInternal(unsignedKey, value, getBitShuffledKey(unsignedKey), bitmapEntryIndex,
             readEntry(bitmapEntryIndex), 0 /* level */);
 }
+/**
+ * Iterate next entry in a certain level.
+ *
+ * @param iterationState the iteration state that will be read and updated in this method.
+ * @param outKey the output key
+ * @return Result instance. mIsValid is false when all entries are iterated.
+ */
+const TrieMap::Result TrieMap::iterateNext(std::vector<TableIterationState> *const iterationState,
+        int *const outKey) const {
+    while (!iterationState->empty()) {
+        TableIterationState &state = iterationState->back();
+        if (state.mTableSize <= state.mCurrentIndex) {
+            // Move to parent.
+            iterationState->pop_back();
+        } else {
+            const int entryIndex = state.mTableIndex + state.mCurrentIndex;
+            state.mCurrentIndex += 1;
+            const Entry entry = readEntry(entryIndex);
+            if (entry.isBitmapEntry()) {
+                // Move to child.
+                iterationState->emplace_back(popCount(entry.getBitmap()), entry.getTableIndex());
+            } else {
+                if (outKey) {
+                    *outKey = entry.getKey();
+                }
+                if (!entry.hasTerminalLink()) {
+                    return Result(entry.getValue(), true, INVALID_INDEX);
+                }
+                const int valueEntryIndex = entry.getValueEntryIndex();
+                const Entry valueEntry = readEntry(valueEntryIndex);
+                return Result(valueEntry.getValueOfValueEntry(), true, valueEntryIndex + 1);
+            }
+        }
+    }
+    // Visited all entries.
+    return Result(0, false, INVALID_INDEX);
+}
 
 /**
  * Shuffle bits of the key in the fixed order.
diff --git a/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.h b/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.h
index 2a9051f..b5bcc3b 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/utils/trie_map.h
@@ -44,6 +44,117 @@
                   mNextLevelBitmapEntryIndex(nextLevelBitmapEntryIndex) {}
     };
 
+    /**
+     * Struct to record iteration state in a table.
+     */
+    struct TableIterationState {
+        int mTableSize;
+        int mTableIndex;
+        int mCurrentIndex;
+
+        TableIterationState(const int tableSize, const int tableIndex)
+                : mTableSize(tableSize), mTableIndex(tableIndex), mCurrentIndex(0) {}
+    };
+
+    class TrieMapRange;
+    class TrieMapIterator {
+     public:
+        class IterationResult {
+         public:
+            IterationResult(const TrieMap *const trieMap, const int key, const uint64_t value,
+                    const int nextLeveBitmapEntryIndex)
+                    : mTrieMap(trieMap), mKey(key), mValue(value),
+                      mNextLevelBitmapEntryIndex(nextLeveBitmapEntryIndex) {}
+
+            const TrieMapRange getEntriesInNextLevel() const {
+                return TrieMapRange(mTrieMap, mNextLevelBitmapEntryIndex);
+            }
+
+            bool hasNextLevelMap() const {
+                return mNextLevelBitmapEntryIndex != INVALID_INDEX;
+            }
+
+            AK_FORCE_INLINE int key() const {
+                return mKey;
+            }
+
+            AK_FORCE_INLINE uint64_t value() const {
+                return mValue;
+            }
+
+         private:
+            const TrieMap *const mTrieMap;
+            const int mKey;
+            const uint64_t mValue;
+            const int mNextLevelBitmapEntryIndex;
+        };
+
+        TrieMapIterator(const TrieMap *const trieMap, const int bitmapEntryIndex)
+                : mTrieMap(trieMap), mStateStack(), mBaseBitmapEntryIndex(bitmapEntryIndex),
+                  mKey(0), mValue(0), mIsValid(false), mNextLevelBitmapEntryIndex(INVALID_INDEX) {
+            if (!trieMap) {
+                return;
+            }
+            const Entry bitmapEntry = mTrieMap->readEntry(mBaseBitmapEntryIndex);
+            mStateStack.emplace_back(
+                    mTrieMap->popCount(bitmapEntry.getBitmap()), bitmapEntry.getTableIndex());
+            this->operator++();
+        }
+
+        const IterationResult operator*() const {
+            return IterationResult(mTrieMap, mKey, mValue, mNextLevelBitmapEntryIndex);
+        }
+
+        bool operator!=(const TrieMapIterator &other) const {
+            // Caveat: This works only for for loops.
+            return mIsValid || other.mIsValid;
+        }
+
+        const TrieMapIterator &operator++() {
+            const Result result = mTrieMap->iterateNext(&mStateStack, &mKey);
+            mValue = result.mValue;
+            mIsValid = result.mIsValid;
+            mNextLevelBitmapEntryIndex = result.mNextLevelBitmapEntryIndex;
+            return *this;
+        }
+
+     private:
+        DISALLOW_DEFAULT_CONSTRUCTOR(TrieMapIterator);
+        DISALLOW_ASSIGNMENT_OPERATOR(TrieMapIterator);
+
+        const TrieMap *const mTrieMap;
+        std::vector<TrieMap::TableIterationState> mStateStack;
+        const int mBaseBitmapEntryIndex;
+        int mKey;
+        uint64_t mValue;
+        bool mIsValid;
+        int mNextLevelBitmapEntryIndex;
+    };
+
+    /**
+     * Class to support iterating entries in TrieMap by range base for loops.
+     */
+    class TrieMapRange {
+     public:
+        TrieMapRange(const TrieMap *const trieMap, const int bitmapEntryIndex)
+                : mTrieMap(trieMap), mBaseBitmapEntryIndex(bitmapEntryIndex) {};
+
+        TrieMapIterator begin() const {
+            return TrieMapIterator(mTrieMap, mBaseBitmapEntryIndex);
+        }
+
+        const TrieMapIterator end() const {
+            return TrieMapIterator(nullptr, INVALID_INDEX);
+        }
+
+     private:
+        DISALLOW_DEFAULT_CONSTRUCTOR(TrieMapRange);
+        DISALLOW_ASSIGNMENT_OPERATOR(TrieMapRange);
+
+        const TrieMap *const mTrieMap;
+        const int mBaseBitmapEntryIndex;
+    };
+
     static const int INVALID_INDEX;
     static const uint64_t MAX_VALUE;
 
@@ -73,6 +184,14 @@
 
     bool put(const int key, const uint64_t value, const int bitmapEntryIndex);
 
+    const TrieMapRange getEntriesInRootLevel() const {
+        return getEntriesInSpecifiedLevel(ROOT_BITMAP_ENTRY_INDEX);
+    }
+
+    const TrieMapRange getEntriesInSpecifiedLevel(const int bitmapEntryIndex) const {
+        return TrieMapRange(this, bitmapEntryIndex);
+    }
+
  private:
     DISALLOW_COPY_AND_ASSIGN(TrieMap);
 
@@ -171,6 +290,8 @@
     bool addNewEntryByExpandingTable(const uint32_t key, const uint64_t value,
             const int tableIndex, const uint32_t bitmap, const int bitmapEntryIndex,
             const int label);
+    const Result iterateNext(std::vector<TableIterationState> *const iterationState,
+            int *const outKey) const;
 
     AK_FORCE_INLINE const Entry readEntry(const int entryIndex) const {
         return Entry(readField0(entryIndex), readField1(entryIndex));
diff --git a/native/jni/tests/suggest/policyimpl/dictionary/utils/trie_map_test.cpp b/native/jni/tests/suggest/policyimpl/dictionary/utils/trie_map_test.cpp
index 5dd7822..df778b6 100644
--- a/native/jni/tests/suggest/policyimpl/dictionary/utils/trie_map_test.cpp
+++ b/native/jni/tests/suggest/policyimpl/dictionary/utils/trie_map_test.cpp
@@ -54,7 +54,7 @@
         EXPECT_TRUE(trieMap.putRoot(i, i));
     }
     for (int i = 0; i < ELEMENT_COUNT; ++i) {
-        EXPECT_EQ(trieMap.getRoot(i).mValue, static_cast<uint64_t>(i));
+        EXPECT_EQ(static_cast<uint64_t>(i), trieMap.getRoot(i).mValue);
     }
 }
 
@@ -78,7 +78,7 @@
         testKeyValuePairs[key] = value;
     }
     for (const auto &v : testKeyValuePairs) {
-        EXPECT_EQ(trieMap.getRoot(v.first).mValue, v.second);
+        EXPECT_EQ(v.second, trieMap.getRoot(v.first).mValue);
     }
 }
 
@@ -163,6 +163,61 @@
             }
         }
     }
+
+    // Iteration
+    for (const auto &firstLevelEntry : trieMap.getEntriesInRootLevel()) {
+        EXPECT_EQ(trieMap.getRoot(firstLevelEntry.key()).mValue, firstLevelEntry.value());
+        EXPECT_EQ(firstLevelEntries[firstLevelEntry.key()], firstLevelEntry.value());
+        firstLevelEntries.erase(firstLevelEntry.key());
+        for (const auto &secondLevelEntry : firstLevelEntry.getEntriesInNextLevel()) {
+            EXPECT_EQ(twoLevelMap[firstLevelEntry.key()][secondLevelEntry.key()],
+                    secondLevelEntry.value());
+            twoLevelMap[firstLevelEntry.key()].erase(secondLevelEntry.key());
+            for (const auto &thirdLevelEntry : secondLevelEntry.getEntriesInNextLevel()) {
+                EXPECT_EQ(threeLevelMap[firstLevelEntry.key()][secondLevelEntry.key()]
+                        [thirdLevelEntry.key()], thirdLevelEntry.value());
+                threeLevelMap[firstLevelEntry.key()][secondLevelEntry.key()].erase(
+                        thirdLevelEntry.key());
+            }
+        }
+    }
+
+    // Ensure all entries have been traversed.
+    EXPECT_TRUE(firstLevelEntries.empty());
+    for (const auto &secondLevelEntry : twoLevelMap) {
+        EXPECT_TRUE(secondLevelEntry.second.empty());
+    }
+    for (const auto &secondLevelEntry : threeLevelMap) {
+        for (const auto &thirdLevelEntry : secondLevelEntry.second) {
+            EXPECT_TRUE(thirdLevelEntry.second.empty());
+        }
+    }
+}
+
+TEST(TrieMapTest, TestIteration) {
+    static const int ELEMENT_COUNT = 200000;
+    TrieMap trieMap;
+    std::unordered_map<int, uint64_t> testKeyValuePairs;
+
+    // Use the uniform integer distribution [S_INT_MIN, S_INT_MAX].
+    std::uniform_int_distribution<int> keyDistribution(S_INT_MIN, S_INT_MAX);
+    auto keyRandomNumberGenerator = std::bind(keyDistribution, std::mt19937());
+
+    // Use the uniform distribution [0, TrieMap::MAX_VALUE].
+    std::uniform_int_distribution<uint64_t> valueDistribution(0, TrieMap::MAX_VALUE);
+    auto valueRandomNumberGenerator = std::bind(valueDistribution, std::mt19937());
+    for (int i = 0; i < ELEMENT_COUNT; ++i) {
+        const int key = keyRandomNumberGenerator();
+        const uint64_t value = valueRandomNumberGenerator();
+        EXPECT_TRUE(trieMap.putRoot(key, value));
+        testKeyValuePairs[key] = value;
+    }
+    for (const auto &entry : trieMap.getEntriesInRootLevel()) {
+        EXPECT_EQ(trieMap.getRoot(entry.key()).mValue, entry.value());
+        EXPECT_EQ(testKeyValuePairs[entry.key()], entry.value());
+        testKeyValuePairs.erase(entry.key());
+    }
+    EXPECT_TRUE(testKeyValuePairs.empty());
 }
 
 }  // namespace