diff --git a/native/jni/src/unigram_dictionary.cpp b/native/jni/src/unigram_dictionary.cpp
index b6b0210..9f7ab53 100644
--- a/native/jni/src/unigram_dictionary.cpp
+++ b/native/jni/src/unigram_dictionary.cpp
@@ -63,8 +63,8 @@
 
 // TODO: This needs to take a const unsigned short* and not tinker with its contents
 static inline void addWord(
-        unsigned short *word, int length, int frequency, WordsPriorityQueue *queue) {
-    queue->push(frequency, word, length);
+        unsigned short *word, int length, int frequency, WordsPriorityQueue *queue, int type) {
+    queue->push(frequency, word, length, type);
 }
 
 // Return the replacement code point for a digraph, or 0 if none.
@@ -213,8 +213,8 @@
         AKLOGI("Max normalized score = %f", ns);
     }
     const int suggestedWordsCount =
-            queuePool.getMasterQueue()->outputSuggestions(
-                    masterCorrection.getPrimaryInputWord(), codesSize, frequencies, outWords);
+            queuePool.getMasterQueue()->outputSuggestions(masterCorrection.getPrimaryInputWord(),
+                    codesSize, frequencies, outWords, outputTypes);
 
     if (DEBUG_DICT) {
         float ns = queuePool.getMasterQueue()->getHighestNormalizedScore(
@@ -391,7 +391,8 @@
         const int finalProbability =
                 correction->getFinalProbability(probability, &wordPointer, &wordLength);
         if (finalProbability != NOT_A_PROBABILITY) {
-            addWord(wordPointer, wordLength, finalProbability, masterQueue);
+            addWord(wordPointer, wordLength, finalProbability, masterQueue,
+                    Dictionary::KIND_CORRECTION);
 
             const int shortcutProbability = finalProbability > 0 ? finalProbability - 1 : 0;
             // Please note that the shortcut candidates will be added to the master queue only.
@@ -409,7 +410,7 @@
                 const int shortcutTargetStringLength = iterator.getNextShortcutTarget(
                         MAX_WORD_LENGTH_INTERNAL, shortcutTarget);
                 addWord(shortcutTarget, shortcutTargetStringLength, shortcutProbability,
-                        masterQueue);
+                        masterQueue, Dictionary::KIND_CORRECTION);
             }
         }
     }
@@ -424,7 +425,7 @@
         }
         const int finalProbability = correction->getFinalProbabilityForSubQueue(
                 probability, &wordPointer, &wordLength, inputIndex);
-        addWord(wordPointer, wordLength, finalProbability, subQueue);
+        addWord(wordPointer, wordLength, finalProbability, subQueue, Dictionary::KIND_CORRECTION);
     }
 }
 
@@ -572,7 +573,8 @@
             AKLOGI("Split two words: freq = %d, length = %d, %d, isSpace ? %d", pairFreq,
                     inputLength, tempOutputWordLength, isSpaceProximity);
         }
-        addWord(outputWord, tempOutputWordLength, pairFreq, queuePool->getMasterQueue());
+        addWord(outputWord, tempOutputWordLength, pairFreq, queuePool->getMasterQueue(),
+                Dictionary::KIND_CORRECTION);
     }
     return FLAG_MULTIPLE_SUGGEST_CONTINUE;
 }
diff --git a/native/jni/src/words_priority_queue.h b/native/jni/src/words_priority_queue.h
index c0dedb5..8a6da1c 100644
--- a/native/jni/src/words_priority_queue.h
+++ b/native/jni/src/words_priority_queue.h
@@ -33,12 +33,14 @@
         unsigned short mWord[MAX_WORD_LENGTH_INTERNAL];
         int mWordLength;
         bool mUsed;
+        int mType;
 
-        void setParams(int score, unsigned short *word, int wordLength) {
+        void setParams(int score, unsigned short *word, int wordLength, int type) {
             mScore = score;
             mWordLength = wordLength;
             memcpy(mWord, word, sizeof(unsigned short) * wordLength);
             mUsed = true;
+            mType = type;
         }
     };
 
@@ -56,7 +58,7 @@
         delete[] mSuggestedWords;
     }
 
-    void push(int score, unsigned short *word, int wordLength) {
+    void push(int score, unsigned short *word, int wordLength, int type) {
         SuggestedWord *sw = 0;
         if (mSuggestions.size() >= MAX_WORDS) {
             sw = mSuggestions.top();
@@ -69,9 +71,9 @@
             }
         }
         if (sw == 0) {
-            sw = getFreeSuggestedWord(score, word, wordLength);
+            sw = getFreeSuggestedWord(score, word, wordLength, type);
         } else {
-            sw->setParams(score, word, wordLength);
+            sw->setParams(score, word, wordLength, type);
         }
         if (sw == 0) {
             AKLOGE("SuggestedWord is accidentally null.");
@@ -94,7 +96,7 @@
     }
 
     int outputSuggestions(const unsigned short *before, const int beforeLength,
-            int *frequencies, unsigned short *outputChars) {
+            int *frequencies, unsigned short *outputChars, int* outputTypes) {
         mHighestSuggestedWord = 0;
         const unsigned int size = min(
               MAX_WORDS, static_cast<unsigned int>(mSuggestions.size()));
@@ -140,6 +142,7 @@
             const unsigned int wordLength = sw->mWordLength;
             char *targetAdr = (char*) outputChars + i * MAX_WORD_LENGTH * sizeof(short);
             frequencies[i] = sw->mScore;
+            outputTypes[i] = sw->mType;
             memcpy(targetAdr, sw->mWord, (wordLength) * sizeof(short));
             if (wordLength < MAX_WORD_LENGTH) {
                 ((unsigned short*) targetAdr)[wordLength] = 0;
@@ -191,10 +194,10 @@
     };
 
     SuggestedWord *getFreeSuggestedWord(int score, unsigned short *word,
-            int wordLength) {
+            int wordLength, int type) {
         for (unsigned int i = 0; i < MAX_WORD_LENGTH; ++i) {
             if (!mSuggestedWords[i].mUsed) {
-                mSuggestedWords[i].setParams(score, word, wordLength);
+                mSuggestedWords[i].setParams(score, word, wordLength, type);
                 return &mSuggestedWords[i];
             }
         }