Implement reset to original selection.

The UX for smart selection specifies that the user is able to reset
the original selection by tapping once on the word after smart
selection. Tap once, and the selection resets to the word that was
initially selected (before smart selection happened), tap one more
time and the insertion cursor appears as usual. If the user taps
a different word other than the original selection, the insertion
cursor appears as before.

Test: I7456eb4773d40366a3f4aa7bf051a1c7ddda6e9d
Bug: 34777048
Change-Id: If73259ddb67379d5a5deaa01b840b5185cedf4c8
diff --git a/core/java/android/widget/Editor.java b/core/java/android/widget/Editor.java
index dd3b054..ade03e1 100644
--- a/core/java/android/widget/Editor.java
+++ b/core/java/android/widget/Editor.java
@@ -2183,6 +2183,11 @@
     }
 
     void onTouchUpEvent(MotionEvent event) {
+        if (getSelectionActionModeHelper().resetOriginalSelection(
+                getTextView().getOffsetForPosition(event.getX(), event.getY()))) {
+            return;
+        }
+
         boolean selectAllGotFocus = mSelectAllOnFocus && mTextView.didTouchFocusSelect();
         hideCursorAndSpanControllers();
         stopTextActionMode();
@@ -3916,7 +3921,7 @@
         @Override
         public void onDestroyActionMode(ActionMode mode) {
             // Clear mTextActionMode not to recursively destroy action mode by clearing selection.
-            getSelectionActionModeHelper().cancelAsyncTask();
+            getSelectionActionModeHelper().onDestroyActionMode();
             mTextActionMode = null;
             Callback customCallback = getCustomCallback();
             if (customCallback != null) {
diff --git a/core/java/android/widget/SelectionActionModeHelper.java b/core/java/android/widget/SelectionActionModeHelper.java
index 770d9ee..6790532 100644
--- a/core/java/android/widget/SelectionActionModeHelper.java
+++ b/core/java/android/widget/SelectionActionModeHelper.java
@@ -54,6 +54,8 @@
     private TextClassificationResult mTextClassificationResult;
     private AsyncTask mTextClassificationAsyncTask;
 
+    private final SelectionInfo mSelectionInfo = new SelectionInfo();
+
     SelectionActionModeHelper(@NonNull Editor editor) {
         mEditor = Preconditions.checkNotNull(editor);
         final TextView textView = mEditor.getTextView();
@@ -94,12 +96,12 @@
         }
     }
 
-    public void cancelAsyncTask() {
-        if (mTextClassificationAsyncTask != null) {
-            mTextClassificationAsyncTask.cancel(true);
-            mTextClassificationAsyncTask = null;
+    public boolean resetOriginalSelection(int textIndex) {
+        if (mSelectionInfo.resetOriginalSelection(textIndex, mEditor.getTextView().getText())) {
+            invalidateActionModeAsync();
+            return true;
         }
-        mTextClassificationResult = null;
+        return false;
     }
 
     @Nullable
@@ -107,12 +109,28 @@
         return mTextClassificationResult;
     }
 
+    public void onDestroyActionMode() {
+        mSelectionInfo.onSelectionDestroyed();
+        cancelAsyncTask();
+    }
+
+    private void cancelAsyncTask() {
+        if (mTextClassificationAsyncTask != null) {
+            mTextClassificationAsyncTask.cancel(true);
+            mTextClassificationAsyncTask = null;
+        }
+        mTextClassificationResult = null;
+    }
+
     private boolean isNoOpTextClassifier() {
         return mEditor.getTextView().getTextClassifier() == TextClassifier.NO_OP;
     }
 
     private void startActionMode(@Nullable SelectionResult result) {
-        final CharSequence text = mEditor.getTextView().getText();
+        final TextView textView = mEditor.getTextView();
+        final CharSequence text = textView.getText();
+        mSelectionInfo.setOriginalSelection(
+                textView.getSelectionStart(), textView.getSelectionEnd());
         if (result != null && text instanceof Spannable) {
             Selection.setSelection((Spannable) text, result.mStart, result.mEnd);
             mTextClassificationResult = result.mResult;
@@ -124,6 +142,9 @@
             if (controller != null) {
                 controller.show();
             }
+            if (result != null) {
+                mSelectionInfo.onSelectionStarted(result.mStart, result.mEnd);
+            }
         }
         mEditor.setRestartActionModeOnNextRefresh(false);
         mTextClassificationAsyncTask = null;
@@ -135,6 +156,8 @@
         if (actionMode != null) {
             actionMode.invalidate();
         }
+        final TextView textView = mEditor.getTextView();
+        mSelectionInfo.onSelectionUpdated(textView.getSelectionStart(), textView.getSelectionEnd());
         mTextClassificationAsyncTask = null;
     }
 
@@ -145,6 +168,56 @@
     }
 
     /**
+     * Holds information about the selection and uses it to decide on whether or not to update
+     * the selection when resetOriginalSelection is called.
+     * The expected UX here is to allow the user to re-snap the selection back to the original word
+     * that was selected with one tap on that word.
+     */
+    private static final class SelectionInfo {
+
+        private int mOriginalStart;
+        private int mOriginalEnd;
+        private int mSelectionStart;
+        private int mSelectionEnd;
+
+        private boolean mResetOriginal;
+
+        public void setOriginalSelection(int selectionStart, int selectionEnd) {
+            mOriginalStart = selectionStart;
+            mOriginalEnd = selectionEnd;
+            mResetOriginal = false;
+        }
+
+        public void onSelectionStarted(int selectionStart, int selectionEnd) {
+            // Set the reset flag to true if the selection changed.
+            mSelectionStart = selectionStart;
+            mSelectionEnd = selectionEnd;
+            mResetOriginal = mSelectionStart != mOriginalStart || mSelectionEnd != mOriginalEnd;
+        }
+
+        public void onSelectionUpdated(int selectionStart, int selectionEnd) {
+            // If the selection did not change, maintain the reset state. Otherwise, disable reset.
+            mResetOriginal &= selectionStart == mSelectionStart && selectionEnd == mSelectionEnd;
+        }
+
+        public void onSelectionDestroyed() {
+            mResetOriginal = false;
+        }
+
+        public boolean resetOriginalSelection(int textIndex, CharSequence text) {
+            if (mResetOriginal
+                    && textIndex >= mOriginalStart && textIndex <= mOriginalEnd
+                    && text instanceof Spannable) {
+                Selection.setSelection((Spannable) text, mOriginalStart, mOriginalEnd);
+                // Only allow a reset once.
+                mResetOriginal = false;
+                return true;
+            }
+            return false;
+        }
+    }
+
+    /**
      * AsyncTask for running a query on a background thread and returning the result on the
      * UiThread. The AsyncTask times out after a specified time, returning a null result if the
      * query has not yet returned.