diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index fa63885..80c80d7 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -2761,8 +2761,10 @@
         return super.onKeyUp(keyCode, event);
     }
 
-    protected StateHandler<LauncherState>[] createStateHandlers() {
-        return new StateHandler[] { getAllAppsController(), getWorkspace() };
+    @Override
+    protected void collectStateHandlers(List<StateHandler> out) {
+        out.add(getAllAppsController());
+        out.add(getWorkspace());
     }
 
     public TouchController[] createTouchControllers() {
diff --git a/src/com/android/launcher3/allapps/AllAppsInsetTransitionController.java b/src/com/android/launcher3/allapps/AllAppsInsetTransitionController.java
deleted file mode 100644
index b34c8b8..0000000
--- a/src/com/android/launcher3/allapps/AllAppsInsetTransitionController.java
+++ /dev/null
@@ -1,315 +0,0 @@
-/*
- * Copyright (C) 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.launcher3.allapps;
-
-import android.annotation.TargetApi;
-import android.graphics.Insets;
-import android.os.Build;
-import android.util.Log;
-import android.view.View;
-import android.view.WindowInsets;
-import android.view.WindowInsetsAnimationControlListener;
-import android.view.WindowInsetsAnimationController;
-import android.view.animation.Interpolator;
-import android.view.animation.LinearInterpolator;
-
-import androidx.annotation.Nullable;
-
-import com.android.launcher3.Utilities;
-import com.android.launcher3.util.UiThreadHelper;
-
-/**
- * Handles IME over all apps to be synchronously transitioning along with the passed in
- * root inset.
- */
-public class AllAppsInsetTransitionController {
-
-    private static final boolean DEBUG = true;
-    private static final String TAG = "AllAppsInsetTransitionController";
-    private static final Interpolator LINEAR = new LinearInterpolator();
-
-    private WindowInsetsAnimationController mAnimationController;
-    private WindowInsetsAnimationControlListener mCurrentRequest;
-
-    private Runnable mSearchEduRunnable;
-
-    private float mAllAppsHeight;
-
-    private int mDownInsetBottom;
-    private boolean mShownAtDown;
-
-    private int mHiddenInsetBottom;
-    private int mShownInsetBottom;
-
-    private float mDown, mCurrent;
-    private View mApps;
-
-    /**
-     *
-     */
-    public boolean showSearchEduIfNecessary() {
-        if (mSearchEduRunnable == null) {
-            return false;
-        }
-        mSearchEduRunnable.run();
-        return true;
-    }
-
-    public void setSearchEduRunnable(Runnable eduRunnable) {
-        mSearchEduRunnable = eduRunnable;
-    }
-
-    // Only purpose of these states is to keep track of fast fling transition
-    enum State {
-        RESET, DRAG_START_BOTTOM, DRAG_START_BOTTOM_IME_CANCELLED,
-        FLING_END_TOP, FLING_END_TOP_IME_CANCELLED,
-        DRAG_START_TOP, FLING_END_BOTTOM
-    }
-
-    private State mState;
-
-    public AllAppsInsetTransitionController(float allAppsHeight, View appsView) {
-        mAllAppsHeight = allAppsHeight;
-        mApps = appsView;
-    }
-
-    public void show() {
-        mApps.getWindowInsetsController().show(WindowInsets.Type.ime());
-    }
-
-    public void hide() {
-        if (!Utilities.ATLEAST_R) return;
-
-        WindowInsets insets = mApps.getRootWindowInsets();
-        if (insets == null) return;
-
-        boolean imeVisible = insets.isVisible(WindowInsets.Type.ime());
-
-        if (DEBUG) {
-            Log.d(TAG, "\nhide imeVisible=" + imeVisible);
-        }
-        if (insets.isVisible(WindowInsets.Type.ime())) {
-            mApps.getWindowInsetsController().hide(WindowInsets.Type.ime());
-        }
-    }
-
-    /**
-     * Initializes member variables and requests for the {@link WindowInsetsAnimationController}
-     * object.
-     *
-     * @param progress value between 0..1
-     */
-    @TargetApi(Build.VERSION_CODES.R)
-    public void onDragStart(float progress) {
-        if (!Utilities.ATLEAST_R) return;
-
-        // Until getRootWindowInsets().isVisible(...) method returns correct value,
-        // only support InsetController based IME transition during swipe up and
-        // NOT swipe down
-        if (Float.compare(progress, 0f) == 0) return;
-
-        setState(true, false, progress);
-        mDown = progress * mAllAppsHeight;
-
-        // Below two values are sometimes incorrect. Possibly a platform bug
-        // mDownInsetBottom = mApps.getRootWindowInsets().getInsets(WindowInsets.Type.ime()).bottom;
-        // mShownAtDown = mApps.getRootWindowInsets().isVisible(WindowInsets.Type.ime());
-
-        if (DEBUG) {
-            Log.d(TAG, "\nonDragStart progress=" + progress
-                    + " mDownInsets=" + mDownInsetBottom
-                    + " mShownAtDown=" + mShownAtDown);
-        }
-
-        mApps.getWindowInsetsController().controlWindowInsetsAnimation(
-                WindowInsets.Type.ime(), -1 /* no predetermined duration */, LINEAR, null,
-                mCurrentRequest = new WindowInsetsAnimationControlListener() {
-
-                    @Override
-                    public void onReady(WindowInsetsAnimationController controller, int types) {
-                        if (DEBUG) {
-                            Log.d(TAG, "Listener.onReady " + (mCurrentRequest == this));
-                        }
-                        if (controller != null) {
-                            if (mCurrentRequest == this && !handleFinishOnFling(controller)) {
-                                mAnimationController = controller;
-                            } else {
-                                controller.finish(false /* just don't show */);
-                            }
-                        }
-                    }
-
-                    @Override
-                    public void onFinished(WindowInsetsAnimationController controller) {
-                        // when screen lock happens, then this method get called
-                        if (DEBUG) {
-                            Log.d(TAG, "Listener.onFinished ctrl=" + controller
-                                    + " mAnimationController=" + mAnimationController);
-                        }
-                        if (mAnimationController != null) {
-                            mAnimationController.finish(true);
-                            mAnimationController = null;
-                        }
-                    }
-
-                    @Override
-                    public void onCancelled(@Nullable WindowInsetsAnimationController controller) {
-                        if (DEBUG) {
-                            // Keep the verbose logging to chase down IME not showing up issue.
-                            // b/178904132
-                            Log.e(TAG, "Listener.onCancelled ctrl=" + controller
-                                    + " mAnimationController=" + mAnimationController,
-                                    new Exception());
-                        }
-                        if (mState == State.DRAG_START_BOTTOM) {
-                            mState = State.DRAG_START_BOTTOM_IME_CANCELLED;
-                        }
-                        mAnimationController = null;
-                        if (controller != null) {
-                            controller.finish(true);
-                        }
-                    }
-                });
-    }
-
-    /**
-     * If IME bounds after touch sequence finishes, call finish.
-     */
-    private boolean handleFinishOnFling(WindowInsetsAnimationController controller) {
-        if (!Utilities.ATLEAST_R) return false;
-
-        if (mState == State.FLING_END_TOP) {
-            controller.finish(true);
-            return true;
-        } else if (mState == State.FLING_END_BOTTOM) {
-            controller.finish(false);
-            return true;
-        }
-        return false;
-    }
-
-    /**
-     * Handles the translation using the progress.
-     *
-     * @param progress value between 0..1
-     */
-    @TargetApi(Build.VERSION_CODES.R)
-    public void setProgress(float progress) {
-        if (!Utilities.ATLEAST_R) return;
-        // progress that equals to 0 or 1 is error prone. Do not use them.
-        // Instead use onDragStart and onAnimationEnd
-        if (mAnimationController == null || progress <= 0f || progress >= 1f) return;
-
-        mCurrent = progress * mAllAppsHeight;
-        mHiddenInsetBottom = mAnimationController.getHiddenStateInsets().bottom; // 0
-        mShownInsetBottom = mAnimationController.getShownStateInsets().bottom; // 1155
-
-        int shift = mShownAtDown ? 0 : (int) (mAllAppsHeight - mShownInsetBottom);
-
-        int inset = (int) (mDownInsetBottom + (mDown - mCurrent) - shift);
-
-        final int start = mShownAtDown ? mShownInsetBottom : mHiddenInsetBottom;
-        final int end = mShownAtDown ? mHiddenInsetBottom : mShownInsetBottom;
-        inset = Math.max(inset, mHiddenInsetBottom);
-        inset = Math.min(inset, mShownInsetBottom);
-        if (DEBUG && false) {
-            Log.d(TAG, "updateInset mCurrent=" + mCurrent + " mDown="
-                    + mDown + " hidden=" + mHiddenInsetBottom
-                    + " shown=" + mShownInsetBottom
-                    + " mDownInsets.bottom=" + mDownInsetBottom + " inset=" + inset
-                    + " shift= " + shift);
-        }
-
-        mAnimationController.setInsetsAndAlpha(
-                Insets.of(0, 0, 0, inset),
-                1f, (inset - start) / (float) (end - start));
-    }
-
-    /**
-     * Report to the animation controller that we no longer plan to translate anymore.
-     *
-     * @param progress value between 0..1
-     */
-    @TargetApi(Build.VERSION_CODES.R)
-    public void onAnimationEnd(float progress) {
-        if (DEBUG) {
-            Log.d(TAG, "onAnimationEnd progress=" + progress
-                    + " mAnimationController=" + mAnimationController);
-        }
-        if (mState == null) {
-            // only called when launcher restarting.
-            UiThreadHelper.hideKeyboardAsync(mApps.getContext(), mApps.getWindowToken());
-        }
-
-        setState(false, true, progress);
-
-
-        if (mAnimationController == null) {
-            if (mState == State.FLING_END_TOP_IME_CANCELLED) {
-                mApps.getWindowInsetsController().show(WindowInsets.Type.ime());
-            }
-            return;
-        }
-
-        /* handle finish */
-        if (mState == State.FLING_END_TOP) {
-            mAnimationController.finish(true /* show */);
-        } else {
-            if (Float.compare(progress, 1f) == 0 /* bottom */) {
-                mAnimationController.finish(false /* gone */);
-            } else {
-                mAnimationController.finish(mShownAtDown);
-            }
-        }
-        /* handle finish */
-
-        if (DEBUG) {
-            Log.d(TAG, "endTranslation progress=" + progress
-                    + " mAnimationController=" + mAnimationController);
-        }
-        mAnimationController = null;
-        mCurrentRequest = null;
-        setState(false, false, progress);
-    }
-
-    private void setState(boolean start, boolean end, float progress) {
-        State state = State.RESET;
-        if (start && end) {
-            throw new IllegalStateException("drag start and end cannot happen in same call");
-        }
-        if (start) {
-            if (Float.compare(progress, 1f) == 0) {
-                state = State.DRAG_START_BOTTOM;
-            } else if (Float.compare(progress, 0f) == 0) {
-                state = State.DRAG_START_TOP;
-            }
-        } else if (end) {
-            if (Float.compare(progress, 1f) == 0 && mState == State.DRAG_START_TOP) {
-                state = State.FLING_END_BOTTOM;
-            } else if (Float.compare(progress, 0f) == 0) {
-                if (mState == State.DRAG_START_BOTTOM) {
-                    state = State.FLING_END_TOP;
-                } else if (mState == State.DRAG_START_BOTTOM_IME_CANCELLED) {
-                    state = State.FLING_END_TOP_IME_CANCELLED;
-                }
-            }
-        }
-        if (DEBUG) {
-            Log.d(TAG, "setState " + mState + " -> " + state);
-        }
-        mState = state;
-    }
-}
diff --git a/src/com/android/launcher3/allapps/AllAppsTransitionController.java b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
index a48e423..abf63dc 100644
--- a/src/com/android/launcher3/allapps/AllAppsTransitionController.java
+++ b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
@@ -33,18 +33,15 @@
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.animation.ObjectAnimator;
-import android.content.SharedPreferences;
 import android.util.FloatProperty;
 import android.view.View;
 import android.view.animation.Interpolator;
-import android.widget.EditText;
-
-import androidx.core.os.BuildCompat;
 
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.DeviceProfile.OnDeviceProfileChangeListener;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
+import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.AnimationSuccessListener;
 import com.android.launcher3.anim.PendingAnimation;
 import com.android.launcher3.anim.PropertySetter;
@@ -63,8 +60,8 @@
  * If release velocity < THRES1, snap according to either top or bottom depending on whether it's
  * closer to top or closer to the page indicator.
  */
-public class AllAppsTransitionController implements StateHandler<LauncherState>,
-        OnDeviceProfileChangeListener, SharedPreferences.OnSharedPreferenceChangeListener {
+public class AllAppsTransitionController
+        implements StateHandler<LauncherState>, OnDeviceProfileChangeListener {
 
     public static final FloatProperty<AllAppsTransitionController> ALL_APPS_PROGRESS =
             new FloatProperty<AllAppsTransitionController>("allAppsProgress") {
@@ -81,7 +78,6 @@
             };
 
     private static final int APPS_VIEW_ALPHA_CHANNEL_INDEX = 0;
-    private static final String PREF_KEY_SHOW_SEARCH_IME = "pref_search_show_ime";
 
     private AllAppsContainerView mAppsView;
     private ScrimView mScrimView;
@@ -99,8 +95,6 @@
     private float mProgress;        // [0, 1], mShiftRange * mProgress = shiftCurrent
 
     private float mScrollRangeDelta = 0;
-    private AllAppsInsetTransitionController mInsetController;
-    private boolean mSearchImeEnabled;
 
     public AllAppsTransitionController(Launcher l) {
         mLauncher = l;
@@ -109,19 +103,12 @@
 
         mIsVerticalLayout = mLauncher.getDeviceProfile().isVerticalBarLayout();
         mLauncher.addOnDeviceProfileChangeListener(this);
-
-        onSharedPreferenceChanged(mLauncher.getSharedPrefs(), PREF_KEY_SHOW_SEARCH_IME);
-        mLauncher.getSharedPrefs().registerOnSharedPreferenceChangeListener(this);
     }
 
     public float getShiftRange() {
         return mShiftRange;
     }
 
-    public AllAppsInsetTransitionController getInsetController() {
-        return mInsetController;
-    }
-
     @Override
     public void onDeviceProfileChanged(DeviceProfile dp) {
         mIsVerticalLayout = dp.isVerticalBarLayout();
@@ -146,14 +133,7 @@
         mProgress = progress;
 
         mScrimView.setProgress(progress);
-        float shiftCurrent = progress * mShiftRange;
-        mAppsView.setTranslationY(shiftCurrent);
-        if (FeatureFlags.ENABLE_DEVICE_SEARCH.get() && mSearchImeEnabled) {
-            if (mInsetController == null) {
-                setupInsetTransitionController();
-            }
-            mInsetController.setProgress(progress);
-        }
+        mAppsView.setTranslationY(progress * mShiftRange);
     }
 
     public float getProgress() {
@@ -242,18 +222,13 @@
     public void setupViews(AllAppsContainerView appsView, ScrimView scrimView) {
         mAppsView = appsView;
         mScrimView = scrimView;
-        if (FeatureFlags.ENABLE_DEVICE_SEARCH.get() && BuildCompat.isAtLeastR()) {
-            setupInsetTransitionController();
+        if (FeatureFlags.ENABLE_DEVICE_SEARCH.get() && Utilities.ATLEAST_R) {
+            mLauncher.getSystemUiController().updateUiState(UI_STATE_ALLAPPS,
+                    View.SYSTEM_UI_FLAG_LAYOUT_FULLSCREEN
+                            | View.SYSTEM_UI_FLAG_LAYOUT_STABLE);
         }
     }
 
-    private void setupInsetTransitionController() {
-        mInsetController = new AllAppsInsetTransitionController(mShiftRange, mAppsView);
-        mLauncher.getSystemUiController().updateUiState(UI_STATE_ALLAPPS,
-                View.SYSTEM_UI_FLAG_LAYOUT_FULLSCREEN
-                        | View.SYSTEM_UI_FLAG_LAYOUT_STABLE);
-    }
-
     /**
      * Updates the total scroll range but does not update the UI.
      */
@@ -274,23 +249,5 @@
         if (Float.compare(mProgress, 1f) == 0) {
             mAppsView.reset(false /* animate */);
         }
-        if (FeatureFlags.ENABLE_DEVICE_SEARCH.get() && mSearchImeEnabled
-                && BuildCompat.isAtLeastR()) {
-            mInsetController.onAnimationEnd(mProgress);
-            if (Float.compare(mProgress, 0f) == 0) {
-                EditText editText = mAppsView.getSearchUiManager().getEditText();
-                if (editText != null && !mInsetController.showSearchEduIfNecessary()) {
-                    editText.requestFocus();
-                }
-            }
-            // TODO: should make the controller hide synchronously
-        }
-    }
-
-    @Override
-    public void onSharedPreferenceChanged(SharedPreferences sharedPreferences, String s) {
-        if (s.equals(PREF_KEY_SHOW_SEARCH_IME)) {
-            mSearchImeEnabled = sharedPreferences.getBoolean(s, true);
-        }
     }
 }
diff --git a/src/com/android/launcher3/allapps/SearchUiManager.java b/src/com/android/launcher3/allapps/SearchUiManager.java
index 39410a7..0d42950 100644
--- a/src/com/android/launcher3/allapps/SearchUiManager.java
+++ b/src/com/android/launcher3/allapps/SearchUiManager.java
@@ -20,10 +20,10 @@
 import android.graphics.Rect;
 import android.view.KeyEvent;
 import android.view.animation.Interpolator;
-import android.widget.EditText;
 
 import androidx.annotation.Nullable;
 
+import com.android.launcher3.ExtendedEditText;
 import com.android.launcher3.anim.PropertySetter;
 
 /**
@@ -75,7 +75,7 @@
      * @return the edit text object
      */
     @Nullable
-    EditText getEditText();
+    ExtendedEditText getEditText();
 
     /**
      * sets highlight result's title
diff --git a/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java b/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java
index 426fd0c..2261d51 100644
--- a/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java
+++ b/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java
@@ -32,7 +32,6 @@
 import android.view.View;
 import android.view.ViewGroup.MarginLayoutParams;
 import android.view.animation.Interpolator;
-import android.widget.EditText;
 
 import com.android.launcher3.BaseDraggingActivity;
 import com.android.launcher3.DeviceProfile;
@@ -230,7 +229,7 @@
     }
 
     @Override
-    public EditText getEditText() {
+    public ExtendedEditText getEditText() {
         return this;
     }
 }
diff --git a/src/com/android/launcher3/statemanager/StateManager.java b/src/com/android/launcher3/statemanager/StateManager.java
index a18f340..2b51e97 100644
--- a/src/com/android/launcher3/statemanager/StateManager.java
+++ b/src/com/android/launcher3/statemanager/StateManager.java
@@ -26,14 +26,12 @@
 import android.animation.AnimatorSet;
 import android.os.Handler;
 import android.os.Looper;
-import android.util.Log;
 
 import com.android.launcher3.anim.AnimationSuccessListener;
 import com.android.launcher3.anim.AnimatorPlaybackController;
 import com.android.launcher3.anim.PendingAnimation;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.states.StateAnimationConfig.AnimationFlags;
-import com.android.launcher3.testing.TestProtocol;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
@@ -90,7 +88,9 @@
 
     public StateHandler[] getStateHandlers() {
         if (mStateHandlers == null) {
-            mStateHandlers = mActivity.createStateHandlers();
+            ArrayList<StateHandler> handlers = new ArrayList<>();
+            mActivity.collectStateHandlers(handlers);
+            mStateHandlers = handlers.toArray(new StateHandler[handlers.size()]);
         }
         return mStateHandlers;
     }
diff --git a/src/com/android/launcher3/statemanager/StatefulActivity.java b/src/com/android/launcher3/statemanager/StatefulActivity.java
index 7abb653..8a35cb3 100644
--- a/src/com/android/launcher3/statemanager/StatefulActivity.java
+++ b/src/com/android/launcher3/statemanager/StatefulActivity.java
@@ -30,6 +30,8 @@
 import com.android.launcher3.statemanager.StateManager.StateHandler;
 import com.android.launcher3.views.BaseDragLayer;
 
+import java.util.List;
+
 /**
  * Abstract activity with state management
  * @param <STATE_TYPE> Type of state object
@@ -46,7 +48,7 @@
     /**
      * Create handlers to control the property changes for this activity
      */
-    protected abstract StateHandler<STATE_TYPE>[] createStateHandlers();
+    protected abstract void collectStateHandlers(List<StateHandler> out);
 
     /**
      * Returns true if the activity is in the provided state
diff --git a/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java b/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java
index 6f1b2f9..516fc74 100644
--- a/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java
+++ b/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java
@@ -38,24 +38,19 @@
 import android.animation.AnimatorSet;
 import android.animation.ValueAnimator;
 import android.os.SystemClock;
-import android.util.Log;
 import android.view.HapticFeedbackConstants;
 import android.view.MotionEvent;
 
-import androidx.core.os.BuildCompat;
-
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAnimUtils;
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.AnimationSuccessListener;
 import com.android.launcher3.anim.AnimatorPlaybackController;
-import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.logger.LauncherAtom;
 import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.states.StateAnimationConfig.AnimationFlags;
-import com.android.launcher3.testing.TestProtocol;
 import com.android.launcher3.util.FlingBlockCheck;
 import com.android.launcher3.util.TouchController;
 
@@ -265,13 +260,6 @@
         }
         mCanBlockFling = mFromState == NORMAL;
         mFlingBlockCheck.unblockFling();
-        // Must be called after all the animation controllers have been paused
-        if (FeatureFlags.ENABLE_DEVICE_SEARCH.get()
-                && BuildCompat.isAtLeastR()
-                && (mToState == ALL_APPS || mToState == NORMAL)) {
-            mLauncher.getAllAppsController().getInsetController().onDragStart(
-                    mFromState == NORMAL ? 1f : 0f);
-        }
     }
 
     @Override