2 files changed, 23 insertions, 20 deletions

diff --git a/java/src/com/android/inputmethod/research/FixedLogBuffer.java b/java/src/com/android/inputmethod/research/FixedLogBuffer.java
index 78dc59562..641bf7eae 100644
--- a/java/src/com/android/inputmethod/research/FixedLogBuffer.java
+++ b/java/src/com/android/inputmethod/research/FixedLogBuffer.java
@@ -51,10 +51,6 @@ public class FixedLogBuffer extends LogBuffer {
         mNumActualWords = 0;
     }
 
-    protected int getNumActualWords() {
-        return mNumActualWords;
-    }
-
     /**
      * Adds a new LogUnit to the front of the LIFO queue, evicting existing LogUnit's
      * (oldest first) if word capacity is reached.
@@ -119,12 +115,24 @@ public class FixedLogBuffer extends LogBuffer {
         return logUnit;
     }
 
-    protected void shiftOutWords(final int numWords) {
-        final int targetNumWords = mNumActualWords - numWords;
-        final LinkedList<LogUnit> logUnits = getLogUnits();
-        while (mNumActualWords > targetNumWords && !logUnits.isEmpty()) {
-            shiftOut();
+    /**
+     * Remove LogUnits from the front of the LogBuffer until {@code numWords} have been removed.
+     *
+     * If there are less than {@code numWords} word-containing {@link LogUnit}s, shifts out
+     * all {@code LogUnit}s in the buffer.
+     *
+     * @param numWords the number of word-containing {@link LogUnit}s to shift out
+     * @return the number of actual {@code LogUnit}s shifted out
+     */
+    protected int shiftOutWords(final int numWords) {
+        int numWordContainingLogUnitsShiftedOut = 0;
+        for (LogUnit logUnit = shiftOut(); logUnit != null
+                && numWordContainingLogUnitsShiftedOut < numWords; logUnit = shiftOut()) {
+            if (logUnit.hasWord()) {
+                numWordContainingLogUnitsShiftedOut++;
+            }
         }
+        return numWordContainingLogUnitsShiftedOut;
     }
 
     public void shiftOutAll() {
diff --git a/java/src/com/android/inputmethod/research/MainLogBuffer.java b/java/src/com/android/inputmethod/research/MainLogBuffer.java
index 3303d2bdb..cd4c1db6e 100644
--- a/java/src/com/android/inputmethod/research/MainLogBuffer.java
+++ b/java/src/com/android/inputmethod/research/MainLogBuffer.java
@@ -25,7 +25,6 @@ import com.android.inputmethod.latin.define.ProductionFlag;
 
 import java.util.ArrayList;
 import java.util.LinkedList;
-import java.util.Random;
 
 /**
  * MainLogBuffer is a FixedLogBuffer that tracks the state of LogUnits to make privacy guarantees.
@@ -100,10 +99,6 @@ public abstract class MainLogBuffer extends FixedLogBuffer {
         return mSuggest.getMainDictionary();
     }
 
-    public void resetWordCounter() {
-        mNumWordsUntilSafeToSample = mNumWordsBetweenNGrams;
-    }
-
     public void setIsStopping() {
         mIsStopping = true;
     }
@@ -201,7 +196,7 @@ public abstract class MainLogBuffer extends FixedLogBuffer {
             // Good n-gram at the front of the buffer.  Publish it, disclosing details.
             publish(logUnits, true /* canIncludePrivateData */);
             shiftOutWords(N_GRAM_SIZE);
-            resetWordCounter();
+            mNumWordsUntilSafeToSample = mNumWordsBetweenNGrams;
         } else {
             // No good n-gram at front, and buffer is full.  Shift out the first word (or if there
             // is none, the existing logUnits).
@@ -224,13 +219,13 @@ public abstract class MainLogBuffer extends FixedLogBuffer {
             final boolean canIncludePrivateData);
 
     @Override
-    protected void shiftOutWords(final int numWords) {
-        final int oldNumActualWords = getNumActualWords();
-        super.shiftOutWords(numWords);
-        final int numWordsShifted = oldNumActualWords - getNumActualWords();
-        mNumWordsUntilSafeToSample -= numWordsShifted;
+    protected int shiftOutWords(final int numWords) {
+        final int numWordContainingLogUnitsShiftedOut = super.shiftOutWords(numWords);
+        mNumWordsUntilSafeToSample = Math.max(0, mNumWordsUntilSafeToSample
+                - numWordContainingLogUnitsShiftedOut);
         if (DEBUG) {
             Log.d(TAG, "wordsUntilSafeToSample now at " + mNumWordsUntilSafeToSample);
         }
+        return numWordContainingLogUnitsShiftedOut;
     }
 }