Snap for 9509315 from b207102b9891bba67931505361122a474fad428e to tm-qpr3-release

Change-Id: Icfd30f0f12d245bd102ec195ed846898fa4731e9
diff --git a/quickstep/src/com/android/launcher3/popup/QuickstepSystemShortcut.java b/quickstep/src/com/android/launcher3/popup/QuickstepSystemShortcut.java
index 9554bd3..45d1b11 100644
--- a/quickstep/src/com/android/launcher3/popup/QuickstepSystemShortcut.java
+++ b/quickstep/src/com/android/launcher3/popup/QuickstepSystemShortcut.java
@@ -101,7 +101,7 @@
             RecentsView recentsView = mTarget.getOverviewPanel();
             // Check if there is already an instance of this app running, if so, initiate the split
             // using that.
-            recentsView.findLastActiveTaskAndDoSplitOperation(
+            recentsView.findLastActiveTaskAndRunCallback(
                     intent.getComponent(),
                     (Consumer<Task>) foundTask -> {
                         SplitSelectSource source = new SplitSelectSource(mOriginalView,
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 4e795d9..5a81b2f 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -16,6 +16,7 @@
 package com.android.launcher3.taskbar;
 
 import static android.content.pm.PackageManager.FEATURE_PC;
+import static android.os.Trace.TRACE_TAG_APP;
 import static android.view.ViewGroup.LayoutParams.MATCH_PARENT;
 import static android.view.WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE;
 import static android.view.WindowManager.LayoutParams.LAYOUT_IN_DISPLAY_CUTOUT_MODE_ALWAYS;
@@ -46,6 +47,7 @@
 import android.graphics.Rect;
 import android.os.Process;
 import android.os.SystemProperties;
+import android.os.Trace;
 import android.provider.Settings;
 import android.util.Log;
 import android.view.Display;
@@ -94,6 +96,7 @@
 import com.android.launcher3.util.ViewCache;
 import com.android.launcher3.views.ActivityContext;
 import com.android.quickstep.views.RecentsView;
+import com.android.quickstep.views.TaskView;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.rotation.RotationButtonController;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
@@ -101,6 +104,7 @@
 import com.android.systemui.unfold.util.ScopedUnfoldTransitionProgressProvider;
 
 import java.io.PrintWriter;
+import java.util.function.Consumer;
 
 /**
  * The {@link ActivityContext} with which we inflate Taskbar-related Views. This allows UI elements
@@ -266,6 +270,13 @@
         dispatchDeviceProfileChanged();
     }
 
+    @Override
+    public void dispatchDeviceProfileChanged() {
+        super.dispatchDeviceProfileChanged();
+        Trace.instantForTrack(TRACE_TAG_APP, "TaskbarActivityContext#DeviceProfileChanged",
+                getDeviceProfile().toSmallString());
+    }
+
     /**
      * Copy the original DeviceProfile, match the number of hotseat icons and qsb width and update
      * the icon size
@@ -778,12 +789,12 @@
                 });
             });
         } else if (tag instanceof WorkspaceItemInfo) {
+            // Tapping a launchable icon on Taskbar
             WorkspaceItemInfo info = (WorkspaceItemInfo) tag;
             if (!info.isDisabled() || !ItemClickHandler.handleDisabledItemClicked(info, this)) {
                 TaskbarUIController taskbarUIController = mControllers.uiController;
                 RecentsView recents = taskbarUIController.getRecentsView();
-                if (recents != null
-                        && taskbarUIController.getRecentsView().isSplitSelectionActive()) {
+                if (recents != null && recents.isSplitSelectionActive()) {
                     // If we are selecting a second app for split, launch the split tasks
                     taskbarUIController.triggerSecondAppForSplit(info, info.intent, view);
                 } else {
@@ -810,7 +821,7 @@
                             getSystemService(LauncherApps.class)
                                     .startShortcut(packageName, id, null, null, info.user);
                         } else {
-                            startItemInfoActivity(info);
+                            launchFromTaskbarPreservingSplitIfVisible(recents, info);
                         }
 
                         mControllers.uiController.onTaskbarIconLaunched(info);
@@ -825,6 +836,7 @@
                 mControllers.taskbarStashController.updateAndAnimateTransientTaskbar(true);
             }
         } else if (tag instanceof AppInfo) {
+            // Tapping an item in AllApps
             AppInfo info = (AppInfo) tag;
             TaskbarUIController taskbarUIController = mControllers.uiController;
             RecentsView recents = taskbarUIController.getRecentsView();
@@ -833,9 +845,8 @@
                 // If we are selecting a second app for split, launch the split tasks
                 taskbarUIController.triggerSecondAppForSplit(info, info.intent, view);
             } else {
-                // Else launch the selected task
-                startItemInfoActivity((AppInfo) tag);
-                mControllers.uiController.onTaskbarIconLaunched((AppInfo) tag);
+                launchFromTaskbarPreservingSplitIfVisible(recents, info);
+                mControllers.uiController.onTaskbarIconLaunched(info);
             }
             mControllers.taskbarStashController.updateAndAnimateTransientTaskbar(true);
         } else if (tag instanceof ItemClickProxy) {
@@ -847,6 +858,31 @@
         AbstractFloatingView.closeAllOpenViews(this);
     }
 
+    /**
+     * Run when the user taps a Taskbar icon while in Overview. If the tapped app is currently
+     * visible to the user in Overview, or is part of a visible split pair, we expand the TaskView
+     * as if the user tapped on it (preserving the split pair). Otherwise, launch it normally
+     * (potentially breaking a split pair).
+     */
+    private void launchFromTaskbarPreservingSplitIfVisible(RecentsView recents, ItemInfo info) {
+        recents.findLastActiveTaskAndRunCallback(
+                info.getTargetComponent(),
+                (Consumer<Task>) foundTask -> {
+                    if (foundTask != null) {
+                        TaskView foundTaskView =
+                                recents.getTaskViewByTaskId(foundTask.key.id);
+                        if (foundTaskView != null
+                                && foundTaskView.isVisibleToUser()) {
+                            TestLogging.recordEvent(
+                                    TestProtocol.SEQUENCE_MAIN, "start: taskbarAppIcon");
+                            foundTaskView.launchTasks();
+                            return;
+                        }
+                    }
+                    startItemInfoActivity(info);
+                });
+    }
+
     private void startItemInfoActivity(ItemInfo info) {
         Intent intent = new Intent(info.getIntent())
                 .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
index 7b03746..ebb37a8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
@@ -170,7 +170,7 @@
      */
     public void triggerSecondAppForSplit(ItemInfoWithIcon info, Intent intent, View startingView) {
         RecentsView recents = getRecentsView();
-        recents.findLastActiveTaskAndDoSplitOperation(
+        recents.findLastActiveTaskAndRunCallback(
                 info.getTargetComponent(),
                 (Consumer<Task>) foundTask -> {
                     if (foundTask != null) {
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index a7651b6..078865f 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.uioverrides;
 
+import static android.os.Trace.TRACE_TAG_APP;
+import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_OPTIMIZE_MEASURE;
 import static android.view.accessibility.AccessibilityEvent.TYPE_VIEW_FOCUSED;
 
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
@@ -62,6 +64,7 @@
 import android.os.CancellationSignal;
 import android.os.IBinder;
 import android.os.SystemProperties;
+import android.os.Trace;
 import android.view.Display;
 import android.view.HapticFeedbackConstants;
 import android.view.RemoteAnimationTarget;
@@ -541,6 +544,7 @@
         if (FeatureFlags.CONTINUOUS_VIEW_TREE_CAPTURE.get()) {
             mViewCapture = ViewCapture.getInstance().startCapture(getWindow());
         }
+        getWindow().addPrivateFlags(PRIVATE_FLAG_OPTIMIZE_MEASURE);
     }
 
     @Override
@@ -1036,6 +1040,13 @@
         return false;
     }
 
+    @Override
+    public void dispatchDeviceProfileChanged() {
+        super.dispatchDeviceProfileChanged();
+        Trace.instantForTrack(TRACE_TAG_APP, "QuickstepLauncher#DeviceProfileChanged",
+                getDeviceProfile().toSmallString());
+    }
+
     private static final class LauncherTaskViewController extends
             TaskViewTouchController<Launcher> {
 
diff --git a/quickstep/src/com/android/quickstep/RecentsActivity.java b/quickstep/src/com/android/quickstep/RecentsActivity.java
index 6f86bf5..f26189c 100644
--- a/quickstep/src/com/android/quickstep/RecentsActivity.java
+++ b/quickstep/src/com/android/quickstep/RecentsActivity.java
@@ -15,6 +15,7 @@
  */
 package com.android.quickstep;
 
+import static android.os.Trace.TRACE_TAG_APP;
 import static android.view.RemoteAnimationTarget.MODE_CLOSING;
 import static android.view.RemoteAnimationTarget.MODE_OPENING;
 
@@ -36,6 +37,7 @@
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.Looper;
+import android.os.Trace;
 import android.view.Display;
 import android.view.RemoteAnimationAdapter;
 import android.view.RemoteAnimationTarget;
@@ -448,6 +450,13 @@
         return new RecentsAtomicAnimationFactory<>(this);
     }
 
+    @Override
+    public void dispatchDeviceProfileChanged() {
+        super.dispatchDeviceProfileChanged();
+        Trace.instantForTrack(TRACE_TAG_APP, "RecentsActivity#DeviceProfileChanged",
+                getDeviceProfile().toSmallString());
+    }
+
     private AnimatorListenerAdapter resetStateListener() {
         return new AnimatorListenerAdapter() {
             @Override
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index bb97334..00fb7ec 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -29,7 +29,6 @@
 import android.content.pm.ActivityInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.ShortcutInfo;
-import android.graphics.Insets;
 import android.graphics.Rect;
 import android.os.Bundle;
 import android.os.Handler;
@@ -51,10 +50,10 @@
 import androidx.annotation.WorkerThread;
 
 import com.android.internal.logging.InstanceId;
+import com.android.internal.util.ScreenshotRequest;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.SplitConfigurationOptions;
 import com.android.systemui.shared.recents.ISystemUiProxy;
-import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.system.smartspace.ILauncherUnlockAnimationController;
 import com.android.systemui.shared.system.smartspace.ISysuiUnlockAnimationController;
 import com.android.systemui.shared.system.smartspace.SmartspaceState;
@@ -384,14 +383,12 @@
     }
 
     @Override
-    public void handleImageBundleAsScreenshot(Bundle screenImageBundle, Rect locationInScreen,
-            Insets visibleInsets, Task.TaskKey task) {
+    public void takeScreenshot(ScreenshotRequest request) {
         if (mSystemUiProxy != null) {
             try {
-                mSystemUiProxy.handleImageBundleAsScreenshot(screenImageBundle, locationInScreen,
-                    visibleInsets, task);
+                mSystemUiProxy.takeScreenshot(request);
             } catch (RemoteException e) {
-                Log.w(TAG, "Failed call handleImageBundleAsScreenshot");
+                Log.w(TAG, "Failed call takeScreenshot");
             }
         }
     }
diff --git a/quickstep/src/com/android/quickstep/util/ImageActionUtils.java b/quickstep/src/com/android/quickstep/util/ImageActionUtils.java
index 9fe24de..3a1c99b 100644
--- a/quickstep/src/com/android/quickstep/util/ImageActionUtils.java
+++ b/quickstep/src/com/android/quickstep/util/ImageActionUtils.java
@@ -18,6 +18,8 @@
 
 import static android.content.Intent.FLAG_ACTIVITY_NEW_TASK;
 import static android.content.Intent.FLAG_GRANT_READ_URI_PERMISSION;
+import static android.view.WindowManager.ScreenshotSource.SCREENSHOT_OVERVIEW;
+import static android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE;
 
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.THREAD_POOL_EXECUTOR;
@@ -47,7 +49,7 @@
 import androidx.core.content.FileProvider;
 
 import com.android.internal.app.ChooserActivity;
-import com.android.internal.util.ScreenshotHelper;
+import com.android.internal.util.ScreenshotRequest;
 import com.android.launcher3.BuildConfig;
 import com.android.quickstep.SystemUiProxy;
 import com.android.systemui.shared.recents.model.Task;
@@ -74,11 +76,17 @@
      * Saves screenshot to location determine by SystemUiProxy
      */
     public static void saveScreenshot(SystemUiProxy systemUiProxy, Bitmap screenshot,
-            Rect screenshotBounds,
-            Insets visibleInsets, Task.TaskKey task) {
-        systemUiProxy.handleImageBundleAsScreenshot(
-                ScreenshotHelper.HardwareBitmapBundler.hardwareBitmapToBundle(screenshot),
-                screenshotBounds, visibleInsets, task);
+            Rect screenshotBounds, Insets visibleInsets, Task.TaskKey task) {
+        ScreenshotRequest request =
+                new ScreenshotRequest.Builder(TAKE_SCREENSHOT_PROVIDED_IMAGE, SCREENSHOT_OVERVIEW)
+                .setTopComponent(task.sourceComponent)
+                .setTaskId(task.id)
+                .setUserId(task.userId)
+                .setBitmap(screenshot)
+                .setBoundsOnScreen(screenshotBounds)
+                .setInsets(visibleInsets)
+                .build();
+        systemUiProxy.takeScreenshot(request);
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index c11f7ed..5fa2a5c 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1275,14 +1275,14 @@
     }
 
     /**
-     * Pulls the list of active Tasks from RecentModel, and finds the most recently active Task
+     * Pulls the list of active Tasks from RecentsModel, and finds the most recently active Task
      * matching a given ComponentName. Then uses that Task (which could be null) with the given
      * callback.
      *
-     * Used in various splitscreen operations when we need to check if there is a currently running
-     * Task of a certain type and use the most recent one.
+     * Used in various task-switching or splitscreen operations when we need to check if there is a
+     * currently running Task of a certain type and use the most recent one.
      */
-    public void findLastActiveTaskAndDoSplitOperation(ComponentName componentName,
+    public void findLastActiveTaskAndRunCallback(ComponentName componentName,
             Consumer<Task> callback) {
         mModel.getTasks(taskGroups -> {
             Task lastActiveTask = null;
diff --git a/src/com/android/launcher3/DeviceProfile.java b/src/com/android/launcher3/DeviceProfile.java
index 7b540d0..8391b91 100644
--- a/src/com/android/launcher3/DeviceProfile.java
+++ b/src/com/android/launcher3/DeviceProfile.java
@@ -1680,6 +1680,16 @@
         writer.println(prefix + pxToDpStr("getCellLayoutWidth()", getCellLayoutWidth()));
     }
 
+    /** Returns a reduced representation of this DeviceProfile. */
+    public String toSmallString() {
+        return "isTablet:" + isTablet + ", "
+                + "isMultiDisplay:" + isMultiDisplay + ", "
+                + "widthPx:" + widthPx + ", "
+                + "heightPx:" + heightPx + ", "
+                + "insets:" + mInsets + ", "
+                + "rotationHint:" + rotationHint;
+    }
+
     private static Context getContext(Context c, Info info, int orientation, WindowBounds bounds) {
         Configuration config = new Configuration(c.getResources().getConfiguration());
         config.orientation = orientation;
diff --git a/src/com/android/launcher3/allapps/AllAppsTransitionController.java b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
index 6f6f86b..f7b0d96 100644
--- a/src/com/android/launcher3/allapps/AllAppsTransitionController.java
+++ b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
@@ -25,6 +25,9 @@
 import static com.android.launcher3.anim.PropertySetter.NO_ANIM_PROPERTY_SETTER;
 import static com.android.launcher3.states.StateAnimationConfig.ANIM_ALL_APPS_FADE;
 import static com.android.launcher3.states.StateAnimationConfig.ANIM_VERTICAL_PROGRESS;
+import static com.android.launcher3.util.SystemUiController.FLAG_DARK_NAV;
+import static com.android.launcher3.util.SystemUiController.FLAG_LIGHT_NAV;
+import static com.android.launcher3.util.SystemUiController.UI_STATE_ALL_APPS;
 
 import android.animation.Animator;
 import android.animation.Animator.AnimatorListener;
@@ -54,6 +57,7 @@
 import com.android.launcher3.util.MultiPropertyFactory;
 import com.android.launcher3.util.MultiPropertyFactory.MultiProperty;
 import com.android.launcher3.util.MultiValueAlpha;
+import com.android.launcher3.util.Themes;
 import com.android.launcher3.views.ScrimView;
 
 /**
@@ -73,6 +77,8 @@
     public static final float SWIPE_ALL_APPS_TO_HOME_MIN_SCALE = 0.9f;
     private static final int REVERT_SWIPE_ALL_APPS_TO_HOME_ANIMATION_DURATION_MS = 200;
 
+    private static final float NAV_BAR_COLOR_FORCE_UPDATE_THRESHOLD = 0.1f;
+
     public static final FloatProperty<AllAppsTransitionController> ALL_APPS_PROGRESS =
             new FloatProperty<AllAppsTransitionController>("allAppsProgress") {
 
@@ -151,6 +157,8 @@
 
     private final Launcher mLauncher;
     private final AnimatedFloat mAllAppScale = new AnimatedFloat(this::onScaleProgressChanged);
+    private final int mNavScrimFlag;
+
     private boolean mIsVerticalLayout;
 
     // Whether this class should take care of closing the keyboard.
@@ -177,10 +185,13 @@
     public AllAppsTransitionController(Launcher l) {
         mLauncher = l;
         DeviceProfile dp = mLauncher.getDeviceProfile();
-        setShiftRange(dp.allAppsShiftRange);
         mProgress = 1f;
         mIsVerticalLayout = dp.isVerticalBarLayout();
         mIsTablet = dp.isTablet;
+        mNavScrimFlag = Themes.getAttrBoolean(l, R.attr.isMainColorDark)
+                ? FLAG_DARK_NAV : FLAG_LIGHT_NAV;
+
+        setShiftRange(dp.allAppsShiftRange);
         mLauncher.addOnDeviceProfileChangeListener(this);
     }
 
@@ -213,6 +224,11 @@
         mProgress = progress;
         getAppsViewProgressTranslationY().setValue(mProgress * mShiftRange);
         mLauncher.onAllAppsTransition(1 - progress);
+
+        boolean hasScrim = progress < NAV_BAR_COLOR_FORCE_UPDATE_THRESHOLD
+                && mLauncher.getAppsView().getNavBarScrimHeight() > 0;
+        mLauncher.getSystemUiController().updateUiState(
+                UI_STATE_ALL_APPS, hasScrim ? mNavScrimFlag : 0);
     }
 
     public float getProgress() {
diff --git a/src/com/android/launcher3/allapps/BaseAllAppsContainerView.java b/src/com/android/launcher3/allapps/BaseAllAppsContainerView.java
index 1c67691..9bb8250 100644
--- a/src/com/android/launcher3/allapps/BaseAllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/BaseAllAppsContainerView.java
@@ -520,13 +520,20 @@
     /**
      * Returns a padding in case a scrim is shown on the bottom of the view and a padding is needed.
      */
-    protected int getNavBarScrimHeight(WindowInsets insets) {
+    protected int computeNavBarScrimHeight(WindowInsets insets) {
         return 0;
     }
 
+    /**
+     * Returns the current height of nav bar scrim
+     */
+    public int getNavBarScrimHeight() {
+        return mNavBarScrimHeight;
+    }
+
     @Override
     public WindowInsets dispatchApplyWindowInsets(WindowInsets insets) {
-        mNavBarScrimHeight = getNavBarScrimHeight(insets);
+        mNavBarScrimHeight = computeNavBarScrimHeight(insets);
         applyAdapterSideAndBottomPaddings(mActivityContext.getDeviceProfile());
         return super.dispatchApplyWindowInsets(insets);
     }
diff --git a/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java b/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
index 5a5ba2b..aefedae 100644
--- a/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
@@ -41,7 +41,7 @@
     }
 
     @Override
-    protected int getNavBarScrimHeight(WindowInsets insets) {
+    protected int computeNavBarScrimHeight(WindowInsets insets) {
         if (Utilities.ATLEAST_Q) {
             return insets.getTappableElementInsets().bottom;
         } else {
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 6fe0ba7..c328554 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -385,6 +385,11 @@
             "ENABLE_MULTI_INSTANCE", false,
             "Enables creation and filtering of multiple task instances in overview");
 
+    public static final BooleanFlag ENABLE_TASKBAR_PINNING = getDebugFlag(
+            "ENABLE_TASKBAR_PINNING", false,
+            "Enables taskbar pinning to allow user to switch between transient and persistent "
+                    + "taskbar flavors");
+
     public static void initialize(Context context) {
         synchronized (sDebugFlags) {
             for (DebugFlag flag : sDebugFlags) {
diff --git a/src/com/android/launcher3/util/SystemUiController.java b/src/com/android/launcher3/util/SystemUiController.java
index 6945983..df54fd7 100644
--- a/src/com/android/launcher3/util/SystemUiController.java
+++ b/src/com/android/launcher3/util/SystemUiController.java
@@ -35,6 +35,7 @@
     public static final int UI_STATE_SCRIM_VIEW = 1;
     public static final int UI_STATE_WIDGET_BOTTOM_SHEET = 2;
     public static final int UI_STATE_FULLSCREEN_TASK = 3;
+    public static final int UI_STATE_ALL_APPS = 4;
 
     public static final int FLAG_LIGHT_NAV = 1 << 0;
     public static final int FLAG_DARK_NAV = 1 << 1;
@@ -54,7 +55,7 @@
     public @interface SystemUiControllerFlags {}
 
     private final Window mWindow;
-    private final int[] mStates = new int[4];
+    private final int[] mStates = new int[5];
 
     public SystemUiController(Window window) {
         mWindow = window;
diff --git a/tests/tapl/com/android/launcher3/tapl/AllAppsQsb.java b/tests/tapl/com/android/launcher3/tapl/AllAppsQsb.java
new file mode 100644
index 0000000..0931cd4
--- /dev/null
+++ b/tests/tapl/com/android/launcher3/tapl/AllAppsQsb.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2022 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.tapl;
+
+import androidx.test.uiautomator.UiObject2;
+
+/**
+ * Operations on AllApp screen qsb.
+ */
+class AllAppsQsb extends Qsb {
+
+    private final UiObject2 mAllAppsContainer;
+
+    AllAppsQsb(LauncherInstrumentation launcher, UiObject2 allAppsContainer) {
+        super(launcher);
+        mAllAppsContainer = allAppsContainer;
+        waitForQsbObject();
+    }
+
+    @Override
+    protected UiObject2 waitForQsbObject() {
+        return mLauncher.waitForObjectInContainer(mAllAppsContainer, "search_container_all_apps");
+    }
+}
diff --git a/tests/tapl/com/android/launcher3/tapl/HomeAllApps.java b/tests/tapl/com/android/launcher3/tapl/HomeAllApps.java
index 50b03aa..8ac1aef 100644
--- a/tests/tapl/com/android/launcher3/tapl/HomeAllApps.java
+++ b/tests/tapl/com/android/launcher3/tapl/HomeAllApps.java
@@ -115,4 +115,13 @@
             }
         }
     }
+
+    /**
+     * Return the QSB UI object on the AllApps screen.
+     * @return the QSB UI object.
+     */
+    @NonNull
+    public Qsb getQsb() {
+        return new AllAppsQsb(mLauncher, verifyActiveContainer());
+    }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/HomeQsb.java b/tests/tapl/com/android/launcher3/tapl/HomeQsb.java
index c365708..20d09a1 100644
--- a/tests/tapl/com/android/launcher3/tapl/HomeQsb.java
+++ b/tests/tapl/com/android/launcher3/tapl/HomeQsb.java
@@ -15,69 +15,23 @@
  */
 package com.android.launcher3.tapl;
 
-import androidx.annotation.NonNull;
-import androidx.test.uiautomator.By;
-import androidx.test.uiautomator.BySelector;
 import androidx.test.uiautomator.UiObject2;
-import androidx.test.uiautomator.Until;
 
 /**
- * Operations on home screen qsb.
+ * Operations on Home screen qsb.
  */
-public class HomeQsb {
+class HomeQsb extends Qsb {
 
-    private final LauncherInstrumentation mLauncher;
-    private static final String ASSISTANT_APP_PACKAGE = "com.google.android.googlequicksearchbox";
-    private static final String ASSISTANT_ICON_RES_ID = "mic_icon";
+    private final UiObject2 mHotSeat;
 
-
-    HomeQsb(LauncherInstrumentation launcher) {
-        mLauncher = launcher;
-        mLauncher.waitForLauncherObject("search_container_hotseat");
+    HomeQsb(LauncherInstrumentation launcher, UiObject2 hotseat) {
+        super(launcher);
+        mHotSeat = hotseat;
+        waitForQsbObject();
     }
 
-    /**
-     * Launch assistant app by tapping mic icon on qsb.
-     */
-    @NonNull
-    public LaunchedAppState launchAssistant() {
-        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
-                "want to click assistant mic icon button");
-             LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
-            UiObject2 assistantIcon = mLauncher.waitForLauncherObject(ASSISTANT_ICON_RES_ID);
-
-            LauncherInstrumentation.log("HomeQsb.launchAssistant before click "
-                    + assistantIcon.getVisibleCenter() + " in "
-                    + mLauncher.getVisibleBounds(assistantIcon));
-
-            mLauncher.clickLauncherObject(assistantIcon);
-
-            try (LauncherInstrumentation.Closable c2 = mLauncher.addContextLayer("clicked")) {
-                // assert Assistant App Launched
-                BySelector selector = By.pkg(ASSISTANT_APP_PACKAGE);
-                mLauncher.assertTrue(
-                        "assistant app didn't start: (" + selector + ")",
-                        mLauncher.getDevice().wait(Until.hasObject(selector),
-                                LauncherInstrumentation.WAIT_TIME_MS)
-                );
-                return new LaunchedAppState(mLauncher);
-            }
-        }
-    }
-
-    /**
-     * Show search result page from tapping qsb.
-     */
-    public SearchResultFromQsb showSearchResult() {
-        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
-                "want to open search result page");
-             LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
-            mLauncher.clickLauncherObject(
-                    mLauncher.waitForLauncherObject("search_container_hotseat"));
-            try (LauncherInstrumentation.Closable c2 = mLauncher.addContextLayer(
-                    "clicked qsb to open search result page")) {
-                return new SearchResultFromQsb(mLauncher);
-            }
-        }
+    @Override
+    protected UiObject2 waitForQsbObject() {
+        return mLauncher.waitForObjectInContainer(mHotSeat, "search_container_hotseat");
     }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/Qsb.java b/tests/tapl/com/android/launcher3/tapl/Qsb.java
new file mode 100644
index 0000000..6bc4f21
--- /dev/null
+++ b/tests/tapl/com/android/launcher3/tapl/Qsb.java
@@ -0,0 +1,86 @@
+/*
+ * Copyright (C) 2022 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.tapl;
+
+import androidx.annotation.NonNull;
+import androidx.test.uiautomator.By;
+import androidx.test.uiautomator.BySelector;
+import androidx.test.uiautomator.UiObject2;
+import androidx.test.uiautomator.Until;
+
+/**
+ * Operations on qsb from either Home screen or AllApp screen.
+ */
+public abstract class Qsb {
+
+    private static final String ASSISTANT_APP_PACKAGE = "com.google.android.googlequicksearchbox";
+    private static final String ASSISTANT_ICON_RES_ID = "mic_icon";
+    protected final LauncherInstrumentation mLauncher;
+
+    protected Qsb(LauncherInstrumentation launcher) {
+        mLauncher = launcher;
+    }
+
+    // Waits for the quick search box.
+    protected abstract UiObject2 waitForQsbObject();
+    /**
+     * Launch assistant app by tapping mic icon on qsb.
+     */
+
+    @NonNull
+    public LaunchedAppState launchAssistant() {
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to click assistant mic icon button");
+             LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
+            UiObject2 assistantIcon = mLauncher.waitForLauncherObject(ASSISTANT_ICON_RES_ID);
+
+            LauncherInstrumentation.log("Qsb.launchAssistant before click "
+                    + assistantIcon.getVisibleCenter() + " in "
+                    + mLauncher.getVisibleBounds(assistantIcon));
+
+            mLauncher.clickLauncherObject(assistantIcon);
+
+            try (LauncherInstrumentation.Closable c2 = mLauncher.addContextLayer("clicked")) {
+                // assert Assistant App Launched
+                BySelector selector = By.pkg(ASSISTANT_APP_PACKAGE);
+                mLauncher.assertTrue(
+                        "assistant app didn't start: (" + selector + ")",
+                        mLauncher.getDevice().wait(Until.hasObject(selector),
+                                LauncherInstrumentation.WAIT_TIME_MS)
+                );
+                return new LaunchedAppState(mLauncher);
+            }
+        }
+    }
+
+    /**
+     * Show search result page from tapping qsb.
+     */
+    public SearchResultFromQsb showSearchResult() {
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to open search result page");
+             LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
+            mLauncher.clickLauncherObject(waitForQsbObject());
+            // wait for the result rendering to complete
+            mLauncher.waitForIdle();
+            try (LauncherInstrumentation.Closable c2 = mLauncher.addContextLayer(
+                    "clicked qsb to open search result page")) {
+                return new SearchResultFromQsb(mLauncher);
+            }
+        }
+    }
+}
diff --git a/tests/tapl/com/android/launcher3/tapl/SearchResultFromQsb.java b/tests/tapl/com/android/launcher3/tapl/SearchResultFromQsb.java
index 80e4116..80176e9 100644
--- a/tests/tapl/com/android/launcher3/tapl/SearchResultFromQsb.java
+++ b/tests/tapl/com/android/launcher3/tapl/SearchResultFromQsb.java
@@ -23,7 +23,7 @@
 import java.util.ArrayList;
 
 /**
- * Operations on search result page opened from home screen qsb.
+ * Operations on search result page opened from qsb.
  */
 public class SearchResultFromQsb {
     // The input resource id in the search box.
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index 2c82c50..388955c 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -118,10 +118,11 @@
      *
      * The qsb must already be visible when calling this method.
      */
-    public HomeQsb getQsb() {
+    @NonNull
+    public Qsb getQsb() {
         try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
                 "want to get the home qsb")) {
-            return new HomeQsb(mLauncher);
+            return new HomeQsb(mLauncher, mHotseat);
         }
     }