Allow selection change update before beginBatchEdit

content/public/android/java/src/org/chromium/content/browser/input/ThreadedInputConnection.java

Index: content/public/android/java/src/org/chromium/content/browser/input/ThreadedInputConnection.java
diff --git a/content/public/android/java/src/org/chromium/content/browser/input/ThreadedInputConnection.java b/content/public/android/java/src/org/chromium/content/browser/input/ThreadedInputConnection.java
index 12b3906a2dc62f7e41fb9fe28dc649619779124c..30203f46c924e855d264213f10a9496f4f5443b0 100644
--- a/content/public/android/java/src/org/chromium/content/browser/input/ThreadedInputConnection.java
+++ b/content/public/android/java/src/org/chromium/content/browser/input/ThreadedInputConnection.java
@@ -41,7 +41,7 @@ public class ThreadedInputConnection extends BaseInputConnection
     private static final boolean DEBUG_LOGS = false;
 
     private static final TextInputState UNBLOCKER = new TextInputState(
-            "", new Range(0, 0), new Range(-1, -1), false, false /* notFromIme */) {
+            "", new Range(0, 0), new Range(-1, -1), false, false /* notFromIme */, false) {
 
         @Override
         public boolean shouldUnblock() {
@@ -74,6 +74,22 @@ public class ThreadedInputConnection extends BaseInputConnection
         }
     };
 
+    private final Runnable mBeginBatchEdit = new Runnable() {
+        @Override
+        public void run() {
+            boolean result = mImeAdapter.beginBatchEdit();
+            if (!result) unblockOnUiThread();
+        }
+    };
+
+    private final Runnable mEndBatchEdit = new Runnable() {
+        @Override
+        public void run() {
+            boolean result = mImeAdapter.endBatchEdit();
+            if (!result) unblockOnUiThread();
+        }
+    };
+
     private final Runnable mNotifyUserActionRunnable = new Runnable() {
         @Override
         public void run() {
@@ -97,6 +113,7 @@ public class ThreadedInputConnection extends BaseInputConnection
     private final BlockingQueue<TextInputState> mQueue = new LinkedBlockingQueue<>();
     private int mPendingAccent;
     private TextInputState mCachedTextInputState;
+    private boolean mLastInBatchEditMode;
 
     ThreadedInputConnection(View view, ImeAdapter imeAdapter, Handler handler) {
         super(view, true);
@@ -104,28 +121,34 @@ public class ThreadedInputConnection extends BaseInputConnection
         ImeUtils.checkOnUiThread();
         mImeAdapter = imeAdapter;
         mHandler = handler;
+        mImeAdapter.endBatchEdit();
     }
 
     void resetOnUiThread() {
         ImeUtils.checkOnUiThread();
         mNumNestedBatchEdits = 0;
         mPendingAccent = 0;
+        mImeAdapter.endBatchEdit();
     }
 
     @Override
-    public void updateStateOnUiThread(final String text, final int selectionStart,
-            final int selectionEnd, final int compositionStart, final int compositionEnd,
-            boolean singleLine, final boolean isNonImeChange) {
+    public void updateStateOnUiThread(String text, int selectionStart,
+            int selectionEnd, int compositionStart, int compositionEnd,
+            boolean singleLine, boolean isNonImeChange, boolean inBatchEditMode) {
         ImeUtils.checkOnUiThread();
 
         mCachedTextInputState =
                 new TextInputState(text, new Range(selectionStart, selectionEnd),
-                        new Range(compositionStart, compositionEnd), singleLine, !isNonImeChange);
+                        new Range(compositionStart, compositionEnd), singleLine, !isNonImeChange,
+                        inBatchEditMode);
         if (DEBUG_LOGS) Log.w(TAG, "updateState: %s", mCachedTextInputState);
 
         addToQueueOnUiThread(mCachedTextInputState);
-        if (isNonImeChange) {
+        // If state update is caused by explicitly requesting the state update,
+        // or batch edit just finished, then we may need to update state to IMM.
+        if (isNonImeChange || (mLastInBatchEditMode && !inBatchEditMode)) {
             mHandler.post(mProcessPendingInputStatesRunnable);
+            mLastInBatchEditMode = inBatchEditMode;
         }
     }
 
@@ -198,7 +221,7 @@ public class ThreadedInputConnection extends BaseInputConnection
     private void updateSelection(TextInputState textInputState) {
         if (textInputState == null) return;
         assertOnImeThread();
-        if (mNumNestedBatchEdits != 0) return;
+        if (textInputState.inBatchEditMode()) return;
         Range selection = textInputState.selection();
         Range composition = textInputState.composition();
         mImeAdapter.updateSelection(
@@ -385,8 +408,10 @@ public class ThreadedInputConnection extends BaseInputConnection
     @Override
     public boolean beginBatchEdit() {
         if (DEBUG_LOGS) Log.w(TAG, "beginBatchEdit [%b]", (mNumNestedBatchEdits == 0));
-        assertOnImeThread();
         mNumNestedBatchEdits++;
+        if (mNumNestedBatchEdits == 1) {
+            ThreadUtils.postOnUiThread(mBeginBatchEdit);
+        }
         return true;
     }
 
@@ -395,12 +420,11 @@ public class ThreadedInputConnection extends BaseInputConnection
      */
     @Override
     public boolean endBatchEdit() {
-        assertOnImeThread();
         if (mNumNestedBatchEdits == 0) return false;
         --mNumNestedBatchEdits;
         if (DEBUG_LOGS) Log.w(TAG, "endBatchEdit [%b]", (mNumNestedBatchEdits == 0));
         if (mNumNestedBatchEdits == 0) {
-            updateSelection(requestAndWaitForTextInputState());
+            ThreadUtils.postOnUiThread(mEndBatchEdit);
         }
         return mNumNestedBatchEdits != 0;
     }