diff --git a/java/src/com/android/inputmethod/latin/BinaryDictionary.java b/java/src/com/android/inputmethod/latin/BinaryDictionary.java
index 6f3c48c..86c265f 100644
--- a/java/src/com/android/inputmethod/latin/BinaryDictionary.java
+++ b/java/src/com/android/inputmethod/latin/BinaryDictionary.java
@@ -199,6 +199,7 @@
             int[] word, int probability, int timestamp);
     private static native boolean removeNgramEntryNative(long dict,
             int[][] prevWordCodePointArrays, boolean[] isBeginningOfSentenceArray, int[] word);
+    // TODO: Rename to updateEntriesForWordWithNgramContextNative.
     private static native boolean updateCounterNative(long dict,
             int[][] prevWordCodePointArrays, boolean[] isBeginningOfSentenceArray,
             int[] word, boolean isValidWord, int count, int timestamp);
@@ -493,6 +494,25 @@
         return true;
     }
 
+    // Update entries for the word occurrence with the ngramContext.
+    @UsedForTesting
+    public boolean updateEntriesForWordWithNgramContext(final NgramContext ngramContext,
+            final String word, final boolean isValidWord, final int count, final int timestamp) {
+        if (TextUtils.isEmpty(word)) {
+            return false;
+        }
+        final int[][] prevWordCodePointArrays = new int[ngramContext.getPrevWordCount()][];
+        final boolean[] isBeginningOfSentenceArray = new boolean[ngramContext.getPrevWordCount()];
+        ngramContext.outputToArray(prevWordCodePointArrays, isBeginningOfSentenceArray);
+        final int[] wordCodePoints = StringUtils.toCodePointArray(word);
+        if (!updateCounterNative(mNativeDict, prevWordCodePointArrays,
+                isBeginningOfSentenceArray, wordCodePoints, isValidWord, count, timestamp)) {
+            return false;
+        }
+        mHasUpdated = true;
+        return true;
+    }
+
     @UsedForTesting
     public void addMultipleDictionaryEntries(final LanguageModelParam[] languageModelParams) {
         if (!isValidDictionary()) return;
diff --git a/tests/src/com/android/inputmethod/latin/BinaryDictionaryDecayingTests.java b/tests/src/com/android/inputmethod/latin/BinaryDictionaryDecayingTests.java
index 6b6ad21..c3c2bb2 100644
--- a/tests/src/com/android/inputmethod/latin/BinaryDictionaryDecayingTests.java
+++ b/tests/src/com/android/inputmethod/latin/BinaryDictionaryDecayingTests.java
@@ -75,20 +75,25 @@
         return formatVersion > FormatSpec.VERSION401;
     }
 
-    private void addUnigramWord(final BinaryDictionary binaryDictionary, final String word,
-            final int probability) {
-        binaryDictionary.addUnigramEntry(word, probability, "" /* shortcutTarget */,
-                BinaryDictionary.NOT_A_PROBABILITY /* shortcutProbability */,
-                false /* isBeginningOfSentence */, false /* isNotAWord */,
-                false /* isBlacklisted */, mCurrentTime /* timestamp */);
+    private void onInputWord(final BinaryDictionary binaryDictionary, final String word,
+            final boolean isValidWord) {
+        binaryDictionary.updateEntriesForWordWithNgramContext(NgramContext.EMPTY_PREV_WORDS_INFO,
+                word, isValidWord, 1 /* count */, mCurrentTime /* timestamp */);
     }
 
-    private void addBigramWords(final BinaryDictionary binaryDictionary, final String word0,
-            final String word1, final int probability) {
-        binaryDictionary.addNgramEntry(new NgramContext(new WordInfo(word0)), word1, probability,
+    private void onInputWordWithPrevWord(final BinaryDictionary binaryDictionary, final String word,
+            final boolean isValidWord, final String prevWord) {
+        binaryDictionary.updateEntriesForWordWithNgramContext(
+                new NgramContext(new WordInfo(prevWord)), word, isValidWord, 1 /* count */,
                 mCurrentTime /* timestamp */);
     }
 
+    private void onInputWordWithBeginningOfSentenceContext(
+            final BinaryDictionary binaryDictionary, final String word, final boolean isValidWord) {
+        binaryDictionary.updateEntriesForWordWithNgramContext(NgramContext.BEGINNING_OF_SENTENCE,
+                word, isValidWord, 1 /* count */, mCurrentTime /* timestamp */);
+    }
+
     private static boolean isValidBigram(final BinaryDictionary binaryDictionary,
             final String word0, final String word1) {
         return binaryDictionary.isValidNgram(new NgramContext(new WordInfo(word0)), word1);
@@ -175,10 +180,9 @@
         setCurrentTimeForTestMode(mCurrentTime);
         final File dictFile = createEmptyDictionaryAndGetFile(formatVersion);
         final BinaryDictionary binaryDictionary = getBinaryDictionary(dictFile);
-        addUnigramWord(binaryDictionary, "a", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "ab", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "aaa", DUMMY_PROBABILITY);
-        addBigramWords(binaryDictionary, "a", "aaa", DUMMY_PROBABILITY);
+        onInputWord(binaryDictionary, "a", true /* isValidWord */);
+        onInputWord(binaryDictionary, "ab", true /* isValidWord */);
+        onInputWordWithPrevWord(binaryDictionary, "aaa", true /* isValidWord */, "a");
         binaryDictionary.flushWithGC();
         binaryDictionary.close();
 
@@ -229,28 +233,27 @@
         final File dictFile = createEmptyDictionaryAndGetFile(formatVersion);
         final BinaryDictionary binaryDictionary = getBinaryDictionary(dictFile);
 
-        addUnigramWord(binaryDictionary, "a", Dictionary.NOT_A_PROBABILITY);
+        onInputWord(binaryDictionary, "a", false /* isValidWord */);
         assertFalse(binaryDictionary.isValidWord("a"));
-        addUnigramWord(binaryDictionary, "a", Dictionary.NOT_A_PROBABILITY);
-        addUnigramWord(binaryDictionary, "a", Dictionary.NOT_A_PROBABILITY);
+        onInputWord(binaryDictionary, "a", false /* isValidWord */);
+        onInputWord(binaryDictionary, "a", false /* isValidWord */);
         assertTrue(binaryDictionary.isValidWord("a"));
 
-        addUnigramWord(binaryDictionary, "b", DUMMY_PROBABILITY);
+        onInputWord(binaryDictionary, "b", true /* isValidWord */);
         assertTrue(binaryDictionary.isValidWord("b"));
 
-        addBigramWords(binaryDictionary, "a", "b", Dictionary.NOT_A_PROBABILITY);
+        onInputWordWithPrevWord(binaryDictionary, "b", false /* isValidWord */, "a");
         assertFalse(isValidBigram(binaryDictionary, "a", "b"));
-        addBigramWords(binaryDictionary, "a", "b", Dictionary.NOT_A_PROBABILITY);
+        onInputWordWithPrevWord(binaryDictionary, "b", false /* isValidWord */, "a");
         assertTrue(isValidBigram(binaryDictionary, "a", "b"));
 
-        addUnigramWord(binaryDictionary, "c", DUMMY_PROBABILITY);
-        addBigramWords(binaryDictionary, "a", "c", DUMMY_PROBABILITY);
+        onInputWordWithPrevWord(binaryDictionary, "c", true /* isValidWord */, "a");
         assertTrue(isValidBigram(binaryDictionary, "a", "c"));
 
         // Add bigrams of not valid unigrams.
-        addBigramWords(binaryDictionary, "x", "y", Dictionary.NOT_A_PROBABILITY);
+        onInputWordWithPrevWord(binaryDictionary, "y", false /* isValidWord */, "x");
         assertFalse(isValidBigram(binaryDictionary, "x", "y"));
-        addBigramWords(binaryDictionary, "x", "y", DUMMY_PROBABILITY);
+        onInputWordWithPrevWord(binaryDictionary, "y", true /* isValidWord */, "x");
         assertFalse(isValidBigram(binaryDictionary, "x", "y"));
 
         binaryDictionary.close();
@@ -266,36 +269,32 @@
         final File dictFile = createEmptyDictionaryAndGetFile(formatVersion);
         final BinaryDictionary binaryDictionary = getBinaryDictionary(dictFile);
 
-        addUnigramWord(binaryDictionary, "a", DUMMY_PROBABILITY);
+        onInputWord(binaryDictionary, "a", true /* isValidWord */);
         assertTrue(binaryDictionary.isValidWord("a"));
         forcePassingShortTime(binaryDictionary);
         assertFalse(binaryDictionary.isValidWord("a"));
 
-        addUnigramWord(binaryDictionary, "a", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "a", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "a", DUMMY_PROBABILITY);
+        onInputWord(binaryDictionary, "a", true /* isValidWord */);
+        onInputWord(binaryDictionary, "a", true /* isValidWord */);
+        onInputWord(binaryDictionary, "a", true /* isValidWord */);
         assertTrue(binaryDictionary.isValidWord("a"));
         forcePassingShortTime(binaryDictionary);
         assertTrue(binaryDictionary.isValidWord("a"));
         forcePassingLongTime(binaryDictionary);
         assertFalse(binaryDictionary.isValidWord("a"));
 
-        addUnigramWord(binaryDictionary, "a", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "b", DUMMY_PROBABILITY);
-        addBigramWords(binaryDictionary, "a", "b", DUMMY_PROBABILITY);
+        onInputWord(binaryDictionary, "a", true /* isValidWord */);
+        onInputWordWithPrevWord(binaryDictionary, "b", true /* isValidWord */, "a");
         assertTrue(isValidBigram(binaryDictionary, "a", "b"));
         forcePassingShortTime(binaryDictionary);
         assertFalse(isValidBigram(binaryDictionary, "a", "b"));
 
-        addUnigramWord(binaryDictionary, "a", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "b", DUMMY_PROBABILITY);
-        addBigramWords(binaryDictionary, "a", "b", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "a", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "b", DUMMY_PROBABILITY);
-        addBigramWords(binaryDictionary, "a", "b", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "a", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "b", DUMMY_PROBABILITY);
-        addBigramWords(binaryDictionary, "a", "b", DUMMY_PROBABILITY);
+        onInputWord(binaryDictionary, "a", true /* isValidWord */);
+        onInputWordWithPrevWord(binaryDictionary, "b", true /* isValidWord */, "a");
+        onInputWord(binaryDictionary, "a", true /* isValidWord */);
+        onInputWordWithPrevWord(binaryDictionary, "b", true /* isValidWord */, "a");
+        onInputWord(binaryDictionary, "a", true /* isValidWord */);
+        onInputWordWithPrevWord(binaryDictionary, "b", true /* isValidWord */, "a");
         assertTrue(isValidBigram(binaryDictionary, "a", "b"));
         forcePassingShortTime(binaryDictionary);
         assertTrue(isValidBigram(binaryDictionary, "a", "b"));
@@ -333,7 +332,7 @@
                 binaryDictionary.getPropertyForGettingStats(BinaryDictionary.MAX_UNIGRAM_COUNT_QUERY));
         for (int i = 0; i < unigramTypedCount; i++) {
             final String word = words.get(random.nextInt(words.size()));
-            addUnigramWord(binaryDictionary, word, DUMMY_PROBABILITY);
+            onInputWord(binaryDictionary, word, true /* isValidWord */);
 
             if (binaryDictionary.needsToRunGC(true /* mindsBlockByGC */)) {
                 final int unigramCountBeforeGC =
@@ -380,10 +379,10 @@
         final String strong = "strong";
         final String weak = "weak";
         for (int j = 0; j < strongUnigramTypedCount; j++) {
-            addUnigramWord(binaryDictionary, strong, DUMMY_PROBABILITY);
+            onInputWord(binaryDictionary, strong, true /* isValidWord */);
         }
         for (int j = 0; j < weakUnigramTypedCount; j++) {
-            addUnigramWord(binaryDictionary, weak, DUMMY_PROBABILITY);
+            onInputWord(binaryDictionary, weak, true /* isValidWord */);
         }
         assertTrue(binaryDictionary.isValidWord(strong));
         assertTrue(binaryDictionary.isValidWord(weak));
@@ -391,7 +390,7 @@
         for (int i = 0; i < unigramCount; i++) {
             final String word = CodePointUtils.generateWord(random, codePointSet);
             for (int j = 0; j < eachUnigramTypedCount; j++) {
-                addUnigramWord(binaryDictionary, word, DUMMY_PROBABILITY);
+                onInputWord(binaryDictionary, word, true /* isValidWord */);
             }
             if (binaryDictionary.needsToRunGC(true /* mindsBlockByGC */)) {
                 final int unigramCountBeforeGC =
@@ -450,12 +449,13 @@
         }
 
         final int maxBigramCount = Integer.parseInt(
-                binaryDictionary.getPropertyForGettingStats(BinaryDictionary.MAX_BIGRAM_COUNT_QUERY));
+                binaryDictionary.getPropertyForGettingStats(
+                        BinaryDictionary.MAX_BIGRAM_COUNT_QUERY));
         for (int i = 0; i < bigramTypedCount; ++i) {
             final Pair<String, String> bigram = bigrams.get(random.nextInt(bigrams.size()));
-            addUnigramWord(binaryDictionary, bigram.first, DUMMY_PROBABILITY);
-            addUnigramWord(binaryDictionary, bigram.second, DUMMY_PROBABILITY);
-            addBigramWords(binaryDictionary, bigram.first, bigram.second, DUMMY_PROBABILITY);
+            onInputWord(binaryDictionary, bigram.first, true /* isValidWord */);
+            onInputWordWithPrevWord(binaryDictionary, bigram.second, true /* isValidWord */,
+                    bigram.first);
 
             if (binaryDictionary.needsToRunGC(true /* mindsBlockByGC */)) {
                 final int bigramCountBeforeGC =
@@ -506,23 +506,22 @@
             final String word = CodePointUtils.generateWord(random, codePointSet);
             words.add(word);
             for (int j = 0; j < unigramTypedCount; j++) {
-                addUnigramWord(binaryDictionary, word, DUMMY_PROBABILITY);
+                onInputWord(binaryDictionary, word, true /* isValidWord */);
             }
         }
         final String strong = "strong";
         final String weak = "weak";
         final String target = "target";
         for (int j = 0; j < unigramTypedCount; j++) {
-            addUnigramWord(binaryDictionary, strong, DUMMY_PROBABILITY);
-            addUnigramWord(binaryDictionary, weak, DUMMY_PROBABILITY);
-            addUnigramWord(binaryDictionary, target, DUMMY_PROBABILITY);
+            onInputWord(binaryDictionary, weak, true /* isValidWord */);
+            onInputWord(binaryDictionary, strong, true /* isValidWord */);
         }
         binaryDictionary.flushWithGC();
         for (int j = 0; j < strongBigramTypedCount; j++) {
-            addBigramWords(binaryDictionary, strong, target, DUMMY_PROBABILITY);
+            onInputWordWithPrevWord(binaryDictionary, target, true /* isValidWord */, strong);
         }
         for (int j = 0; j < weakBigramTypedCount; j++) {
-            addBigramWords(binaryDictionary, weak, target, DUMMY_PROBABILITY);
+            onInputWordWithPrevWord(binaryDictionary, target, true /* isValidWord */, weak);
         }
         assertTrue(isValidBigram(binaryDictionary, strong, target));
         assertTrue(isValidBigram(binaryDictionary, weak, target));
@@ -535,7 +534,7 @@
             final String word1 = words.get(word1Index);
 
             for (int j = 0; j < eachBigramTypedCount; j++) {
-                addBigramWords(binaryDictionary, word0, word1, DUMMY_PROBABILITY);
+                onInputWordWithPrevWord(binaryDictionary, word1, true /* isValidWord */, word0);
             }
             if (binaryDictionary.needsToRunGC(true /* mindsBlockByGC */)) {
                 final int bigramCountBeforeGC =
@@ -563,19 +562,18 @@
         setCurrentTimeForTestMode(mCurrentTime);
         final File dictFile = createEmptyDictionaryAndGetFile(fromFormatVersion);
         final BinaryDictionary binaryDictionary = getBinaryDictionary(dictFile);
-        addUnigramWord(binaryDictionary, "aaa", DUMMY_PROBABILITY);
+        onInputWord(binaryDictionary, "aaa", true /* isValidWord */);
         assertTrue(binaryDictionary.isValidWord("aaa"));
-        addUnigramWord(binaryDictionary, "bbb", Dictionary.NOT_A_PROBABILITY);
-        assertFalse(binaryDictionary.isValidWord("bbb"));
-        addUnigramWord(binaryDictionary, "ccc", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "ccc", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "ccc", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "ccc", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "ccc", DUMMY_PROBABILITY);
-        addUnigramWord(binaryDictionary, "abc", DUMMY_PROBABILITY);
-        addBigramWords(binaryDictionary, "aaa", "abc", DUMMY_PROBABILITY);
+        onInputWord(binaryDictionary, "ccc", true /* isValidWord */);
+        onInputWord(binaryDictionary, "ccc", true /* isValidWord */);
+        onInputWord(binaryDictionary, "ccc", true /* isValidWord */);
+        onInputWord(binaryDictionary, "ccc", true /* isValidWord */);
+        onInputWord(binaryDictionary, "ccc", true /* isValidWord */);
+
+        onInputWordWithPrevWord(binaryDictionary, "abc", true /* isValidWord */, "aaa");
         assertTrue(isValidBigram(binaryDictionary, "aaa", "abc"));
-        addBigramWords(binaryDictionary, "aaa", "bbb", Dictionary.NOT_A_PROBABILITY);
+        onInputWordWithPrevWord(binaryDictionary, "bbb", false /* isValidWord */, "aaa");
+        assertFalse(binaryDictionary.isValidWord("bbb"));
         assertFalse(isValidBigram(binaryDictionary, "aaa", "bbb"));
 
         assertEquals(fromFormatVersion, binaryDictionary.getFormatVersion());
@@ -585,11 +583,11 @@
         assertTrue(binaryDictionary.isValidWord("aaa"));
         assertFalse(binaryDictionary.isValidWord("bbb"));
         assertTrue(binaryDictionary.getFrequency("aaa") < binaryDictionary.getFrequency("ccc"));
-        addUnigramWord(binaryDictionary, "bbb", Dictionary.NOT_A_PROBABILITY);
+        onInputWord(binaryDictionary, "bbb", false /* isValidWord */);
         assertTrue(binaryDictionary.isValidWord("bbb"));
         assertTrue(isValidBigram(binaryDictionary, "aaa", "abc"));
         assertFalse(isValidBigram(binaryDictionary, "aaa", "bbb"));
-        addBigramWords(binaryDictionary, "aaa", "bbb", Dictionary.NOT_A_PROBABILITY);
+        onInputWordWithPrevWord(binaryDictionary, "bbb", false /* isValidWord */, "aaa");
         assertTrue(isValidBigram(binaryDictionary, "aaa", "bbb"));
         binaryDictionary.close();
     }
@@ -612,28 +610,19 @@
                 true /* isBeginningOfSentence */, true /* isNotAWord */, false /* isBlacklisted */,
                 mCurrentTime);
         final NgramContext beginningOfSentenceContext = NgramContext.BEGINNING_OF_SENTENCE;
-        addUnigramWord(binaryDictionary, "aaa", DUMMY_PROBABILITY);
-        binaryDictionary.addNgramEntry(beginningOfSentenceContext, "aaa", DUMMY_PROBABILITY,
-                mCurrentTime);
+        onInputWordWithBeginningOfSentenceContext(binaryDictionary, "aaa", true /* isValidWord */);
         assertTrue(binaryDictionary.isValidNgram(beginningOfSentenceContext, "aaa"));
-        binaryDictionary.addNgramEntry(beginningOfSentenceContext, "aaa", DUMMY_PROBABILITY,
-                mCurrentTime);
-        addUnigramWord(binaryDictionary, "bbb", DUMMY_PROBABILITY);
-        binaryDictionary.addNgramEntry(beginningOfSentenceContext, "bbb", DUMMY_PROBABILITY,
-                mCurrentTime);
+        onInputWordWithBeginningOfSentenceContext(binaryDictionary, "aaa", true /* isValidWord */);
+        onInputWordWithBeginningOfSentenceContext(binaryDictionary, "bbb", true /* isValidWord */);
+        onInputWordWithBeginningOfSentenceContext(binaryDictionary, "bbb", true /* isValidWord */);
         assertTrue(binaryDictionary.isValidNgram(beginningOfSentenceContext, "aaa"));
         assertTrue(binaryDictionary.isValidNgram(beginningOfSentenceContext, "bbb"));
-
         forcePassingLongTime(binaryDictionary);
         assertFalse(binaryDictionary.isValidNgram(beginningOfSentenceContext, "aaa"));
         assertFalse(binaryDictionary.isValidNgram(beginningOfSentenceContext, "bbb"));
 
-        addUnigramWord(binaryDictionary, "aaa", DUMMY_PROBABILITY);
-        binaryDictionary.addNgramEntry(beginningOfSentenceContext, "aaa", DUMMY_PROBABILITY,
-                mCurrentTime);
-        addUnigramWord(binaryDictionary, "bbb", DUMMY_PROBABILITY);
-        binaryDictionary.addNgramEntry(beginningOfSentenceContext, "bbb", DUMMY_PROBABILITY,
-                mCurrentTime);
+        onInputWordWithBeginningOfSentenceContext(binaryDictionary, "aaa", true /* isValidWord */);
+        onInputWordWithBeginningOfSentenceContext(binaryDictionary, "bbb", true /* isValidWord */);
         assertTrue(binaryDictionary.isValidNgram(beginningOfSentenceContext, "aaa"));
         assertTrue(binaryDictionary.isValidNgram(beginningOfSentenceContext, "bbb"));
         binaryDictionary.close();
@@ -651,10 +640,10 @@
         final File dictFile = createEmptyDictionaryAndGetFile(formatVersion);
         final BinaryDictionary binaryDictionary = getBinaryDictionary(dictFile);
 
-        addUnigramWord(binaryDictionary, "aaa", Dictionary.NOT_A_PROBABILITY);
+        onInputWord(binaryDictionary, "aaa", false /* isValidWord */);
         assertFalse(binaryDictionary.isValidWord("aaa"));
         for (int i = 0; i < unigramInputCount; i++) {
-            addUnigramWord(binaryDictionary, "aaa", Dictionary.NOT_A_PROBABILITY);
+            onInputWord(binaryDictionary, "aaa", false /* isValidWord */);
         }
         assertTrue(binaryDictionary.isValidWord("aaa"));
         assertTrue(binaryDictionary.removeUnigramEntry("aaa"));