diff --git a/Android.bp b/Android.bp
index 9d675a4..cca48ce 100644
--- a/Android.bp
+++ b/Android.bp
@@ -12,6 +12,23 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+package {
+    default_applicable_licenses: ["packages_apps_Launcher3_license"],
+}
+
+// Added automatically by a large-scale-change
+// See: http://go/android-license-faq
+license {
+    name: "packages_apps_Launcher3_license",
+    visibility: [":__subpackages__"],
+    license_kinds: [
+        "SPDX-license-identifier-Apache-2.0",
+    ],
+    license_text: [
+        "NOTICE",
+    ],
+}
+
 android_library {
     name: "launcher-aosp-tapl",
     static_libs: [
diff --git a/Android.mk b/Android.mk
index 127df79..304935b 100644
--- a/Android.mk
+++ b/Android.mk
@@ -46,6 +46,9 @@
 LOCAL_SDK_VERSION := current
 LOCAL_MIN_SDK_VERSION := 26
 LOCAL_MODULE := Launcher3CommonDepsLib
+LOCAL_LICENSE_KINDS := SPDX-license-identifier-Apache-2.0
+LOCAL_LICENSE_CONDITIONS := notice
+LOCAL_NOTICE_FILE := $(LOCAL_PATH)/NOTICE
 LOCAL_PRIVILEGED_MODULE := true
 LOCAL_MANIFEST_FILE := AndroidManifest-common.xml
 
@@ -137,6 +140,9 @@
   LOCAL_MIN_SDK_VERSION := 26
 endif
 LOCAL_MODULE := Launcher3QuickStepLib
+LOCAL_LICENSE_KINDS := SPDX-license-identifier-Apache-2.0
+LOCAL_LICENSE_CONDITIONS := notice
+LOCAL_NOTICE_FILE := $(LOCAL_PATH)/NOTICE
 LOCAL_PRIVILEGED_MODULE := true
 LOCAL_STATIC_ANDROID_LIBRARIES := Launcher3CommonDepsLib
 
diff --git a/AndroidManifest.xml b/AndroidManifest.xml
index 56a2595..1fad72d 100644
--- a/AndroidManifest.xml
+++ b/AndroidManifest.xml
@@ -49,7 +49,7 @@
             android:stateNotNeeded="true"
             android:windowSoftInputMode="adjustPan"
             android:screenOrientation="unspecified"
-            android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout|smallestScreenSize"
+            android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout"
             android:resizeableActivity="true"
             android:resumeWhilePausing="true"
             android:taskAffinity=""
diff --git a/go/AndroidManifest-launcher.xml b/go/AndroidManifest-launcher.xml
index 2223036..d2575b6 100644
--- a/go/AndroidManifest-launcher.xml
+++ b/go/AndroidManifest-launcher.xml
@@ -49,7 +49,7 @@
             android:stateNotNeeded="true"
             android:windowSoftInputMode="adjustPan"
             android:screenOrientation="unspecified"
-            android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout|smallestScreenSize|uiMode"
+            android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout|uiMode"
             android:resizeableActivity="true"
             android:resumeWhilePausing="true"
             android:taskAffinity=""
diff --git a/quickstep/AndroidManifest-launcher.xml b/quickstep/AndroidManifest-launcher.xml
index 53910e3..7fe9b08 100644
--- a/quickstep/AndroidManifest-launcher.xml
+++ b/quickstep/AndroidManifest-launcher.xml
@@ -49,7 +49,7 @@
             android:stateNotNeeded="true"
             android:windowSoftInputMode="adjustPan"
             android:screenOrientation="unspecified"
-            android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout|smallestScreenSize"
+            android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout"
             android:resizeableActivity="true"
             android:resumeWhilePausing="true"
             android:taskAffinity=""
diff --git a/quickstep/AndroidManifest.xml b/quickstep/AndroidManifest.xml
index 97f4a21..5e5cf73 100644
--- a/quickstep/AndroidManifest.xml
+++ b/quickstep/AndroidManifest.xml
@@ -71,7 +71,7 @@
              android:stateNotNeeded="true"
              android:theme="@style/LauncherTheme"
              android:screenOrientation="unspecified"
-             android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout|smallestScreenSize"
+             android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout"
              android:resizeableActivity="true"
              android:resumeWhilePausing="true"
              android:taskAffinity=""/>
diff --git a/quickstep/src/com/android/launcher3/appprediction/AppsDividerView.java b/quickstep/src/com/android/launcher3/appprediction/AppsDividerView.java
index b891120..66b1a86 100644
--- a/quickstep/src/com/android/launcher3/appprediction/AppsDividerView.java
+++ b/quickstep/src/com/android/launcher3/appprediction/AppsDividerView.java
@@ -42,7 +42,6 @@
 import com.android.launcher3.allapps.FloatingHeaderRow;
 import com.android.launcher3.allapps.FloatingHeaderView;
 import com.android.launcher3.anim.PropertySetter;
-import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.util.Themes;
 
@@ -183,11 +182,6 @@
     }
 
     private void updateViewVisibility() {
-        // hide divider since we have item decoration for prediction row
-        if (FeatureFlags.ENABLE_DEVICE_SEARCH.get()) {
-            setVisibility(GONE);
-            return;
-        }
         setVisibility(mDividerType == DividerType.NONE
                 ? GONE
                 : (mIsScrolledOut ? INVISIBLE : VISIBLE));
diff --git a/quickstep/src/com/android/launcher3/appprediction/PredictionRowView.java b/quickstep/src/com/android/launcher3/appprediction/PredictionRowView.java
index 7d494c2..b570d55 100644
--- a/quickstep/src/com/android/launcher3/appprediction/PredictionRowView.java
+++ b/quickstep/src/com/android/launcher3/appprediction/PredictionRowView.java
@@ -108,7 +108,8 @@
 
     AllAppsSectionDecorator.SectionDecorationHandler mDecorationHandler;
 
-    @Nullable private List<ItemInfo> mPendingPredictedItems;
+    @Nullable
+    private List<ItemInfo> mPendingPredictedItems;
 
     public PredictionRowView(@NonNull Context context) {
         this(context, null);
@@ -181,7 +182,7 @@
 
     @Override
     public boolean shouldDraw() {
-        return getVisibility() == VISIBLE;
+        return getVisibility() != GONE;
     }
 
     @Override
@@ -189,6 +190,11 @@
         return mPredictionsEnabled;
     }
 
+    @Override
+    public boolean isVisible() {
+        return getVisibility() == VISIBLE;
+    }
+
     /**
      * Returns the predicted apps.
      */
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarVisibilityController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarVisibilityController.java
index 4cf55d8..6d20d97 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarVisibilityController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarVisibilityController.java
@@ -59,6 +59,9 @@
         boolean isImeVisible = (SystemUiProxy.INSTANCE.get(mLauncher).getLastSystemUiStateFlags()
                 & QuickStepContract.SYSUI_STATE_IME_SHOWING) != 0;
         mTaskbarVisibilityAlphaForIme.updateValue(isImeVisible ? 0f : 1f);
+
+        onTaskbarBackgroundAlphaChanged();
+        updateVisibilityAlpha();
     }
 
     protected void cleanup() {
diff --git a/quickstep/src/com/android/launcher3/uioverrides/DeviceFlag.java b/quickstep/src/com/android/launcher3/uioverrides/DeviceFlag.java
index 010694b..bb1f6fc 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/DeviceFlag.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/DeviceFlag.java
@@ -60,6 +60,14 @@
         mListeners.add(r);
     }
 
+    @Override
+    public void removeChangeListener(Runnable r) {
+        if (mListeners == null) {
+            return;
+        }
+        mListeners.remove(r);
+    }
+
     private void registerDeviceConfigChangedListener(Context context) {
         DeviceConfig.addOnPropertiesChangedListener(
                 NAMESPACE_LAUNCHER,
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 808f26e..5f8fc6d 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -1748,7 +1748,9 @@
         if (mWindowTransitionController != null) {
             mWindowTransitionController.setProgress(mCurrentShift.value, mDragLengthFactor);
         }
-        if (mRecentsAnimationTargets != null) {
+        // No need to apply any transform if there is ongoing swipe-pip-to-home animator since
+        // that animator handles the leash solely.
+        if (mRecentsAnimationTargets != null && !mIsSwipingPipToHome) {
             if (mRecentsViewScrollLinked) {
                 mTaskViewSimulator.setScroll(mRecentsView.getScrollOffset());
             }
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index e243715..e4c8b6f 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -37,18 +37,23 @@
 import android.content.Intent;
 import android.content.SharedPreferences;
 import android.content.res.Configuration;
+import android.graphics.Point;
 import android.graphics.Rect;
 import android.graphics.Region;
 import android.graphics.drawable.Icon;
+import android.hardware.display.DisplayManager;
 import android.os.Build;
 import android.os.Bundle;
 import android.os.IBinder;
 import android.os.Looper;
 import android.os.SystemClock;
+import android.os.SystemProperties;
 import android.util.Log;
 import android.view.Choreographer;
+import android.view.Display;
 import android.view.InputEvent;
 import android.view.MotionEvent;
+import android.view.Surface;
 import android.view.accessibility.AccessibilityManager;
 
 import androidx.annotation.BinderThread;
@@ -91,6 +96,7 @@
 import com.android.systemui.shared.recents.IOverviewProxy;
 import com.android.systemui.shared.recents.ISystemUiProxy;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
+import com.android.systemui.shared.system.InputChannelCompat;
 import com.android.systemui.shared.system.InputChannelCompat.InputEventReceiver;
 import com.android.systemui.shared.system.InputConsumerController;
 import com.android.systemui.shared.system.InputMonitorCompat;
@@ -121,6 +127,9 @@
      */
     private static final int SYSTEM_ACTION_ID_ALL_APPS = 14;
 
+    public static final boolean ENABLE_PER_WINDOW_INPUT_ROTATION =
+            SystemProperties.getBoolean("persist.debug.per_window_input_rotation", false);
+
     private int mBackGestureNotificationCounter = -1;
     @Nullable
     private OverscrollPlugin mOverscrollPlugin;
@@ -248,6 +257,8 @@
     private InputMonitorCompat mInputMonitorCompat;
     private InputEventReceiver mInputEventReceiver;
 
+    private DisplayManager mDisplayManager;
+
     @Override
     public void onCreate() {
         super.onCreate();
@@ -261,6 +272,7 @@
         mDeviceState.addOneHandedModeChangedCallback(this::onOneHandedModeOverlayChanged);
         mDeviceState.runOnUserUnlocked(this::onUserUnlocked);
         ProtoTracer.INSTANCE.get(this).add(this);
+        mDisplayManager = getSystemService(DisplayManager.class);
 
         sConnected = true;
     }
@@ -426,6 +438,15 @@
             return;
         }
         MotionEvent event = (MotionEvent) ev;
+        if (ENABLE_PER_WINDOW_INPUT_ROTATION) {
+            final Display display = mDisplayManager.getDisplay(mDeviceState.getDisplayId());
+            int rotation = display.getRotation();
+            Point sz = new Point();
+            display.getRealSize(sz);
+            if (rotation != Surface.ROTATION_0) {
+                event.transform(InputChannelCompat.createRotationMatrix(rotation, sz.x, sz.y));
+            }
+        }
 
         TestLogging.recordMotionEvent(
                 TestProtocol.SEQUENCE_TIS, "TouchInteractionService.onInputEvent", event);
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
index f2f9fb7..5baf518 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
@@ -38,10 +38,12 @@
 import android.content.ContextWrapper;
 import android.content.Intent;
 import android.graphics.PointF;
+import android.hardware.display.DisplayManager;
 import android.os.Build;
 import android.os.Handler;
 import android.os.Looper;
 import android.util.Log;
+import android.view.Display;
 import android.view.MotionEvent;
 import android.view.VelocityTracker;
 import android.view.ViewConfiguration;
@@ -64,6 +66,7 @@
 import com.android.quickstep.RotationTouchHelper;
 import com.android.quickstep.TaskAnimationManager;
 import com.android.quickstep.TaskUtils;
+import com.android.quickstep.TouchInteractionService;
 import com.android.quickstep.util.ActiveGestureLog;
 import com.android.quickstep.util.CachedEventDispatcher;
 import com.android.quickstep.util.MotionPauseDetector;
@@ -113,6 +116,8 @@
     private final PointF mLastPos = new PointF();
     private int mActivePointerId = INVALID_POINTER_ID;
 
+    private int mLastRotation = -1;
+
     // Distance after which we start dragging the window.
     private final float mTouchSlop;
 
@@ -130,6 +135,8 @@
     // Might be displacement in X or Y, depending on the direction we are swiping from the nav bar.
     private float mStartDisplacement;
 
+    private final DisplayManager mDisplayManager;
+
     private Handler mMainThreadHandler;
     private Runnable mCancelRecentsAnimationRunnable = () -> {
         ActivityManagerWrapper.getInstance().cancelRecentsAnimation(
@@ -172,6 +179,7 @@
         mPassedPilferInputSlop = mPassedWindowMoveSlop = continuingPreviousGesture;
         mDisableHorizontalSwipe = !mPassedPilferInputSlop && disableHorizontalSwipe;
         mRotationTouchHelper = mDeviceState.getRotationTouchHelper();
+        mDisplayManager = getSystemService(DisplayManager.class);
     }
 
     @Override
@@ -197,6 +205,17 @@
             return;
         }
 
+        if (TouchInteractionService.ENABLE_PER_WINDOW_INPUT_ROTATION) {
+            final Display display = mDisplayManager.getDisplay(mDeviceState.getDisplayId());
+            final int rotation = display.getRotation();
+            if (rotation != mLastRotation) {
+                // If rotation changes, reset tracking to avoid degenerate velocities.
+                mLastPos.set(ev.getX(), ev.getY());
+                mVelocityTracker.clear();
+                mLastRotation = rotation;
+            }
+        }
+
         // Proxy events to recents view
         if (mPassedWindowMoveSlop && mInteractionHandler != null
                 && !mRecentsViewDispatcher.hasConsumer()) {
diff --git a/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java b/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java
index 6630aed..25ae055 100644
--- a/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java
+++ b/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java
@@ -71,6 +71,7 @@
     private ViewOutlineProvider mOldBannerOutlineProvider;
     private float mBannerOffsetPercentage;
     private float mBannerAlpha = 1f;
+    private float mVerticalOffset = 0f;
 
     public DigitalWellBeingToast(BaseDraggingActivity activity, TaskView taskView) {
         mActivity = activity;
@@ -275,16 +276,17 @@
             @Override
             public void getOutline(View view, Outline outline) {
                 mOldBannerOutlineProvider.getOutline(view, outline);
-                outline.offset(0, -Math.round(view.getTranslationY()));
+                outline.offset(0, Math.round(-view.getTranslationY() + mVerticalOffset));
             }
         });
         mBanner.setClipToOutline(true);
     }
 
-    void updateBannerOffset(float offsetPercentage) {
+    void updateBannerOffset(float offsetPercentage, float verticalOffset) {
         if (mBanner != null && mBannerOffsetPercentage != offsetPercentage) {
+            mVerticalOffset = verticalOffset;
             mBannerOffsetPercentage = offsetPercentage;
-            mBanner.setTranslationY(offsetPercentage * mBanner.getHeight());
+            mBanner.setTranslationY(offsetPercentage * mBanner.getHeight() + mVerticalOffset);
             mBanner.invalidateOutline();
         }
     }
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 361c545..b72e05c 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -537,7 +537,7 @@
     @Override
     protected void onWindowVisibilityChanged(int visibility) {
         super.onWindowVisibilityChanged(visibility);
-        if (visibility == GONE && ENABLE_QUICKSTEP_LIVE_TILE.get()) {
+        if (visibility != VISIBLE && ENABLE_QUICKSTEP_LIVE_TILE.get()) {
             finishRecentsAnimation(true /* toRecents */, null);
         }
         updateTaskStackListenerState();
@@ -1752,7 +1752,7 @@
         if (alpha > 0) {
             setVisibility(VISIBLE);
         } else if (!mFreezeViewVisibility) {
-            setVisibility(GONE);
+            setVisibility(INVISIBLE);
         }
     }
 
@@ -1764,7 +1764,7 @@
         if (mFreezeViewVisibility != freezeViewVisibility) {
             mFreezeViewVisibility = freezeViewVisibility;
             if (!mFreezeViewVisibility) {
-                setVisibility(mContentAlpha > 0 ? VISIBLE : GONE);
+                setVisibility(mContentAlpha > 0 ? VISIBLE : INVISIBLE);
             }
         }
     }
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index e891c95..dfbe6ce 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -48,7 +48,6 @@
 import android.content.Context;
 import android.content.Intent;
 import android.graphics.Outline;
-import android.graphics.Point;
 import android.graphics.Rect;
 import android.graphics.RectF;
 import android.graphics.drawable.Drawable;
@@ -429,7 +428,9 @@
             mContextualChip.setScaleX(comp(modalness));
             mContextualChip.setScaleY(comp(modalness));
         }
-        mDigitalWellBeingToast.updateBannerOffset(modalness);
+        mDigitalWellBeingToast.updateBannerOffset(modalness,
+                mCurrentFullscreenParams.mCurrentDrawnInsets.top
+                        + mCurrentFullscreenParams.mCurrentDrawnInsets.bottom);
     }
 
     public TaskMenuView getMenuView() {
@@ -666,7 +667,9 @@
             mContextualChip.setScaleX(scale);
             mContextualChip.setScaleY(scale);
         }
-        mDigitalWellBeingToast.updateBannerOffset(1f - scale);
+        mDigitalWellBeingToast.updateBannerOffset(1f - scale,
+                mCurrentFullscreenParams.mCurrentDrawnInsets.top
+                        + mCurrentFullscreenParams.mCurrentDrawnInsets.bottom);
     }
 
     public void setIconScaleAnimStartProgress(float startProgress) {
diff --git a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
index 35383d2..67840d1 100644
--- a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
+++ b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
@@ -200,8 +200,8 @@
                 () -> launcher.getNavigationModel() == expectedMode, WAIT_TIME_MS, launcher);
 
         Wait.atMost(() -> "Switching nav mode: "
-                        + launcher.getNavigationModeMismatchError(),
-                () -> launcher.getNavigationModeMismatchError() == null,
+                        + launcher.getNavigationModeMismatchError(false),
+                () -> launcher.getNavigationModeMismatchError(false) == null,
                 WAIT_TIME_MS, launcher);
         AbstractLauncherUiTest.checkDetectedLeaks(launcher);
         return true;
diff --git a/res/drawable/bg_widgets_searchbox.xml b/res/drawable/bg_widgets_searchbox.xml
new file mode 100644
index 0000000..81dd2aa
--- /dev/null
+++ b/res/drawable/bg_widgets_searchbox.xml
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2021 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.
+-->
+<shape xmlns:android="http://schemas.android.com/apk/res/android" android:shape="rectangle">
+    <solid android:color="#FFFFF7" />
+    <corners android:radius="24dp" />
+</shape>
\ No newline at end of file
diff --git a/res/layout/live_preview_widget_cell.xml b/res/layout/live_preview_widget_cell.xml
index 7a42d19..1e1ce6e 100644
--- a/res/layout/live_preview_widget_cell.xml
+++ b/res/layout/live_preview_widget_cell.xml
@@ -15,14 +15,16 @@
 -->
 <com.android.launcher3.dragndrop.LivePreviewWidgetCell
     xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="wrap_content"
+    android:layout_width="0dp"
     android:layout_height="wrap_content"
+    android:paddingHorizontal="@dimen/widget_cell_horizontal_padding"
+    android:paddingVertical="@dimen/widget_cell_vertical_padding"
     android:layout_weight="1"
     android:orientation="vertical"
     android:focusable="true"
     android:background="?android:attr/colorPrimaryDark"
     android:gravity="center_horizontal">
 
-    <include layout="@layout/widget_cell_content"  />
+    <include layout="@layout/widget_cell_content" />
 
 </com.android.launcher3.dragndrop.LivePreviewWidgetCell>
\ No newline at end of file
diff --git a/res/layout/personal_work_tabs.xml b/res/layout/personal_work_tabs.xml
index 8f29997..5fb5bcb 100644
--- a/res/layout/personal_work_tabs.xml
+++ b/res/layout/personal_work_tabs.xml
@@ -23,6 +23,7 @@
     android:layout_marginLeft="@dimen/all_apps_tabs_side_padding"
     android:layout_marginRight="@dimen/all_apps_tabs_side_padding"
     android:orientation="horizontal"
+    android:elevation="2dp"
     style="@style/TextHeadline">
 
     <Button
diff --git a/res/layout/widget_cell.xml b/res/layout/widget_cell.xml
index 148a99b..73a5737 100644
--- a/res/layout/widget_cell.xml
+++ b/res/layout/widget_cell.xml
@@ -15,8 +15,10 @@
 -->
 <com.android.launcher3.widget.WidgetCell
     xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="wrap_content"
+    android:layout_width="0dp"
     android:layout_height="wrap_content"
+    android:paddingHorizontal="@dimen/widget_cell_horizontal_padding"
+    android:paddingVertical="@dimen/widget_cell_vertical_padding"
     android:layout_weight="1"
     android:orientation="vertical"
     android:focusable="true"
diff --git a/res/layout/widget_cell_content.xml b/res/layout/widget_cell_content.xml
index 64f2362..c3e37e9 100644
--- a/res/layout/widget_cell_content.xml
+++ b/res/layout/widget_cell_content.xml
@@ -17,47 +17,37 @@
     android:layout_width="wrap_content"
     android:layout_height="wrap_content">
 
-    <LinearLayout
-        android:layout_width="match_parent"
-        android:layout_height="wrap_content"
-        android:paddingTop="@dimen/widget_preview_label_vertical_padding"
-        android:paddingBottom="@dimen/widget_preview_label_vertical_padding"
-        android:paddingLeft="@dimen/widget_preview_label_horizontal_padding"
-        android:paddingRight="@dimen/widget_preview_label_horizontal_padding"
-        android:orientation="horizontal">
-
-        <!-- The name of the widget. -->
-        <TextView
-            android:id="@+id/widget_name"
-            android:layout_width="0dp"
-            android:layout_height="wrap_content"
-            android:layout_weight="1"
-            android:ellipsize="end"
-            android:fadingEdge="horizontal"
-            android:gravity="start"
-            android:singleLine="true"
-            android:maxLines="1"
-            android:textColor="?android:attr/textColorSecondary"
-            android:textSize="14sp" />
-
-        <!-- The original dimensions of the widget (can't be the same text as above due to different
-             style. -->
-        <TextView
-            android:id="@+id/widget_dims"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
-            android:layout_marginStart="5dp"
-            android:layout_marginLeft="5dp"
-            android:textColor="?android:attr/textColorSecondary"
-            android:textSize="14sp"
-            android:alpha="0.8" />
-    </LinearLayout>
-
     <!-- The image of the widget. This view does not support padding. Any placement adjustment
          should be done using margins. -->
     <com.android.launcher3.widget.WidgetImageView
         android:id="@+id/widget_preview"
-        android:layout_width="match_parent"
+        android:layout_width="wrap_content"
         android:layout_height="0dp"
-        android:layout_weight="1" />
+        android:layout_weight="1"
+        android:layout_marginVertical="8dp" />
+
+    <!-- The name of the widget. -->
+    <TextView
+        android:id="@+id/widget_name"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:ellipsize="end"
+        android:fadingEdge="horizontal"
+        android:gravity="center_horizontal"
+        android:singleLine="true"
+        android:maxLines="1"
+        android:textColor="?android:attr/textColorSecondary"
+        android:textSize="14sp" />
+
+    <!-- The original dimensions of the widget (can't be the same text as above due to different
+         style. -->
+    <TextView
+        android:id="@+id/widget_dims"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:gravity="center_horizontal"
+        android:textColor="?android:attr/textColorSecondary"
+        android:textSize="14sp"
+        android:alpha="0.8" />
+
 </merge>
\ No newline at end of file
diff --git a/res/layout/widgets_bottom_sheet.xml b/res/layout/widgets_bottom_sheet.xml
index 3fdfc96..c1b2cbf 100644
--- a/res/layout/widgets_bottom_sheet.xml
+++ b/res/layout/widgets_bottom_sheet.xml
@@ -44,11 +44,16 @@
         android:textSize="14sp"
         android:text="@string/long_press_widget_to_add"/>
 
-    <include layout="@layout/widgets_scroll_container"
-         android:id="@+id/widgets"
-         android:layout_width="wrap_content"
-         android:layout_height="wrap_content"
-         android:layout_marginTop="45dp"
-         android:layout_marginBottom="40dp"/>
+    <ScrollView
+        android:id="@+id/widgets_table_scroll_view"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:layout_marginTop="45dp"
+        android:layout_marginBottom="40dp">
+        <include layout="@layout/widgets_table_container"
+            android:layout_width="match_parent"
+            android:layout_height="wrap_content"
+            android:layout_gravity="center_horizontal" />
+    </ScrollView>
 
 </com.android.launcher3.widget.WidgetsBottomSheet>
\ No newline at end of file
diff --git a/res/layout/widgets_full_sheet_paged_view.xml b/res/layout/widgets_full_sheet_paged_view.xml
index cfbb6dd..8125db8 100644
--- a/res/layout/widgets_full_sheet_paged_view.xml
+++ b/res/layout/widgets_full_sheet_paged_view.xml
@@ -16,13 +16,15 @@
 <merge xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:launcher="http://schemas.android.com/apk/res-auto">
 
-    <include layout="@layout/personal_work_tabs" />
+    <include layout="@layout/personal_work_tabs"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:layout_marginHorizontal="16dp" />
 
     <com.android.launcher3.workprofile.PersonalWorkPagedView
         android:id="@+id/widgets_view_pager"
         android:layout_width="match_parent"
         android:layout_height="match_parent"
-        android:layout_below="@+id/tabs"
         android:clipToPadding="false"
         android:descendantFocusability="afterDescendants"
         launcher:pageIndicator="@+id/tabs">
diff --git a/res/layout/widgets_full_sheet_search_and_recommendations.xml b/res/layout/widgets_full_sheet_search_and_recommendations.xml
new file mode 100644
index 0000000..9a6f922
--- /dev/null
+++ b/res/layout/widgets_full_sheet_search_and_recommendations.xml
@@ -0,0 +1,49 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2021 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.
+-->
+<LinearLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/search_and_recommendations_container"
+    android:layout_width="match_parent"
+    android:layout_height="wrap_content"
+    android:padding="16dp"
+    android:orientation="vertical">
+    <View
+        android:id="@+id/collapse_handle"
+        android:layout_width="48dp"
+        android:layout_height="2dp"
+        android:layout_gravity="center_horizontal"
+        android:background="@color/popup_color_primary_dark"/>
+    <TextView
+        android:id="@+id/title"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:gravity="center_horizontal"
+        android:textSize="24sp"
+        android:layout_marginTop="16dp"
+        android:text="@string/widget_button_text"/>
+    <!-- Disable the search bar because it has not been implemented. -->
+    <EditText
+        android:id="@+id/widgets_search_bar"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:visibility="gone"
+        android:layout_marginTop="16dp"
+        android:background="@drawable/bg_widgets_searchbox"
+        android:drawablePadding="8dp"
+        android:drawableStart="@drawable/ic_allapps_search"
+        android:hint="@string/widgets_full_sheet_search_bar_hint"
+        android:padding="12dp" />
+</LinearLayout>
diff --git a/res/layout/widgets_list_row_view.xml b/res/layout/widgets_list_row_view.xml
index eec57a5..5942ba6 100644
--- a/res/layout/widgets_list_row_view.xml
+++ b/res/layout/widgets_list_row_view.xml
@@ -45,5 +45,5 @@
         launcher:iconSizeOverride="@dimen/widget_section_icon_size"
         launcher:layoutHorizontal="true" />
 
-    <include layout="@layout/widgets_scroll_container" />
+    <include layout="@layout/widgets_table_container" />
 </LinearLayout>
diff --git a/res/layout/widgets_table_container.xml b/res/layout/widgets_table_container.xml
new file mode 100644
index 0000000..c4dfe7e
--- /dev/null
+++ b/res/layout/widgets_table_container.xml
@@ -0,0 +1,22 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2021 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.
+-->
+<TableLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/widgets_table"
+    android:layout_width="match_parent"
+    android:layout_height="wrap_content"
+    android:layout_marginHorizontal="8dp"
+    android:background="?android:attr/colorPrimaryDark" />
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index acc6466..5f78192 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -102,8 +102,8 @@
     <dimen name="work_profile_footer_text_size">16sp</dimen>
 
 <!-- Widget tray -->
-    <dimen name="widget_preview_label_vertical_padding">8dp</dimen>
-    <dimen name="widget_preview_label_horizontal_padding">16dp</dimen>
+    <dimen name="widget_cell_vertical_padding">8dp</dimen>
+    <dimen name="widget_cell_horizontal_padding">16dp</dimen>
 
     <dimen name="widget_preview_shadow_blur">0.5dp</dimen>
     <dimen name="widget_preview_key_shadow_distance">1dp</dimen>
@@ -183,13 +183,11 @@
     <dimen name="popup_padding_start">10dp</dimen>
     <dimen name="popup_padding_end">16dp</dimen>
     <dimen name="popup_vertical_padding">4dp</dimen>
-    <dimen name="popup_arrow_width">10dp</dimen>
-    <dimen name="popup_arrow_height">8dp</dimen>
-    <dimen name="popup_arrow_vertical_offset">-2dp</dimen>
+    <dimen name="popup_arrow_width">12dp</dimen>
+    <dimen name="popup_arrow_height">10dp</dimen>
+    <dimen name="popup_arrow_vertical_offset">-1dp</dimen>
     <!-- popup_padding_start + deep_shortcut_icon_size / 2 -->
-    <dimen name="popup_arrow_horizontal_center_start">28dp</dimen>
-    <!-- popup_padding_end + deep_shortcut_drag_handle_size / 2 -->
-    <dimen name="popup_arrow_horizontal_center_end">24dp</dimen>
+    <dimen name="popup_arrow_horizontal_center_offset">28dp</dimen>
     <dimen name="popup_arrow_corner_radius">2dp</dimen>
     <!-- popup_padding_start + icon_size + 10dp -->
     <dimen name="deep_shortcuts_text_padding_start">56dp</dimen>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 5a9def7..73f9e53 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -57,6 +57,12 @@
         <item quantity="one"><xliff:g id="widget_count" example="1">%1$d</xliff:g> widget</item>
         <item quantity="other"><xliff:g id="widget_count" example="2">%1$d</xliff:g> widgets</item>
     </plurals>
+    <!-- Text for both the tile of a popup view, which shows all available widgets installed on
+         the device, and the text of a button, which opens this popup view. [CHAR LIMIT=30]-->
+    <string name="widget_button_text">Widgets</string>
+    <!-- Search bar text shown in the popup view showing all available widgets installed on the
+         device. [CHAR_LIMIT=50] -->
+    <string name="widgets_full_sheet_search_bar_hint">Search</string>
 
     <!-- All Apps -->
     <!-- Search bar text in the apps view. [CHAR_LIMIT=50] -->
@@ -185,8 +191,6 @@
     <string name="folder_name_format_overflow">Folder: <xliff:g id="name" example="Games">%1$s</xliff:g>, <xliff:g id="size" example="2">%2$d</xliff:g> or more items</string>
 
     <!-- Strings for the customization mode -->
-    <!-- Text for widget add button [CHAR LIMIT=30]-->
-    <string name="widget_button_text">Widgets</string>
     <!-- Text for wallpaper change button [CHAR LIMIT=30]-->
     <string name="wallpaper_button_text">Wallpapers</string>
     <!-- Text for wallpaper change button [CHAR LIMIT=30]-->
diff --git a/robolectric_tests/Android.mk b/robolectric_tests/Android.mk
index 836ded5..405a458 100644
--- a/robolectric_tests/Android.mk
+++ b/robolectric_tests/Android.mk
@@ -19,6 +19,9 @@
 include $(CLEAR_VARS)
 
 LOCAL_MODULE := LauncherRoboTests
+LOCAL_LICENSE_KINDS := SPDX-license-identifier-Apache-2.0
+LOCAL_LICENSE_CONDITIONS := notice
+LOCAL_NOTICE_FILE := $(LOCAL_PATH)/../NOTICE
 LOCAL_MODULE_CLASS := JAVA_LIBRARIES
 
 LOCAL_SDK_VERSION := system_current
@@ -50,6 +53,9 @@
 include $(CLEAR_VARS)
 
 LOCAL_MODULE := RunLauncherRoboTests
+LOCAL_LICENSE_KINDS := SPDX-license-identifier-Apache-2.0
+LOCAL_LICENSE_CONDITIONS := notice
+LOCAL_NOTICE_FILE := $(LOCAL_PATH)/../NOTICE
 LOCAL_SDK_VERSION := system_current
 LOCAL_JAVA_LIBRARIES := LauncherRoboTests
 
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
similarity index 89%
rename from robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java
rename to robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
index ec9fde3..358e6e0 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
@@ -32,6 +32,7 @@
 import android.view.View.OnClickListener;
 import android.view.View.OnLongClickListener;
 import android.widget.FrameLayout;
+import android.widget.TableRow;
 import android.widget.TextView;
 
 import com.android.launcher3.DeviceProfile;
@@ -65,12 +66,12 @@
 import java.util.List;
 
 @RunWith(RobolectricTestRunner.class)
-public final class WidgetsListRowViewHolderBinderTest {
+public final class WidgetsListTableViewHolderBinderTest {
     private static final String TEST_PACKAGE = "com.google.test";
     private static final String APP_NAME = "Test app";
 
     private Context mContext;
-    private WidgetsListRowViewHolderBinder mViewHolderBinder;
+    private WidgetsListTableViewHolderBinder mViewHolderBinder;
     private InvariantDeviceProfile mTestProfile;
     // Replace ActivityController with ActivityScenario, which is the recommended way for activity
     // testing.
@@ -105,7 +106,7 @@
             return componentWithLabel.getComponent().getShortClassName();
         }).when(mIconCache).getTitleNoCache(any());
 
-        mViewHolderBinder = new WidgetsListRowViewHolderBinder(
+        mViewHolderBinder = new WidgetsListTableViewHolderBinder(
                 mContext,
                 LayoutInflater.from(mTestActivity),
                 mOnIconClickListener,
@@ -129,16 +130,17 @@
         mViewHolderBinder.bindViewHolder(viewHolder, entry);
         shadowOf(getMainLooper()).idle();
 
-        // THEN the cell container has 5 children: 3 widgets + 2 separators
-        // Index:        0       1        2       3        4
+        // THEN the table container has one row, which contains 3 widgets.
         // View:  .SampleWidget0 | .SampleWidget1 | .SampleWidget2
-        assertThat(viewHolder.cellContainer.getChildCount()).isEqualTo(5);
+        assertThat(viewHolder.mTableContainer.getChildCount()).isEqualTo(1);
+        TableRow row = (TableRow) viewHolder.mTableContainer.getChildAt(0);
+        assertThat(row.getChildCount()).isEqualTo(3);
         // Widget 0 label is .SampleWidget0.
-        assertWidgetCellWithLabel(viewHolder.cellContainer.getChildAt(0), ".SampleWidget0");
+        assertWidgetCellWithLabel(row.getChildAt(0), ".SampleWidget0");
         // Widget 1 label is .SampleWidget1.
-        assertWidgetCellWithLabel(viewHolder.cellContainer.getChildAt(2), ".SampleWidget1");
+        assertWidgetCellWithLabel(row.getChildAt(1), ".SampleWidget1");
         // Widget 2 label is .SampleWidget2.
-        assertWidgetCellWithLabel(viewHolder.cellContainer.getChildAt(4), ".SampleWidget2");
+        assertWidgetCellWithLabel(row.getChildAt(2), ".SampleWidget2");
     }
 
     private WidgetsListContentEntry generateSampleAppWithWidgets(String appName, String packageName,
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
new file mode 100644
index 0000000..5922223
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
@@ -0,0 +1,201 @@
+/*
+ * Copyright (C) 2021 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.widget.picker.util;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.robolectric.Shadows.shadowOf;
+
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.graphics.Point;
+import android.graphics.drawable.Drawable;
+import android.os.UserHandle;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.icons.ComponentWithLabel;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.pm.ShortcutConfigActivityInfo;
+import com.android.launcher3.widget.util.WidgetsTableUtils;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@RunWith(RobolectricTestRunner.class)
+public final class WidgetsTableUtilsTest {
+    private static final String TEST_PACKAGE = "com.google.test";
+
+    @Mock
+    private IconCache mIconCache;
+
+    private Context mContext;
+    private InvariantDeviceProfile mTestProfile;
+    private WidgetItem mWidget1x1;
+    private WidgetItem mWidget2x2;
+    private WidgetItem mWidget2x3;
+    private WidgetItem mWidget2x4;
+    private WidgetItem mWidget4x4;
+
+    private WidgetItem mShortcut1;
+    private WidgetItem mShortcut2;
+    private WidgetItem mShortcut3;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+
+        mContext = RuntimeEnvironment.application;
+
+        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile.numRows = 5;
+        mTestProfile.numColumns = 5;
+
+        initTestWidgets();
+        initTestShortcuts();
+
+        doAnswer(invocation -> ((ComponentWithLabel) invocation.getArgument(0))
+                .getComponent().getPackageName())
+                .when(mIconCache).getTitleNoCache(any());
+    }
+
+
+    @Test
+    public void groupWidgetItemsIntoTable_widgetsOnly_maxSpansPerRow5_shouldGroupWidgetsInTable() {
+        List<WidgetItem> widgetItems = List.of(mWidget4x4, mWidget2x3, mWidget1x1, mWidget2x4,
+                mWidget2x2);
+
+        List<ArrayList<WidgetItem>> widgetItemInTable = WidgetsTableUtils.groupWidgetItemsIntoTable(
+                widgetItems, /* maxSpansPerRow= */ 5);
+
+        // Row 0: 1x1, 2x2, 2x3
+        // Row 1: 2x4
+        // Row 2: 4x4
+        assertThat(widgetItemInTable).hasSize(3);
+        assertThat(widgetItemInTable.get(0)).containsExactly(mWidget1x1, mWidget2x2, mWidget2x3);
+        assertThat(widgetItemInTable.get(1)).containsExactly(mWidget2x4);
+        assertThat(widgetItemInTable.get(2)).containsExactly(mWidget4x4);
+    }
+
+    @Test
+    public void groupWidgetItemsIntoTable_widgetsOnly_maxSpansPerRow4_shouldGroupWidgetsInTable() {
+        List<WidgetItem> widgetItems = List.of(mWidget4x4, mWidget2x3, mWidget1x1, mWidget2x4,
+                mWidget2x2);
+
+        List<ArrayList<WidgetItem>> widgetItemInTable = WidgetsTableUtils.groupWidgetItemsIntoTable(
+                widgetItems, /* maxSpansPerRow= */ 4);
+
+        // Row 0: 1x1, 2x2
+        // Row 1: 2x3, 2x4
+        // Row 2: 4x4
+        assertThat(widgetItemInTable).hasSize(3);
+        assertThat(widgetItemInTable.get(0)).containsExactly(mWidget1x1, mWidget2x2);
+        assertThat(widgetItemInTable.get(1)).containsExactly(mWidget2x3, mWidget2x4);
+        assertThat(widgetItemInTable.get(2)).containsExactly(mWidget4x4);
+    }
+
+    @Test
+    public void groupWidgetItemsIntoTable_mixItems_maxSpansPerRow4_shouldGroupWidgetsInTable() {
+        List<WidgetItem> widgetItems = List.of(mWidget4x4, mShortcut3, mWidget2x3, mShortcut1,
+                mWidget1x1, mShortcut2, mWidget2x4, mWidget2x2);
+
+        List<ArrayList<WidgetItem>> widgetItemInTable = WidgetsTableUtils.groupWidgetItemsIntoTable(
+                widgetItems, /* maxSpansPerRow= */ 4);
+
+        // Row 0: 1x1, 2x2
+        // Row 1: 2x3, 2x4
+        // Row 2: 4x4
+        // Row 3: shortcut3, shortcut1, shortcut2
+        assertThat(widgetItemInTable).hasSize(4);
+        assertThat(widgetItemInTable.get(0)).containsExactly(mWidget1x1, mWidget2x2);
+        assertThat(widgetItemInTable.get(1)).containsExactly(mWidget2x3, mWidget2x4);
+        assertThat(widgetItemInTable.get(2)).containsExactly(mWidget4x4);
+        assertThat(widgetItemInTable.get(3)).containsExactly(mShortcut3, mShortcut2, mShortcut1);
+    }
+
+    private void initTestWidgets() {
+        List<Point> widgetSizes = List.of(new Point(1, 1), new Point(2, 2), new Point(2, 3),
+                new Point(2, 4), new Point(4, 4));
+
+        ArrayList<WidgetItem> widgetItems = new ArrayList<>();
+        widgetSizes.stream().forEach(
+                widgetSize -> {
+                    ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+                    AppWidgetProviderInfo info = new AppWidgetProviderInfo();
+                    info.provider = ComponentName.createRelative(TEST_PACKAGE,
+                            ".WidgetProvider_" + widgetSize.x + "x" + widgetSize.y);
+                    LauncherAppWidgetProviderInfo widgetInfo =
+                            LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, info);
+                    widgetInfo.spanX = widgetSize.x;
+                    widgetInfo.spanY = widgetSize.y;
+                    ReflectionHelpers.setField(widgetInfo, "providerInfo",
+                            packageManager.addReceiverIfNotPresent(widgetInfo.provider));
+                    widgetItems.add(new WidgetItem(widgetInfo, mTestProfile, mIconCache));
+                }
+        );
+        mWidget1x1 = widgetItems.get(0);
+        mWidget2x2 = widgetItems.get(1);
+        mWidget2x3 = widgetItems.get(2);
+        mWidget2x4 = widgetItems.get(3);
+        mWidget4x4 = widgetItems.get(4);
+    }
+
+    private void initTestShortcuts() {
+        PackageManager packageManager = mContext.getPackageManager();
+        mShortcut1 = new WidgetItem(new TestShortcutConfigActivityInfo(
+                ComponentName.createRelative(TEST_PACKAGE, ".shortcut1"), UserHandle.CURRENT),
+                mIconCache, packageManager);
+        mShortcut2 = new WidgetItem(new TestShortcutConfigActivityInfo(
+                ComponentName.createRelative(TEST_PACKAGE, ".shortcut2"), UserHandle.CURRENT),
+                mIconCache, packageManager);
+        mShortcut3 = new WidgetItem(new TestShortcutConfigActivityInfo(
+                ComponentName.createRelative(TEST_PACKAGE, ".shortcut3"), UserHandle.CURRENT),
+                mIconCache, packageManager);
+
+    }
+
+    private final class TestShortcutConfigActivityInfo extends ShortcutConfigActivityInfo {
+
+        TestShortcutConfigActivityInfo(ComponentName componentName, UserHandle user) {
+            super(componentName, user);
+        }
+
+        @Override
+        public Drawable getFullResIcon(IconCache cache) {
+            return null;
+        }
+
+        @Override
+        public CharSequence getLabel(PackageManager pm) {
+            return null;
+        }
+    }
+}
diff --git a/src/com/android/launcher3/AppWidgetResizeFrame.java b/src/com/android/launcher3/AppWidgetResizeFrame.java
index 1330ed4..be92c5d 100644
--- a/src/com/android/launcher3/AppWidgetResizeFrame.java
+++ b/src/com/android/launcher3/AppWidgetResizeFrame.java
@@ -2,6 +2,8 @@
 
 import static com.android.launcher3.LauncherAnimUtils.LAYOUT_HEIGHT;
 import static com.android.launcher3.LauncherAnimUtils.LAYOUT_WIDTH;
+import static com.android.launcher3.Utilities.ATLEAST_S;
+import static com.android.launcher3.config.FeatureFlags.ENABLE_FOUR_COLUMNS;
 import static com.android.launcher3.views.BaseDragLayer.LAYOUT_X;
 import static com.android.launcher3.views.BaseDragLayer.LAYOUT_Y;
 
@@ -11,14 +13,19 @@
 import android.appwidget.AppWidgetHostView;
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
+import android.content.res.Configuration;
 import android.graphics.Point;
+import android.graphics.PointF;
 import android.graphics.Rect;
+import android.os.Bundle;
 import android.util.AttributeSet;
 import android.view.KeyEvent;
 import android.view.MotionEvent;
 import android.view.View;
 import android.view.ViewGroup;
 
+import androidx.annotation.Nullable;
+
 import com.android.launcher3.accessibility.DragViewStateAnnouncer;
 import com.android.launcher3.dragndrop.DragLayer;
 import com.android.launcher3.util.FocusLogic;
@@ -352,33 +359,99 @@
     }
 
     public static void updateWidgetSizeRanges(AppWidgetHostView widgetView, Launcher launcher,
-            int spanX, int spanY) {
-        getWidgetSizeRanges(launcher, spanX, spanY, sTmpRect);
-        widgetView.updateAppWidgetSize(null, sTmpRect.left, sTmpRect.top,
-                sTmpRect.right, sTmpRect.bottom);
+                                              int spanX, int spanY) {
+        List<PointF> sizes = getWidgetSizes(launcher, spanX, spanY);
+        if (ATLEAST_S) {
+            widgetView.updateAppWidgetSize(new Bundle(), sizes);
+        } else {
+            Rect bounds = getMinMaxSizes(sizes, null /* outRect */);
+            widgetView.updateAppWidgetSize(new Bundle(), bounds.left, bounds.top, bounds.right,
+                    bounds.bottom);
+        }
     }
 
-    public static Rect getWidgetSizeRanges(Context context, int spanX, int spanY, Rect rect) {
-        if (rect == null) {
-            rect = new Rect();
-        }
+    private static PointF getWidgetSize(Context context, Point cellSize, int spanX, int spanY) {
         final float density = context.getResources().getDisplayMetrics().density;
+        float hBorderSpacing = 0;
+        float vBorderSpacing = 0;
+        if (ENABLE_FOUR_COLUMNS.get()) {
+            final int borderSpacing = context.getResources()
+                    .getDimensionPixelSize(R.dimen.dynamic_grid_cell_border_spacing);
+            hBorderSpacing = (spanX - 1) * borderSpacing;
+            vBorderSpacing = (spanY - 1) * borderSpacing;
+        }
+        PointF widgetSize = new PointF();
+        widgetSize.x = ((spanX * cellSize.x) + hBorderSpacing) / density;
+        widgetSize.y = ((spanY * cellSize.y) + vBorderSpacing) / density;
+        return widgetSize;
+    }
+
+    /** Returns the actual widget size given its span. */
+    public static PointF getWidgetSize(Context context, int spanX, int spanY) {
+        final Point[] cellSize = CELL_SIZE.get(context);
+        if (isLandscape(context)) {
+            return getWidgetSize(context, cellSize[0], spanX, spanY);
+        }
+        return getWidgetSize(context, cellSize[1], spanX, spanY);
+    }
+
+    /** Returns true if the screen is in landscape mode. */
+    private static boolean isLandscape(Context context) {
+        return context.getResources().getConfiguration().orientation
+                == Configuration.ORIENTATION_LANDSCAPE;
+    }
+
+    /** Returns the list of sizes for a widget of given span, in dp. */
+    public static ArrayList<PointF> getWidgetSizes(Context context, int spanX, int spanY) {
         final Point[] cellSize = CELL_SIZE.get(context);
 
-        final int borderSpacing = context.getResources()
-                .getDimensionPixelSize(R.dimen.dynamic_grid_cell_border_spacing);
-        final float hBorderSpacing = (spanX - 1) * borderSpacing;
-        final float vBorderSpacing = (spanY - 1) * borderSpacing;
+        PointF landSize = getWidgetSize(context, cellSize[0], spanX, spanY);
+        PointF portSize = getWidgetSize(context, cellSize[1], spanX, spanY);
 
-        // Compute landscape size
-        int landWidth = (int) (((spanX * cellSize[0].x) + hBorderSpacing) / density);
-        int landHeight = (int) (((spanY * cellSize[0].y) + vBorderSpacing) / density);
+        ArrayList<PointF> sizes = new ArrayList<>(2);
+        sizes.add(landSize);
+        sizes.add(portSize);
+        return sizes;
+    }
 
-        // Compute portrait size
-        int portWidth = (int) (((spanX * cellSize[1].x) + hBorderSpacing) / density);
-        int portHeight = (int) (((spanY * cellSize[1].y) + vBorderSpacing) / density);
-        rect.set(portWidth, landHeight, landWidth, portHeight);
-        return rect;
+    /**
+     * Returns the min and max widths and heights given a list of sizes, in dp.
+     *
+     * @param sizes List of sizes to get the min/max from.
+     * @param outRect Rectangle in which the result can be stored, to avoid extra allocations. If
+     *               null, a new rectangle will be allocated.
+     * @return A rectangle with the left (resp. top) is used for the min width (resp. height) and
+     * the right (resp. bottom) for the max. The returned rectangle is set with 0s if the list is
+     * empty.
+     */
+    public static Rect getMinMaxSizes(List<PointF> sizes, @Nullable Rect outRect) {
+        if (outRect == null) {
+            outRect = new Rect();
+        }
+        if (sizes.isEmpty()) {
+            outRect.set(0, 0, 0, 0);
+        } else {
+            PointF first = sizes.get(0);
+            outRect.set((int) first.x, (int) first.y, (int) first.x, (int) first.y);
+            for (int i = 1; i < sizes.size(); i++) {
+                outRect.union((int) sizes.get(i).x, (int) sizes.get(i).y);
+            }
+        }
+        return outRect;
+    }
+
+    /**
+     * Returns the range of sizes a widget may be displayed, given its span.
+     *
+     * @param context Context in which the View is rendered.
+     * @param spanX Width of the widget, in cells.
+     * @param spanY Height of the widget, in cells.
+     * @param outRect Rectangle in which the result can be stored, to avoid extra allocations. If
+     *               null, a new rectangle will be allocated.
+     */
+    public static Rect getWidgetSizeRanges(Context context, int spanX, int spanY,
+            @Nullable Rect outRect) {
+        return getMinMaxSizes(getWidgetSizes(context, spanX, spanY), outRect);
     }
 
     @Override
diff --git a/src/com/android/launcher3/BaseRecyclerView.java b/src/com/android/launcher3/BaseRecyclerView.java
index c55b46b..9369bdc 100644
--- a/src/com/android/launcher3/BaseRecyclerView.java
+++ b/src/com/android/launcher3/BaseRecyclerView.java
@@ -23,6 +23,7 @@
 import android.view.ViewGroup;
 import android.view.accessibility.AccessibilityNodeInfo;
 
+import androidx.recyclerview.widget.LinearLayoutManager;
 import androidx.recyclerview.widget.RecyclerView;
 
 import com.android.launcher3.compat.AccessibilityManagerCompat;
@@ -188,4 +189,21 @@
         super.onInitializeAccessibilityNodeInfo(info);
         if (isLayoutSuppressed()) info.setScrollable(false);
     }
+
+    /**
+     * Scrolls this recycler view to the top.
+     */
+    public void scrollToTop() {
+        if (mScrollbar != null) {
+            mScrollbar.reattachThumbToScroll();
+        }
+        if (getLayoutManager() instanceof LinearLayoutManager) {
+            LinearLayoutManager layoutManager = (LinearLayoutManager) getLayoutManager();
+            if (layoutManager.findFirstCompletelyVisibleItemPosition() == 0) {
+                // We are at the top, so don't scrollToPosition (would cause unnecessary relayout).
+                return;
+            }
+        }
+        scrollToPosition(0);
+    }
 }
\ No newline at end of file
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 49adf1f..f7ff262 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -42,6 +42,7 @@
 import static com.android.launcher3.logging.StatsLogManager.LAUNCHER_STATE_BACKGROUND;
 import static com.android.launcher3.logging.StatsLogManager.LAUNCHER_STATE_HOME;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_ENTRY;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_ENTRY_WITH_DEVICE_SEARCH;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_EXIT;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ONRESUME;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ONSTOP;
@@ -1067,7 +1068,11 @@
         if (ALL_APPS.equals(state)) {
             // creates new instance ID since new all apps session is started.
             mAllAppsSessionLogId = new InstanceIdSequence().newInstanceId();
-            getStatsLogManager().logger().log(LAUNCHER_ALLAPPS_ENTRY);
+            getStatsLogManager()
+                    .logger()
+                    .log(FeatureFlags.ENABLE_DEVICE_SEARCH.get()
+                            ? LAUNCHER_ALLAPPS_ENTRY_WITH_DEVICE_SEARCH
+                            : LAUNCHER_ALLAPPS_ENTRY);
         } else if (ALL_APPS.equals(mPrevLauncherState)
                 // Check if mLogInstanceId is not null to make sure exit event is logged only once.
                 && mAllAppsSessionLogId != null) {
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index e89b9b0..699495c 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -615,7 +615,9 @@
                         + "\" bitmapIcon=" + info.bitmap.icon
                         + " componentName=" + info.componentName.getPackageName());
             }
+            writer.println();
         }
+        mModelDelegate.dump(prefix, fd, writer, args);
         mBgDataModel.dump(prefix, fd, writer, args);
     }
 
diff --git a/src/com/android/launcher3/accessibility/LauncherAccessibilityDelegate.java b/src/com/android/launcher3/accessibility/LauncherAccessibilityDelegate.java
index cd4616a..db7fd3f 100644
--- a/src/com/android/launcher3/accessibility/LauncherAccessibilityDelegate.java
+++ b/src/com/android/launcher3/accessibility/LauncherAccessibilityDelegate.java
@@ -357,10 +357,8 @@
         }
 
         layout.markCellsAsOccupiedForView(host);
-        Rect sizeRange = new Rect();
-        AppWidgetResizeFrame.getWidgetSizeRanges(mLauncher, info.spanX, info.spanY, sizeRange);
-        ((LauncherAppWidgetHostView) host).updateAppWidgetSize(null,
-                sizeRange.left, sizeRange.top, sizeRange.right, sizeRange.bottom);
+        AppWidgetResizeFrame.updateWidgetSizeRanges(((LauncherAppWidgetHostView) host), mLauncher,
+                info.spanX, info.spanY);
         host.requestLayout();
         mLauncher.getModelWriter().updateItemInDatabase(info);
         announceConfirmation(mLauncher.getString(R.string.widget_resized, info.spanX, info.spanY));
diff --git a/src/com/android/launcher3/allapps/AllAppsContainerView.java b/src/com/android/launcher3/allapps/AllAppsContainerView.java
index 746bfba..a92e1aa 100644
--- a/src/com/android/launcher3/allapps/AllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsContainerView.java
@@ -16,6 +16,8 @@
 package com.android.launcher3.allapps;
 
 import static com.android.launcher3.allapps.AllAppsGridAdapter.AdapterItem;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_TAP_ON_PERSONAL_TAB;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_TAP_ON_WORK_TAB;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_HAS_SHORTCUT_PERMISSION;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_QUIET_MODE_CHANGE_PERMISSION;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_QUIET_MODE_ENABLED;
@@ -32,6 +34,7 @@
 import android.text.Selection;
 import android.text.SpannableStringBuilder;
 import android.util.AttributeSet;
+import android.util.Log;
 import android.view.KeyEvent;
 import android.view.LayoutInflater;
 import android.view.MotionEvent;
@@ -61,6 +64,7 @@
 import com.android.launcher3.keyboard.FocusedItemDecorator;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.testing.TestProtocol;
 import com.android.launcher3.util.ItemInfoMatcher;
 import com.android.launcher3.util.MultiValueAlpha;
 import com.android.launcher3.util.MultiValueAlpha.AlphaProperty;
@@ -432,9 +436,19 @@
             mAH[AdapterHolder.WORK].setup(mViewPager.getChildAt(1), mWorkMatcher);
             mViewPager.getPageIndicator().setActiveMarker(AdapterHolder.MAIN);
             findViewById(R.id.tab_personal)
-                    .setOnClickListener((View view) -> mViewPager.snapToPage(AdapterHolder.MAIN));
+                    .setOnClickListener((View view) -> {
+                        if (mViewPager.snapToPage(AdapterHolder.MAIN)) {
+                            mLauncher.getStatsLogManager().logger()
+                                    .log(LAUNCHER_ALLAPPS_TAP_ON_PERSONAL_TAB);
+                        }
+                    });
             findViewById(R.id.tab_work)
-                    .setOnClickListener((View view) -> mViewPager.snapToPage(AdapterHolder.WORK));
+                    .setOnClickListener((View view) -> {
+                        if (mViewPager.snapToPage(AdapterHolder.WORK)) {
+                            mLauncher.getStatsLogManager().logger()
+                                    .log(LAUNCHER_ALLAPPS_TAP_ON_WORK_TAB);
+                        }
+                    });
             onActivePageChanged(mViewPager.getNextPage());
         } else {
             mAH[AdapterHolder.MAIN].setup(findViewById(R.id.apps_list_view), null);
@@ -481,6 +495,10 @@
         int layout = showTabs ? R.layout.all_apps_tabs : R.layout.all_apps_rv_layout;
         View newView = getLayoutInflater().inflate(layout, this, false);
         addView(newView, index);
+        if (TestProtocol.sDebugTracing) {
+            Log.d(TestProtocol.WORK_PROFILE_REMOVED, "should show tabs:" + showTabs,
+                    new Exception());
+        }
         if (showTabs) {
             mViewPager = (AllAppsPagedView) newView;
             mViewPager.initParentViews(this);
@@ -757,6 +775,7 @@
                 int bottomOffset = mWorkModeSwitch != null && mIsWork ? switchH : 0;
                 recyclerView.setPadding(padding.left, padding.top, padding.right,
                         padding.bottom + bottomOffset);
+                recyclerView.scrollToTop();
             }
         }
 
diff --git a/src/com/android/launcher3/allapps/AllAppsPagedView.java b/src/com/android/launcher3/allapps/AllAppsPagedView.java
index 647402b..14e3b51 100644
--- a/src/com/android/launcher3/allapps/AllAppsPagedView.java
+++ b/src/com/android/launcher3/allapps/AllAppsPagedView.java
@@ -15,9 +15,13 @@
  */
 package com.android.launcher3.allapps;
 
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_SWIPE_TO_PERSONAL_TAB;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_SWIPE_TO_WORK_TAB;
+
 import android.content.Context;
 import android.util.AttributeSet;
 
+import com.android.launcher3.Launcher;
 import com.android.launcher3.PagedView;
 import com.android.launcher3.R;
 import com.android.launcher3.config.FeatureFlags;
@@ -44,4 +48,16 @@
                         R.dimen.all_apps_header_top_padding);
         setPadding(0, topPadding, 0, 0);
     }
+
+    @Override
+    protected boolean snapToPageWithVelocity(int whichPage, int velocity) {
+        boolean resp = super.snapToPageWithVelocity(whichPage, velocity);
+        if (resp && whichPage != mCurrentPage) {
+            Launcher.getLauncher(getContext()).getStatsLogManager().logger()
+                    .log(mCurrentPage < whichPage
+                            ? LAUNCHER_ALLAPPS_SWIPE_TO_WORK_TAB
+                            : LAUNCHER_ALLAPPS_SWIPE_TO_PERSONAL_TAB);
+        }
+        return resp;
+    }
 }
diff --git a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
index e61b95d..ace9938 100644
--- a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
@@ -36,7 +36,6 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.R;
-import com.android.launcher3.allapps.AllAppsGridAdapter.AppsGridLayoutManager;
 import com.android.launcher3.views.RecyclerViewFastScroller;
 
 import java.util.ArrayList;
@@ -109,23 +108,6 @@
         mViewHeights.put(AllAppsGridAdapter.VIEW_TYPE_ICON, grid.allAppsCellHeightPx);
     }
 
-    /**
-     * Scrolls this recycler view to the top.
-     */
-    public void scrollToTop() {
-        // Ensure we reattach the scrollbar if it was previously detached while fast-scrolling
-        if (mScrollbar != null) {
-            mScrollbar.reattachThumbToScroll();
-        }
-        if (getLayoutManager() instanceof AppsGridLayoutManager) {
-            AppsGridLayoutManager layoutManager = (AppsGridLayoutManager) getLayoutManager();
-            if (layoutManager.findFirstCompletelyVisibleItemPosition() == 0) {
-                // We are at the top, so don't scrollToPosition (would cause unnecessary relayout).
-                return;
-            }
-        }
-        scrollToPosition(0);
-    }
 
     @Override
     public void onDraw(Canvas c) {
diff --git a/src/com/android/launcher3/allapps/AllAppsTransitionController.java b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
index a4e1f27..dc58c99 100644
--- a/src/com/android/launcher3/allapps/AllAppsTransitionController.java
+++ b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
@@ -33,6 +33,7 @@
 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;
@@ -63,7 +64,7 @@
  * closer to top or closer to the page indicator.
  */
 public class AllAppsTransitionController implements StateHandler<LauncherState>,
-        OnDeviceProfileChangeListener {
+        OnDeviceProfileChangeListener, SharedPreferences.OnSharedPreferenceChangeListener {
 
     public static final FloatProperty<AllAppsTransitionController> ALL_APPS_PROGRESS =
             new FloatProperty<AllAppsTransitionController>("allAppsProgress") {
@@ -80,6 +81,7 @@
             };
 
     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;
@@ -98,6 +100,7 @@
 
     private float mScrollRangeDelta = 0;
     private AllAppsInsetTransitionController mInsetController;
+    private boolean mSearchImeEnabled;
 
     public AllAppsTransitionController(Launcher l) {
         mLauncher = l;
@@ -106,6 +109,9 @@
 
         mIsVerticalLayout = mLauncher.getDeviceProfile().isVerticalBarLayout();
         mLauncher.addOnDeviceProfileChangeListener(this);
+
+        onSharedPreferenceChanged(mLauncher.getSharedPrefs(), PREF_KEY_SHOW_SEARCH_IME);
+        mLauncher.getSharedPrefs().registerOnSharedPreferenceChangeListener(this);
     }
 
     public float getShiftRange() {
@@ -142,8 +148,7 @@
         float shiftCurrent = progress * mShiftRange;
 
         mAppsView.setTranslationY(shiftCurrent);
-        if (FeatureFlags.ENABLE_DEVICE_SEARCH.get()
-                && !FeatureFlags.DISABLE_INITIAL_IME_IN_ALLAPPS.get()) {
+        if (FeatureFlags.ENABLE_DEVICE_SEARCH.get() && mSearchImeEnabled) {
             mInsetController.setProgress(progress);
         }
     }
@@ -234,9 +239,7 @@
     public void setupViews(AllAppsContainerView appsView, ScrimView scrimView) {
         mAppsView = appsView;
         mScrimView = scrimView;
-        if (FeatureFlags.ENABLE_DEVICE_SEARCH.get()
-                && !FeatureFlags.DISABLE_INITIAL_IME_IN_ALLAPPS.get()
-                && BuildCompat.isAtLeastR()) {
+        if (FeatureFlags.ENABLE_DEVICE_SEARCH.get() && BuildCompat.isAtLeastR()) {
             mInsetController = new AllAppsInsetTransitionController(mShiftRange, mAppsView);
             mLauncher.getSystemUiController().updateUiState(UI_STATE_ALLAPPS,
                     View.SYSTEM_UI_FLAG_LAYOUT_FULLSCREEN
@@ -264,8 +267,8 @@
         if (Float.compare(mProgress, 1f) == 0) {
             mAppsView.reset(false /* animate */);
         }
-        if (FeatureFlags.ENABLE_DEVICE_SEARCH.get()
-                && !FeatureFlags.DISABLE_INITIAL_IME_IN_ALLAPPS.get() && BuildCompat.isAtLeastR()) {
+        if (FeatureFlags.ENABLE_DEVICE_SEARCH.get() && mSearchImeEnabled
+                && BuildCompat.isAtLeastR()) {
             mInsetController.onAnimationEnd(mProgress);
             if (Float.compare(mProgress, 0f) == 0) {
                 EditText editText = mAppsView.getSearchUiManager().getEditText();
@@ -276,4 +279,11 @@
             // 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/FloatingHeaderRow.java b/src/com/android/launcher3/allapps/FloatingHeaderRow.java
index e357f61..31c6cc7 100644
--- a/src/com/android/launcher3/allapps/FloatingHeaderRow.java
+++ b/src/com/android/launcher3/allapps/FloatingHeaderRow.java
@@ -61,4 +61,11 @@
      * Returns a child that has focus to be launched by the IME.
      */
     View getFocusedChild();
+
+    /**
+     * Returns true if view is currently visible
+     */
+    default boolean isVisible() {
+        return shouldDraw();
+    }
 }
diff --git a/src/com/android/launcher3/allapps/FloatingHeaderView.java b/src/com/android/launcher3/allapps/FloatingHeaderView.java
index 813db7d..9056e8a 100644
--- a/src/com/android/launcher3/allapps/FloatingHeaderView.java
+++ b/src/com/android/launcher3/allapps/FloatingHeaderView.java
@@ -200,7 +200,7 @@
     public View getFocusedChild() {
         if (FeatureFlags.ENABLE_DEVICE_SEARCH.get()) {
             for (FloatingHeaderRow row : mAllRows) {
-                if (row.hasVisibleContent() && row.shouldDraw()) {
+                if (row.hasVisibleContent() && row.isVisible()) {
                     return row.getFocusedChild();
                 }
             }
diff --git a/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java b/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
index bc2e66c..13ddc12 100644
--- a/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
@@ -16,8 +16,6 @@
 package com.android.launcher3.allapps;
 
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_KEYBOARD_CLOSED;
-import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_SWITCHED_TO_MAIN_TAB;
-import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_SWITCHED_TO_WORK_TAB;
 
 import android.content.Context;
 import android.graphics.Rect;
@@ -90,14 +88,6 @@
     public void onActivePageChanged(int currentActivePage) {
         super.onActivePageChanged(currentActivePage);
         if (mUsingTabs) {
-            // Log tab switches only when the launcher is in AllApps state
-            if (mLauncher.getStateManager().getCurrentStableState() == LauncherState.ALL_APPS) {
-                mLauncher.getStatsLogManager().logger()
-                        .log(currentActivePage == AdapterHolder.WORK
-                                ? LAUNCHER_ALLAPPS_SWITCHED_TO_WORK_TAB
-                                : LAUNCHER_ALLAPPS_SWITCHED_TO_MAIN_TAB);
-            }
-
             if (currentActivePage == AdapterHolder.WORK) {
                 WorkEduView.showWorkEduIfNeeded(mLauncher);
             } else {
diff --git a/src/com/android/launcher3/allapps/SearchUiManager.java b/src/com/android/launcher3/allapps/SearchUiManager.java
index f926086..39410a7 100644
--- a/src/com/android/launcher3/allapps/SearchUiManager.java
+++ b/src/com/android/launcher3/allapps/SearchUiManager.java
@@ -61,7 +61,8 @@
     /**
      * Called when activity is destroyed. Used to close search system services
      */
-    default void destroy(){}
+    default void destroy() {
+    }
 
     /**
      * Returns true if the QSB should be visible for the given set of visible elements
@@ -75,4 +76,9 @@
      */
     @Nullable
     EditText getEditText();
+
+    /**
+     * sets highlight result's title
+     */
+    default void setFocusedResultTitle(@Nullable  CharSequence title) { }
 }
diff --git a/src/com/android/launcher3/allapps/search/AllAppsSearchBarController.java b/src/com/android/launcher3/allapps/search/AllAppsSearchBarController.java
index 198c4b2..3319018 100644
--- a/src/com/android/launcher3/allapps/search/AllAppsSearchBarController.java
+++ b/src/com/android/launcher3/allapps/search/AllAppsSearchBarController.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.allapps.search;
 
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_FOCUSED_ITEM_SELECTED_WITH_IME;
+
 import android.text.Editable;
 import android.text.TextUtils;
 import android.text.TextWatcher;
@@ -31,10 +33,10 @@
 import com.android.launcher3.Utilities;
 import com.android.launcher3.allapps.AllAppsGridAdapter.AdapterItem;
 import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.search.SearchAlgorithm;
+import com.android.launcher3.search.SearchCallback;
 import com.android.launcher3.util.PackageManagerHelper;
 
-import java.util.ArrayList;
-
 /**
  * An interface to a search box that AllApps can command.
  */
@@ -43,11 +45,11 @@
         OnFocusChangeListener {
 
     protected BaseDraggingActivity mLauncher;
-    protected Callbacks mCb;
+    protected SearchCallback<AdapterItem> mCallback;
     protected ExtendedEditText mInput;
     protected String mQuery;
 
-    protected SearchAlgorithm mSearchAlgorithm;
+    protected SearchAlgorithm<AdapterItem> mSearchAlgorithm;
 
     public void setVisibility(int visibility) {
         mInput.setVisibility(visibility);
@@ -57,9 +59,9 @@
      * Sets the references to the apps model and the search result callback.
      */
     public final void initialize(
-            SearchAlgorithm searchAlgorithm, ExtendedEditText input,
-            BaseDraggingActivity launcher, Callbacks cb) {
-        mCb = cb;
+            SearchAlgorithm<AdapterItem> searchAlgorithm, ExtendedEditText input,
+            BaseDraggingActivity launcher, SearchCallback<AdapterItem> callback) {
+        mCallback = callback;
         mLauncher = launcher;
 
         mInput = input;
@@ -85,10 +87,10 @@
         mQuery = s.toString();
         if (mQuery.isEmpty()) {
             mSearchAlgorithm.cancel(true);
-            mCb.clearSearchResult();
+            mCallback.clearSearchResult();
         } else {
             mSearchAlgorithm.cancel(false);
-            mSearchAlgorithm.doSearch(mQuery, mCb);
+            mSearchAlgorithm.doSearch(mQuery, mCallback);
         }
     }
 
@@ -98,13 +100,15 @@
         }
         // If play store continues auto updating an app, we want to show partial result.
         mSearchAlgorithm.cancel(false);
-        mSearchAlgorithm.doSearch(mQuery, mCb);
+        mSearchAlgorithm.doSearch(mQuery, mCallback);
     }
 
     @Override
     public boolean onEditorAction(TextView v, int actionId, KeyEvent event) {
         if (FeatureFlags.ENABLE_DEVICE_SEARCH.get()) {
             if (actionId == EditorInfo.IME_ACTION_SEARCH || actionId == EditorInfo.IME_ACTION_GO) {
+                mLauncher.getStatsLogManager().logger()
+                        .log(LAUNCHER_ALLAPPS_FOCUSED_ITEM_SELECTED_WITH_IME);
                 // selectFocusedView should return SearchTargetEvent that is passed onto onClick
                 if (Launcher.getLauncher(mLauncher).getAppsView().selectFocusedView(v)) {
                     return true;
@@ -149,7 +153,7 @@
      * Resets the search bar state.
      */
     public void reset() {
-        mCb.clearSearchResult();
+        mCallback.clearSearchResult();
         mInput.reset();
         mQuery = null;
     }
@@ -167,31 +171,4 @@
     public boolean isSearchFieldFocused() {
         return mInput.isFocused();
     }
-
-    /**
-     * Callback for getting search results.
-     */
-    public interface Callbacks {
-
-        /**
-         * Called when the search from primary source is complete.
-         *
-         * @param items sorted list of search result adapter items
-         */
-        void onSearchResult(String query, ArrayList<AdapterItem> items);
-
-        /**
-         * Called when the search from secondary source is complete.
-         *
-         * @param items sorted list of search result adapter items
-         */
-        void onAppendSearchResult(String query, ArrayList<AdapterItem> items);
-
-        /**
-         * Called when the search results should be cleared.
-         */
-        void clearSearchResult();
-    }
-
-
 }
\ No newline at end of file
diff --git a/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java b/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java
index aef32d7..426fd0c 100644
--- a/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java
+++ b/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java
@@ -47,6 +47,7 @@
 import com.android.launcher3.allapps.SearchUiManager;
 import com.android.launcher3.anim.PropertySetter;
 import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.search.SearchCallback;
 
 import java.util.ArrayList;
 
@@ -54,7 +55,7 @@
  * Layout to contain the All-apps search UI.
  */
 public class AppsSearchContainerLayout extends ExtendedEditText
-        implements SearchUiManager, AllAppsSearchBarController.Callbacks,
+        implements SearchUiManager, SearchCallback<AdapterItem>,
         AllAppsStore.OnUpdateListener, Insettable {
 
     private final BaseDraggingActivity mLauncher;
diff --git a/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java b/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java
index 66bbd2e..4e213b0 100644
--- a/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java
+++ b/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java
@@ -19,14 +19,17 @@
 import android.os.Handler;
 
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.allapps.AllAppsGridAdapter.AdapterItem;
 import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.search.SearchAlgorithm;
+import com.android.launcher3.search.SearchCallback;
 
 import java.text.Collator;
 
 /**
  * The default search implementation.
  */
-public class DefaultAppSearchAlgorithm implements SearchAlgorithm {
+public class DefaultAppSearchAlgorithm implements SearchAlgorithm<AdapterItem> {
 
     protected final Handler mResultHandler;
     private final AppsSearchPipeline mAppsSearchPipeline;
@@ -45,7 +48,7 @@
 
     @Override
     public void doSearch(final String query,
-            final AllAppsSearchBarController.Callbacks callback) {
+            final SearchCallback<AdapterItem> callback) {
         mAppsSearchPipeline.query(query,
                 results -> mResultHandler.post(
                         () -> callback.onSearchResult(query, results)),
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 556aff6..e406e9b 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -95,9 +95,6 @@
     public static final BooleanFlag ENABLE_DEVICE_SEARCH = new DeviceFlag(
             "ENABLE_DEVICE_SEARCH", false, "Allows on device search in all apps");
 
-    public static final BooleanFlag DISABLE_INITIAL_IME_IN_ALLAPPS = getDebugFlag(
-            "DISABLE_INITIAL_IME_IN_ALLAPPS", false, "Disable default IME state in all apps");
-
     public static final BooleanFlag FOLDER_NAME_SUGGEST = new DeviceFlag(
             "FOLDER_NAME_SUGGEST", true,
             "Suggests folder names instead of blank text.");
@@ -161,7 +158,7 @@
             "ENABLE_SMARTSPACE_UNIVERSAL", false,
             "Replace Smartspace with a version rendered by System UI.");
 
-    public static final BooleanFlag ENABLE_SMARTSPACE_ENHANCED = getDebugFlag(
+    public static final BooleanFlag ENABLE_SMARTSPACE_ENHANCED = new DeviceFlag(
             "ENABLE_SMARTSPACE_ENHANCED", false,
             "Replace Smartspace with the enhanced version. "
               + "Ignored if ENABLE_SMARTSPACE_UNIVERSAL is enabled.");
@@ -266,6 +263,8 @@
         }
 
         public void addChangeListener(Context context, Runnable r) { }
+
+        public void removeChangeListener(Runnable r) {}
     }
 
     public static class DebugFlag extends BooleanFlag {
diff --git a/src/com/android/launcher3/dragndrop/LivePreviewWidgetCell.java b/src/com/android/launcher3/dragndrop/LivePreviewWidgetCell.java
index 71e10a8..be6a07f 100644
--- a/src/com/android/launcher3/dragndrop/LivePreviewWidgetCell.java
+++ b/src/com/android/launcher3/dragndrop/LivePreviewWidgetCell.java
@@ -1,6 +1,9 @@
 package com.android.launcher3.dragndrop;
 
+import static com.android.launcher3.Utilities.ATLEAST_S;
+
 import android.content.Context;
+import android.content.res.Resources;
 import android.graphics.Bitmap;
 import android.util.AttributeSet;
 import android.view.View;
@@ -62,9 +65,10 @@
 
     @Override
     public void applyFromCellItem(WidgetItem item, WidgetPreviewLoader loader) {
-        if (mPreview == null
+        if (ATLEAST_S
+                && mPreview == null
                 && item.widgetInfo != null
-                && item.widgetInfo.previewLayout != View.NO_ID) {
+                && item.widgetInfo.previewLayout != Resources.ID_NULL) {
             mPreview = new RemoteViews(item.widgetInfo.provider.getPackageName(),
                     item.widgetInfo.previewLayout);
         }
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index bf420d9..d554bb9 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -362,11 +362,11 @@
         @UiEvent(doc = "User closed the AllApps keyboard.")
         LAUNCHER_ALLAPPS_KEYBOARD_CLOSED(694),
 
-        @UiEvent(doc = "User switched to Main tab in AllApps screen.")
-        LAUNCHER_ALLAPPS_SWITCHED_TO_MAIN_TAB(695),
+        @UiEvent(doc = "User switched to AllApps Main/Personal tab by swiping left.")
+        LAUNCHER_ALLAPPS_SWIPE_TO_PERSONAL_TAB(695),
 
-        @UiEvent(doc = "User switched to Work tab in AllApps screen.")
-        LAUNCHER_ALLAPPS_SWITCHED_TO_WORK_TAB(696),
+        @UiEvent(doc = "User switched to AllApps Work tab by swiping right.")
+        LAUNCHER_ALLAPPS_SWIPE_TO_WORK_TAB(696),
 
         @UiEvent(doc = "Default event when dedicated UI event is not available for the user action"
                 + " on slice .")
@@ -392,6 +392,21 @@
 
         @UiEvent(doc = "User selected from a selection row of Slice.")
         LAUNCHER_SLICE_SELECTION_ACTION(707),
+
+        @UiEvent(doc = "IME is used for selecting the focused item on the AllApps screen.")
+        LAUNCHER_ALLAPPS_FOCUSED_ITEM_SELECTED_WITH_IME(718),
+
+        @UiEvent(doc = "User long-pressed on an AllApps item.")
+        LAUNCHER_ALLAPPS_ITEM_LONG_PRESSED(719),
+
+        @UiEvent(doc = "Launcher entered into AllApps state with device search enabled.")
+        LAUNCHER_ALLAPPS_ENTRY_WITH_DEVICE_SEARCH(720),
+
+        @UiEvent(doc = "User switched to AllApps Main/Personal tab by tapping on it.")
+        LAUNCHER_ALLAPPS_TAP_ON_PERSONAL_TAB(721),
+
+        @UiEvent(doc = "User switched to AllApps Work tab by tapping on it.")
+        LAUNCHER_ALLAPPS_TAP_ON_WORK_TAB(722),
         ;
 
         // ADD MORE
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 8438622..2cc1439 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -45,6 +45,7 @@
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.text.TextUtils;
+import android.util.ArrayMap;
 import android.util.Log;
 import android.util.LongSparseArray;
 import android.util.TimingLogger;
@@ -126,7 +127,7 @@
 
     private final UserManagerState mUserManagerState = new UserManagerState();
 
-    protected Map<ComponentKey, AppWidgetProviderInfo> mWidgetProvidersMap;
+    protected final Map<ComponentKey, AppWidgetProviderInfo> mWidgetProvidersMap = new ArrayMap<>();
 
     private boolean mStopped;
 
@@ -664,12 +665,13 @@
                             final boolean wasProviderReady = !c.hasRestoreFlag(
                                     LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY);
 
-                            if (mWidgetProvidersMap == null) {
-                                mWidgetProvidersMap = WidgetManagerHelper.getAllProvidersMap(
-                                        context);
+                            ComponentKey providerKey = new ComponentKey(component, c.user);
+                            if (!mWidgetProvidersMap.containsKey(providerKey)) {
+                                mWidgetProvidersMap.put(providerKey,
+                                        widgetHelper.findProvider(component, c.user));
                             }
-                            final AppWidgetProviderInfo provider = mWidgetProvidersMap.get(
-                                    new ComponentKey(component, c.user));
+                            final AppWidgetProviderInfo provider =
+                                    mWidgetProvidersMap.get(providerKey);
 
                             final boolean isProviderReady = isValidProvider(provider);
                             if (!isSafeMode && !customWidget &&
@@ -873,7 +875,6 @@
                     mBgDataModel.itemsIdMap.remove(folderId);
                 }
             }
-
             // Remove any ghost widgets
             LauncherSettings.Settings.call(contentResolver,
                     LauncherSettings.Settings.METHOD_REMOVE_GHOST_WIDGETS);
diff --git a/src/com/android/launcher3/model/ModelDelegate.java b/src/com/android/launcher3/model/ModelDelegate.java
index 92bea5b..13ec1ec 100644
--- a/src/com/android/launcher3/model/ModelDelegate.java
+++ b/src/com/android/launcher3/model/ModelDelegate.java
@@ -27,6 +27,8 @@
 import com.android.launcher3.shortcuts.ShortcutKey;
 import com.android.launcher3.util.ResourceBasedOverride;
 
+import java.io.FileDescriptor;
+import java.io.PrintWriter;
 import java.util.Map;
 
 /**
@@ -89,4 +91,11 @@
     @WorkerThread
     public void destroy() { }
 
+    /**
+     * Add data to a dumpsys request for Launcher (e.g. for bug reports).
+     *
+     * @see com.android.launcher3.Launcher#dump(java.lang.String, java.io.FileDescriptor,
+     *                                          java.io.PrintWriter, java.lang.String[])
+     **/
+    public void dump(String prefix, FileDescriptor fd, PrintWriter writer, String[] args) { }
 }
diff --git a/src/com/android/launcher3/model/WidgetItem.java b/src/com/android/launcher3/model/WidgetItem.java
index 37c089e..de2481a 100644
--- a/src/com/android/launcher3/model/WidgetItem.java
+++ b/src/com/android/launcher3/model/WidgetItem.java
@@ -43,4 +43,20 @@
         activityInfo = info;
         spanX = spanY = 1;
     }
+
+    /**
+     * Returns {@code true} if this {@link WidgetItem} has the same type as the given
+     * {@code otherItem}.
+     *
+     * For example, both items are widgets or both items are shortcuts.
+     */
+    public boolean hasSameType(WidgetItem otherItem) {
+        if (widgetInfo != null && otherItem.widgetInfo != null) {
+            return true;
+        }
+        if (activityInfo != null && otherItem.activityInfo != null) {
+            return true;
+        }
+        return false;
+    }
 }
diff --git a/src/com/android/launcher3/model/data/SearchActionItemInfo.java b/src/com/android/launcher3/model/data/SearchActionItemInfo.java
index e3b3b09..b3057d5 100644
--- a/src/com/android/launcher3/model/data/SearchActionItemInfo.java
+++ b/src/com/android/launcher3/model/data/SearchActionItemInfo.java
@@ -37,6 +37,7 @@
     public static final int FLAG_SHOULD_START_FOR_RESULT = FLAG_SHOULD_START | 1 << 2;
     public static final int FLAG_BADGE_WITH_PACKAGE = 1 << 3;
     public static final int FLAG_PRIMARY_ICON_FROM_TITLE = 1 << 4;
+    public static final int FLAG_BADGE_WITH_COMPONENT_NAME = 1 << 5;
 
     private final String mFallbackPackageName;
     private int mFlags = 0;
diff --git a/src/com/android/launcher3/popup/ArrowPopup.java b/src/com/android/launcher3/popup/ArrowPopup.java
index 56438d0..5a34d2a 100644
--- a/src/com/android/launcher3/popup/ArrowPopup.java
+++ b/src/com/android/launcher3/popup/ArrowPopup.java
@@ -26,11 +26,8 @@
 import android.animation.ValueAnimator;
 import android.content.Context;
 import android.content.res.Resources;
-import android.graphics.CornerPathEffect;
 import android.graphics.Outline;
-import android.graphics.Paint;
 import android.graphics.Rect;
-import android.graphics.drawable.ShapeDrawable;
 import android.util.AttributeSet;
 import android.util.Pair;
 import android.view.Gravity;
@@ -51,7 +48,6 @@
 import com.android.launcher3.anim.RevealOutlineAnimation;
 import com.android.launcher3.anim.RoundedRectRevealOutlineProvider;
 import com.android.launcher3.dragndrop.DragLayer;
-import com.android.launcher3.graphics.TriangleShape;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.views.BaseDragLayer;
 
@@ -72,7 +68,11 @@
     protected final T mLauncher;
     protected final boolean mIsRtl;
 
-    private final int mArrowOffset;
+    private final int mArrowOffsetVertical;
+    private final int mArrowOffsetHorizontal;
+    private final int mArrowWidth;
+    private final int mArrowHeight;
+    private final int mArrowPointRadius;
     private final View mArrow;
 
     protected boolean mIsLeftAligned;
@@ -103,11 +103,14 @@
 
         // Initialize arrow view
         final Resources resources = getResources();
-        final int arrowWidth = resources.getDimensionPixelSize(R.dimen.popup_arrow_width);
-        final int arrowHeight = resources.getDimensionPixelSize(R.dimen.popup_arrow_height);
+        mArrowWidth = resources.getDimensionPixelSize(R.dimen.popup_arrow_width);
+        mArrowHeight = resources.getDimensionPixelSize(R.dimen.popup_arrow_height);
         mArrow = new View(context);
-        mArrow.setLayoutParams(new DragLayer.LayoutParams(arrowWidth, arrowHeight));
-        mArrowOffset = resources.getDimensionPixelSize(R.dimen.popup_arrow_vertical_offset);
+        mArrow.setLayoutParams(new DragLayer.LayoutParams(mArrowWidth, mArrowHeight));
+        mArrowOffsetVertical = resources.getDimensionPixelSize(R.dimen.popup_arrow_vertical_offset);
+        mArrowOffsetHorizontal = resources.getDimensionPixelSize(
+                R.dimen.popup_arrow_horizontal_center_offset) - (mArrowWidth / 2);
+        mArrowPointRadius = resources.getDimensionPixelSize(R.dimen.popup_arrow_corner_radius);
     }
 
     public ArrowPopup(Context context, AttributeSet attrs) {
@@ -200,48 +203,33 @@
         orientAboutObject();
     }
 
-    private void addArrow() {
-        final Resources res = getResources();
-        final int arrowCenterOffset = res.getDimensionPixelSize(isAlignedWithStart()
-                ? R.dimen.popup_arrow_horizontal_center_start
-                : R.dimen.popup_arrow_horizontal_center_end);
-        final int halfArrowWidth = res.getDimensionPixelSize(R.dimen.popup_arrow_width) / 2;
-        getPopupContainer().addView(mArrow);
-        DragLayer.LayoutParams arrowLp = (DragLayer.LayoutParams) mArrow.getLayoutParams();
+    private int getArrowLeft() {
         if (mIsLeftAligned) {
-            mArrow.setX(getX() + arrowCenterOffset - halfArrowWidth);
-        } else {
-            mArrow.setX(getX() + getMeasuredWidth() - arrowCenterOffset - halfArrowWidth);
+            return mArrowOffsetHorizontal;
         }
+        return getMeasuredWidth() - mArrowOffsetHorizontal - mArrowWidth;
+    }
+
+    private void addArrow() {
+        getPopupContainer().addView(mArrow);
+        mArrow.setX(getX() + getArrowLeft());
 
         if (Gravity.isVertical(mGravity)) {
             // This is only true if there wasn't room for the container next to the icon,
             // so we centered it instead. In that case we don't want to showDefaultOptions the arrow.
             mArrow.setVisibility(INVISIBLE);
         } else {
-            ShapeDrawable arrowDrawable = new ShapeDrawable(TriangleShape.create(
-                    arrowLp.width, arrowLp.height, !mIsAboveIcon));
-            Paint arrowPaint = arrowDrawable.getPaint();
-            arrowPaint.setColor(Themes.getAttrColor(getContext(), R.attr.popupColorPrimary));
-            // The corner path effect won't be reflected in the shadow, but shouldn't be noticeable.
-            int radius = getResources().getDimensionPixelSize(R.dimen.popup_arrow_corner_radius);
-            arrowPaint.setPathEffect(new CornerPathEffect(radius));
-            mArrow.setBackground(arrowDrawable);
-            // Clip off the part of the arrow that is underneath the popup.
-            if (mIsAboveIcon) {
-                mArrow.setClipBounds(new Rect(0, -mArrowOffset, arrowLp.width, arrowLp.height));
-            } else {
-                mArrow.setClipBounds(new Rect(0, 0, arrowLp.width, arrowLp.height + mArrowOffset));
-            }
+            mArrow.setBackground(new RoundedArrowDrawable(
+                    mArrowWidth, mArrowHeight, mArrowPointRadius,
+                    mOutlineRadius, getMeasuredWidth(), getMeasuredHeight(),
+                    mArrowOffsetHorizontal, -mArrowOffsetVertical,
+                    !mIsAboveIcon, mIsLeftAligned,
+                    Themes.getAttrColor(getContext(), R.attr.popupColorPrimary)));
             mArrow.setElevation(getElevation());
         }
 
-        mArrow.setPivotX(arrowLp.width / 2);
-        mArrow.setPivotY(mIsAboveIcon ? arrowLp.height : 0);
-    }
-
-    protected boolean isAlignedWithStart() {
-        return mIsLeftAligned && !mIsRtl || !mIsLeftAligned && mIsRtl;
+        mArrow.setPivotX(mArrowWidth / 2.0f);
+        mArrow.setPivotY(mIsAboveIcon ? mArrowHeight : 0);
     }
 
     /**
@@ -274,8 +262,9 @@
      */
     private void orientAboutObject(boolean allowAlignLeft, boolean allowAlignRight) {
         measure(MeasureSpec.UNSPECIFIED, MeasureSpec.UNSPECIFIED);
+
         int width = getMeasuredWidth();
-        int extraVerticalSpace = mArrow.getLayoutParams().height + mArrowOffset
+        int extraVerticalSpace = mArrowHeight + mArrowOffsetVertical
                 + getResources().getDimensionPixelSize(R.dimen.popup_vertical_padding);
         int height = getMeasuredHeight() + extraVerticalSpace;
 
@@ -291,22 +280,7 @@
 
         // Offset x so that the arrow and shortcut icons are center-aligned with the original icon.
         int iconWidth = mTempRect.width();
-        Resources resources = getResources();
-        int xOffset;
-        if (isAlignedWithStart()) {
-            // Aligning with the shortcut icon.
-            int shortcutIconWidth = resources.getDimensionPixelSize(R.dimen.deep_shortcut_icon_size);
-            int shortcutPaddingStart = resources.getDimensionPixelSize(
-                    R.dimen.popup_padding_start);
-            xOffset = iconWidth / 2 - shortcutIconWidth / 2 - shortcutPaddingStart;
-        } else {
-            // Aligning with the drag handle.
-            int shortcutDragHandleWidth = resources.getDimensionPixelSize(
-                    R.dimen.deep_shortcut_drag_handle_size);
-            int shortcutPaddingEnd = resources.getDimensionPixelSize(
-                    R.dimen.popup_padding_end);
-            xOffset = iconWidth / 2 - shortcutDragHandleWidth / 2 - shortcutPaddingEnd;
-        }
+        int xOffset = iconWidth / 2 - mArrowOffsetHorizontal - mArrowWidth / 2;
         x += mIsLeftAligned ? xOffset : -xOffset;
 
         // Check whether we can still align as we originally wanted, now that we've calculated x.
@@ -375,12 +349,14 @@
         FrameLayout.LayoutParams arrowLp = (FrameLayout.LayoutParams) mArrow.getLayoutParams();
         if (mIsAboveIcon) {
             arrowLp.gravity = lp.gravity = Gravity.BOTTOM;
-            lp.bottomMargin = getPopupContainer().getHeight() - y - getMeasuredHeight() - insets.top;
-            arrowLp.bottomMargin = lp.bottomMargin - arrowLp.height - mArrowOffset - insets.bottom;
+            lp.bottomMargin =
+                    getPopupContainer().getHeight() - y - getMeasuredHeight() - insets.top;
+            arrowLp.bottomMargin =
+                    lp.bottomMargin - arrowLp.height - mArrowOffsetVertical - insets.bottom;
         } else {
             arrowLp.gravity = lp.gravity = Gravity.TOP;
             lp.topMargin = y + insets.top;
-            arrowLp.topMargin = lp.topMargin - insets.top - arrowLp.height - mArrowOffset;
+            arrowLp.topMargin = lp.topMargin - insets.top - arrowLp.height - mArrowOffsetVertical;
         }
     }
 
@@ -529,22 +505,13 @@
     protected void onCreateCloseAnimation(AnimatorSet anim) { }
 
     private RoundedRectRevealOutlineProvider createOpenCloseOutlineProvider() {
-        Resources res = getResources();
-        int arrowCenterX = res.getDimensionPixelSize(mIsLeftAligned ^ mIsRtl ?
-                R.dimen.popup_arrow_horizontal_center_start:
-                R.dimen.popup_arrow_horizontal_center_end);
-        int halfArrowWidth = res.getDimensionPixelSize(R.dimen.popup_arrow_width) / 2;
-        float arrowCornerRadius = res.getDimension(R.dimen.popup_arrow_corner_radius);
-        if (!mIsLeftAligned) {
-            arrowCenterX = getMeasuredWidth() - arrowCenterX;
-        }
+        int arrowLeft = getArrowLeft();
         int arrowCenterY = mIsAboveIcon ? getMeasuredHeight() : 0;
 
-        mStartRect.set(arrowCenterX - halfArrowWidth, arrowCenterY, arrowCenterX + halfArrowWidth,
-                arrowCenterY);
+        mStartRect.set(arrowLeft, arrowCenterY, arrowLeft + mArrowWidth, arrowCenterY);
 
-        return new RoundedRectRevealOutlineProvider
-                (arrowCornerRadius, mOutlineRadius, mStartRect, mEndRect);
+        return new RoundedRectRevealOutlineProvider(
+                mArrowPointRadius, mOutlineRadius, mStartRect, mEndRect);
     }
 
     /**
diff --git a/src/com/android/launcher3/popup/RoundedArrowDrawable.java b/src/com/android/launcher3/popup/RoundedArrowDrawable.java
new file mode 100644
index 0000000..e662d5c
--- /dev/null
+++ b/src/com/android/launcher3/popup/RoundedArrowDrawable.java
@@ -0,0 +1,161 @@
+/*
+ * Copyright (C) 2021 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.popup;
+
+import static java.lang.Math.atan;
+import static java.lang.Math.cos;
+import static java.lang.Math.sin;
+import static java.lang.Math.toDegrees;
+
+import android.graphics.Canvas;
+import android.graphics.ColorFilter;
+import android.graphics.Matrix;
+import android.graphics.Outline;
+import android.graphics.Paint;
+import android.graphics.Path;
+import android.graphics.PixelFormat;
+import android.graphics.drawable.Drawable;
+
+/**
+ * A drawable for a very specific purpose. Used for the caret arrow on a rounded rectangle popup
+ * bubble.
+ * Draws a triangle with one rounded tip, the opposite edge is clipped by the body of the popup
+ * so there is no overlap when drawing them together.
+ */
+public class RoundedArrowDrawable extends Drawable {
+
+    private final Path mPath;
+    private final Paint mPaint;
+
+    /**
+     * Default constructor.
+     *
+     * @param width of the arrow.
+     * @param height of the arrow.
+     * @param radius of the tip of the arrow.
+     * @param popupRadius of the rect to clip this by.
+     * @param popupWidth of the rect to clip this by.
+     * @param popupHeight of the rect to clip this by.
+     * @param arrowOffsetX from the edge of the popup to the arrow.
+     * @param arrowOffsetY how much the arrow will overlap the popup.
+     * @param isPointingUp or not.
+     * @param leftAligned or false for right aligned.
+     * @param color to draw the triangle.
+     */
+    public RoundedArrowDrawable(float width, float height, float radius, float popupRadius,
+            float popupWidth, float popupHeight,
+            float arrowOffsetX, float arrowOffsetY, boolean isPointingUp, boolean leftAligned,
+            int color) {
+        mPath = new Path();
+        mPaint = new Paint();
+        mPaint.setColor(color);
+        mPaint.setStyle(Paint.Style.FILL);
+        mPaint.setAntiAlias(true);
+
+        // Make the drawable with the triangle pointing down and positioned on the left..
+        addDownPointingRoundedTriangleToPath(width, height, radius, mPath);
+        clipPopupBodyFromPath(popupRadius, popupWidth, popupHeight, arrowOffsetX, arrowOffsetY,
+                mPath);
+
+        // ... then flip it horizontal or vertical based on where it will be used.
+        Matrix pathTransform = new Matrix();
+        pathTransform.setScale(
+                leftAligned ? 1 : -1, isPointingUp ? -1 : 1, width * 0.5f, height * 0.5f);
+        mPath.transform(pathTransform);
+    }
+
+    @Override
+    public void draw(Canvas canvas) {
+        canvas.drawPath(mPath, mPaint);
+    }
+
+    @Override
+    public void getOutline(Outline outline) {
+        outline.setPath(mPath);
+    }
+
+    @Override
+    public int getOpacity() {
+        return PixelFormat.TRANSLUCENT;
+    }
+
+    @Override
+    public void setAlpha(int i) {
+        mPaint.setAlpha(i);
+    }
+
+    @Override
+    public void setColorFilter(ColorFilter colorFilter) {
+        mPaint.setColorFilter(colorFilter);
+    }
+
+    private static void addDownPointingRoundedTriangleToPath(float width, float height,
+            float radius, Path path) {
+        // Calculated for the arrow pointing down, will be flipped later if needed.
+
+        // Theta is half of the angle inside the triangle tip
+        float tanTheta = width / (2.0f * height);
+        float theta = (float) atan(tanTheta);
+
+        // Some trigonometry to find the center of the circle for the rounded tip
+        float roundedPointCenterY = (float) (height - (radius / sin(theta)));
+
+        // p is the distance along the triangle side to the intersection with the point circle
+        float p = radius / tanTheta;
+        float lineRoundPointIntersectFromCenter = (float) (p * sin(theta));
+        float lineRoundPointIntersectFromTop = (float) (height - (p * cos(theta)));
+
+        float centerX = width / 2.0f;
+        float thetaDeg = (float) toDegrees(theta);
+
+        path.reset();
+        path.moveTo(0, 0);
+        // Draw the top
+        path.lineTo(width, 0);
+        // Draw the right side up to the circle intersection
+        path.lineTo(
+                centerX + lineRoundPointIntersectFromCenter,
+                lineRoundPointIntersectFromTop);
+        // Draw the rounded point
+        path.arcTo(
+                centerX - radius,
+                roundedPointCenterY - radius,
+                centerX + radius,
+                roundedPointCenterY + radius,
+                thetaDeg,
+                180 - (2 * thetaDeg),
+                false);
+        // Draw the left edge to close
+        path.lineTo(0, 0);
+        path.close();
+    }
+
+    private static void clipPopupBodyFromPath(float popupRadius, float popupWidth,
+            float popupHeight, float arrowOffsetX, float arrowOffsetY, Path path) {
+        // Make a path that is used to clip the triangle, this represents the body of the popup
+        Path clipPiece = new Path();
+        clipPiece.addRoundRect(
+                0, 0, popupWidth, popupHeight,
+                popupRadius, popupRadius, Path.Direction.CW);
+        // clipping is performed as if the arrow is pointing down and positioned on the left, the
+        // resulting path will be flipped as needed later.
+        // The extra 0.5 in the vertical offset is to close the gap between this anti-aliased object
+        // and the anti-aliased body of the popup.
+        clipPiece.offset(-arrowOffsetX, -popupHeight + arrowOffsetY - 0.5f);
+        path.op(clipPiece, Path.Op.DIFFERENCE);
+    }
+}
diff --git a/src/com/android/launcher3/qsb/QsbContainerView.java b/src/com/android/launcher3/qsb/QsbContainerView.java
index 289e0d8..459aefe 100644
--- a/src/com/android/launcher3/qsb/QsbContainerView.java
+++ b/src/com/android/launcher3/qsb/QsbContainerView.java
@@ -20,6 +20,8 @@
 import static android.appwidget.AppWidgetManager.EXTRA_APPWIDGET_ID;
 import static android.appwidget.AppWidgetManager.EXTRA_APPWIDGET_PROVIDER;
 
+import static com.android.launcher3.Utilities.ATLEAST_S;
+
 import android.app.Activity;
 import android.app.Fragment;
 import android.app.SearchManager;
@@ -30,6 +32,7 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
+import android.graphics.PointF;
 import android.graphics.Rect;
 import android.os.Bundle;
 import android.provider.Settings;
@@ -50,6 +53,8 @@
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.graphics.FragmentWithPreview;
 
+import java.util.ArrayList;
+
 /**
  * A frame layout which contains a QSB. This internally uses fragment to bind the view, which
  * allows it to contain the logic for {@link Fragment#startActivityForResult(Intent, int)}.
@@ -294,12 +299,16 @@
             InvariantDeviceProfile idp = LauncherAppState.getIDP(getContext());
 
             Bundle opts = new Bundle();
-            Rect size = AppWidgetResizeFrame.getWidgetSizeRanges(getContext(),
-                    idp.numColumns, 1, null);
+            ArrayList<PointF> sizes = AppWidgetResizeFrame
+                    .getWidgetSizes(getContext(), idp.numColumns, 1);
+            Rect size = AppWidgetResizeFrame.getMinMaxSizes(sizes, null /* outRect */);
             opts.putInt(AppWidgetManager.OPTION_APPWIDGET_MIN_WIDTH, size.left);
             opts.putInt(AppWidgetManager.OPTION_APPWIDGET_MIN_HEIGHT, size.top);
             opts.putInt(AppWidgetManager.OPTION_APPWIDGET_MAX_WIDTH, size.right);
             opts.putInt(AppWidgetManager.OPTION_APPWIDGET_MAX_HEIGHT, size.bottom);
+            if (ATLEAST_S) {
+                opts.putParcelableArrayList(AppWidgetManager.OPTION_APPWIDGET_SIZES, sizes);
+            }
             return opts;
         }
 
diff --git a/src/com/android/launcher3/allapps/search/SearchAlgorithm.java b/src/com/android/launcher3/search/SearchAlgorithm.java
similarity index 76%
rename from src/com/android/launcher3/allapps/search/SearchAlgorithm.java
rename to src/com/android/launcher3/search/SearchAlgorithm.java
index c409b1c..1665354 100644
--- a/src/com/android/launcher3/allapps/search/SearchAlgorithm.java
+++ b/src/com/android/launcher3/search/SearchAlgorithm.java
@@ -13,17 +13,19 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.launcher3.allapps.search;
+package com.android.launcher3.search;
 
 /**
  * An interface for handling search.
+ *
+ * @param <T> Search Result type
  */
-public interface SearchAlgorithm {
+public interface SearchAlgorithm<T> {
 
     /**
-     * Performs search and sends the result to the callback.
+     * Performs search and sends the result to {@link SearchCallback}.
      */
-    void doSearch(String query, AllAppsSearchBarController.Callbacks callback);
+    void doSearch(String query, SearchCallback<T> callback);
 
     /**
      * Cancels any active request.
diff --git a/src/com/android/launcher3/search/SearchCallback.java b/src/com/android/launcher3/search/SearchCallback.java
new file mode 100644
index 0000000..5796116
--- /dev/null
+++ b/src/com/android/launcher3/search/SearchCallback.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2017 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.search;
+
+import java.util.ArrayList;
+
+/**
+ * An interface for receiving search results.
+ *
+ * @param <T> Search Result type
+ */
+public interface SearchCallback<T> {
+
+    /**
+     * Called when the search from primary source is complete.
+     *
+     * @param items list of search results
+     */
+    void onSearchResult(String query, ArrayList<T> items);
+
+    /**
+     * Called when the search from secondary source is complete.
+     *
+     * @param items list of search results
+     */
+    void onAppendSearchResult(String query, ArrayList<T> items);
+
+    /**
+     * Called when the search results should be cleared.
+     */
+    void clearSearchResult();
+}
+
diff --git a/src/com/android/launcher3/settings/SettingsActivity.java b/src/com/android/launcher3/settings/SettingsActivity.java
index ac8dac5..f03065c 100644
--- a/src/com/android/launcher3/settings/SettingsActivity.java
+++ b/src/com/android/launcher3/settings/SettingsActivity.java
@@ -129,6 +129,7 @@
         private String mHighLightKey;
         private boolean mPreferenceHighlighted = false;
         private NotificationDotsPreference mNotificationSettingsChangedListener;
+        private Preference mDeveloperOptionPref;
 
         @Override
         public void onCreatePreferences(Bundle savedInstanceState, String rootKey) {
@@ -202,18 +203,37 @@
                     return FeatureFlags.showFlagTogglerUi(getContext());
 
                 case DEVELOPER_OPTIONS_KEY:
-                    // Show if plugins are enabled or flag UI is enabled.
-                    return FeatureFlags.showFlagTogglerUi(getContext()) ||
-                            PluginManagerWrapper.hasPlugins(getContext());
+                    mDeveloperOptionPref = preference;
+                    return updateDeveloperOption();
             }
 
             return true;
         }
 
+        /**
+         * Show if plugins are enabled or flag UI is enabled.
+         * @return True if we should show the preference option.
+         */
+        private boolean updateDeveloperOption() {
+            boolean showPreference = FeatureFlags.showFlagTogglerUi(getContext())
+                    || PluginManagerWrapper.hasPlugins(getContext());
+            if (mDeveloperOptionPref != null) {
+                mDeveloperOptionPref.setEnabled(showPreference);
+                if (showPreference) {
+                    getPreferenceScreen().addPreference(mDeveloperOptionPref);
+                } else {
+                    getPreferenceScreen().removePreference(mDeveloperOptionPref);
+                }
+            }
+            return showPreference;
+        }
+
         @Override
         public void onResume() {
             super.onResume();
 
+            updateDeveloperOption();
+
             if (isAdded() && !mPreferenceHighlighted) {
                 PreferenceHighlighter highlighter = createHighlighter();
                 if (highlighter != null) {
diff --git a/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java b/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java
index 31adc08..65df614 100644
--- a/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java
+++ b/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java
@@ -270,7 +270,6 @@
         mFlingBlockCheck.unblockFling();
         // Must be called after all the animation controllers have been paused
         if (FeatureFlags.ENABLE_DEVICE_SEARCH.get()
-                && !FeatureFlags.DISABLE_INITIAL_IME_IN_ALLAPPS.get()
                 && BuildCompat.isAtLeastR()
                 && (mToState == ALL_APPS || mToState == NORMAL)) {
             mLauncher.getAllAppsController().getInsetController().onDragStart(
diff --git a/src/com/android/launcher3/touch/ItemLongClickListener.java b/src/com/android/launcher3/touch/ItemLongClickListener.java
index 7baeab8..919673f 100644
--- a/src/com/android/launcher3/touch/ItemLongClickListener.java
+++ b/src/com/android/launcher3/touch/ItemLongClickListener.java
@@ -21,6 +21,7 @@
 import static com.android.launcher3.LauncherState.ALL_APPS;
 import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.LauncherState.OVERVIEW;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_ITEM_LONG_PRESSED;
 
 import android.view.View;
 import android.view.View.OnLongClickListener;
@@ -32,6 +33,7 @@
 import com.android.launcher3.dragndrop.DragController;
 import com.android.launcher3.dragndrop.DragOptions;
 import com.android.launcher3.folder.Folder;
+import com.android.launcher3.logging.StatsLogManager.StatsLogger;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.TestProtocol;
@@ -86,6 +88,12 @@
         if (!launcher.isInState(ALL_APPS) && !launcher.isInState(OVERVIEW)) return false;
         if (launcher.getWorkspace().isSwitchingState()) return false;
 
+        StatsLogger logger = launcher.getStatsLogManager().logger();
+        if (v.getTag() instanceof ItemInfo) {
+            logger.withItemInfo((ItemInfo) v.getTag());
+        }
+        logger.log(LAUNCHER_ALLAPPS_ITEM_LONG_PRESSED);
+
         // Start the drag
         final DragController dragController = launcher.getDragController();
         dragController.addDragListener(new DragController.DragListener() {
diff --git a/src/com/android/launcher3/views/RecyclerViewFastScroller.java b/src/com/android/launcher3/views/RecyclerViewFastScroller.java
index 804fb3e..ae34257 100644
--- a/src/com/android/launcher3/views/RecyclerViewFastScroller.java
+++ b/src/com/android/launcher3/views/RecyclerViewFastScroller.java
@@ -34,6 +34,7 @@
 import android.view.ViewConfiguration;
 import android.widget.TextView;
 
+import androidx.annotation.Nullable;
 import androidx.recyclerview.widget.RecyclerView;
 
 import com.android.launcher3.BaseRecyclerView;
@@ -99,6 +100,7 @@
     private boolean mIsThumbDetached;
     private final boolean mCanThumbDetach;
     private boolean mIgnoreDragGesture;
+    private boolean mIsRecyclerViewFirstChildInParent = true;
 
     // This is the offset from the top of the scrollbar when the user first starts touching.  To
     // prevent jumping, this offset is applied as the user scrolls.
@@ -112,6 +114,7 @@
 
     protected BaseRecyclerView mRv;
     private RecyclerView.OnScrollListener mOnScrollListener;
+    @Nullable private OnFastScrollChangeListener mOnFastScrollChangeListener;
 
     private int mDownX;
     private int mDownY;
@@ -188,6 +191,9 @@
         updatePopupY(y);
         mThumbOffsetY = y;
         invalidate();
+        if (mOnFastScrollChangeListener != null) {
+            mOnFastScrollChangeListener.onThumbOffsetYChanged(mThumbOffsetY);
+        }
     }
 
     public int getThumbOffsetY() {
@@ -391,7 +397,9 @@
             return false;
         }
         getHitRect(sTempRect);
-        sTempRect.top += mRv.getScrollBarTop();
+        if (mIsRecyclerViewFirstChildInParent) {
+            sTempRect.top += mRv.getScrollBarTop();
+        }
         if (outOffset != null) {
             outOffset.set(sTempRect.left, sTempRect.top);
         }
@@ -404,4 +412,23 @@
         // alpha is so low, it does not matter.
         return false;
     }
+
+    public void setIsRecyclerViewFirstChildInParent(boolean isRecyclerViewFirstChildInParent) {
+        mIsRecyclerViewFirstChildInParent = isRecyclerViewFirstChildInParent;
+    }
+
+    public void setOnFastScrollChangeListener(
+            @Nullable OnFastScrollChangeListener onFastScrollChangeListener) {
+        mOnFastScrollChangeListener = onFastScrollChangeListener;
+    }
+
+    /**
+     * A callback that is invoked when there is a scroll change in {@link RecyclerViewFastScroller}.
+     */
+    public interface OnFastScrollChangeListener {
+        /**
+         * Called when the thumb offset vertical position, in pixels, has changed to {@code y}.
+         */
+        void onThumbOffsetYChanged(int y);
+    }
 }
diff --git a/src/com/android/launcher3/views/WorkEduView.java b/src/com/android/launcher3/views/WorkEduView.java
index 03d3026..c8cf627 100644
--- a/src/com/android/launcher3/views/WorkEduView.java
+++ b/src/com/android/launcher3/views/WorkEduView.java
@@ -168,8 +168,8 @@
     }
 
     private AllAppsPagedView getAllAppsPagedView() {
-        View v =  mLauncher.getAppsView().getContentView();
-        return  (v instanceof AllAppsPagedView)  ? (AllAppsPagedView) v : null;
+        View v = mLauncher.getAppsView().getContentView();
+        return (v instanceof AllAppsPagedView) ? (AllAppsPagedView) v : null;
     }
 
     /**
diff --git a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
index 780a1a1..41098f9 100644
--- a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
@@ -16,9 +16,12 @@
 
 package com.android.launcher3.widget;
 
+import static com.android.launcher3.Utilities.ATLEAST_S;
+
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
 import android.content.res.Configuration;
+import android.graphics.PointF;
 import android.os.Handler;
 import android.os.SystemClock;
 import android.util.SparseBooleanArray;
@@ -70,8 +73,6 @@
     private boolean mIsAutoAdvanceRegistered;
     private Runnable mAutoAdvanceRunnable;
 
-
-
     public LauncherAppWidgetHostView(Context context) {
         super(context);
         mLauncher = Launcher.getLauncher(context);
@@ -219,6 +220,16 @@
     }
 
     @Override
+    protected void onSizeChanged(int w, int h, int oldw, int oldh) {
+        super.onSizeChanged(w, h, oldw, oldh);
+
+        if (ATLEAST_S) {
+            float density = getContext().getResources().getDisplayMetrics().density;
+            setCurrentSize(new PointF(w / density, h / density));
+        }
+    }
+
+    @Override
     public void onInitializeAccessibilityNodeInfo(AccessibilityNodeInfo info) {
         super.onInitializeAccessibilityNodeInfo(info);
         info.setClassName(getClass().getName());
diff --git a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
index ca47728..8c3206d 100644
--- a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
@@ -22,6 +22,7 @@
 import android.content.Context;
 import android.graphics.Canvas;
 import android.graphics.Color;
+import android.graphics.PointF;
 import android.graphics.PorterDuff;
 import android.graphics.Rect;
 import android.graphics.drawable.Drawable;
@@ -46,6 +47,8 @@
 import com.android.launcher3.touch.ItemClickHandler;
 import com.android.launcher3.util.Themes;
 
+import java.util.List;
+
 public class PendingAppWidgetHostView extends LauncherAppWidgetHostView
         implements OnClickListener, ItemInfoUpdateReceiver {
     private static final float SETUP_ICON_SIZE_FACTOR = 2f / 5;
@@ -109,6 +112,11 @@
     }
 
     @Override
+    public void updateAppWidgetSize(Bundle newOptions, List<PointF> sizes) {
+        // No-op
+    }
+
+    @Override
     protected View getDefaultView() {
         View defaultView = mInflater.inflate(R.layout.appwidget_not_ready, this, false);
         defaultView.setOnClickListener(this);
diff --git a/src/com/android/launcher3/widget/WidgetCell.java b/src/com/android/launcher3/widget/WidgetCell.java
index bef91d2..c6593f8 100644
--- a/src/com/android/launcher3/widget/WidgetCell.java
+++ b/src/com/android/launcher3/widget/WidgetCell.java
@@ -24,7 +24,6 @@
 import android.view.MotionEvent;
 import android.view.View;
 import android.view.View.OnLayoutChangeListener;
-import android.view.ViewGroup;
 import android.view.ViewPropertyAnimator;
 import android.view.accessibility.AccessibilityNodeInfo;
 import android.widget.LinearLayout;
@@ -60,6 +59,8 @@
     /** Widget preview width is calculated by multiplying this factor to the widget cell width. */
     private static final float PREVIEW_SCALE = 0.8f;
 
+    private int mPreviewWidth;
+    private int mPreviewHeight;
     protected int mPresetPreviewSize;
     private int mCellSize;
 
@@ -106,6 +107,7 @@
     private void setContainerWidth() {
         mCellSize = (int) (mDeviceProfile.allAppsIconSizePx * WIDTH_SCALE);
         mPresetPreviewSize = (int) (mCellSize * PREVIEW_SCALE);
+        mPreviewWidth = mPreviewHeight = mPresetPreviewSize;
     }
 
     @Override
@@ -128,6 +130,7 @@
         mWidgetImage.setBitmap(null, null);
         mWidgetName.setText(null);
         mWidgetDims.setText(null);
+        mPreviewWidth = mPreviewHeight = mPresetPreviewSize;
 
         if (mActiveRequest != null) {
             mActiveRequest.cancel();
@@ -181,6 +184,11 @@
             return;
         }
         if (bitmap != null) {
+            LayoutParams layoutParams = (LayoutParams) mWidgetImage.getLayoutParams();
+            layoutParams.width = bitmap.getWidth();
+            layoutParams.height = bitmap.getHeight();
+            mWidgetImage.setLayoutParams(layoutParams);
+
             mWidgetImage.setBitmap(bitmap, mWidgetPreviewLoader.getBadgeForUser(mItem.user,
                     BaseIconFactory.getBadgeSizeForIconSize(mDeviceProfile.allAppsIconSizePx)));
             if (mAnimatePreview) {
@@ -197,8 +205,16 @@
         if (mActiveRequest != null) {
             return;
         }
-        mActiveRequest = mWidgetPreviewLoader.getPreview(
-                mItem, mPresetPreviewSize, mPresetPreviewSize, this);
+        mActiveRequest = mWidgetPreviewLoader.getPreview(mItem, mPreviewWidth, mPreviewHeight,
+                this);
+    }
+
+    /** Sets the widget preview image size in number of cells. */
+    public void setPreviewSize(int spanX, int spanY) {
+        int padding = 2 * getResources()
+                .getDimensionPixelSize(R.dimen.widget_preview_shortcut_padding);
+        mPreviewWidth = mDeviceProfile.cellWidthPx * spanX + padding;
+        mPreviewHeight = mDeviceProfile.cellHeightPx * spanY + padding;
     }
 
     @Override
@@ -233,12 +249,6 @@
     }
 
     @Override
-    public void setLayoutParams(ViewGroup.LayoutParams params) {
-        params.width = params.height = mCellSize;
-        super.setLayoutParams(params);
-    }
-
-    @Override
     public CharSequence getAccessibilityClassName() {
         return WidgetCell.class.getName();
     }
diff --git a/src/com/android/launcher3/widget/WidgetHostViewLoader.java b/src/com/android/launcher3/widget/WidgetHostViewLoader.java
index c022374..2438bdf 100644
--- a/src/com/android/launcher3/widget/WidgetHostViewLoader.java
+++ b/src/com/android/launcher3/widget/WidgetHostViewLoader.java
@@ -3,6 +3,7 @@
 import android.appwidget.AppWidgetHostView;
 import android.appwidget.AppWidgetManager;
 import android.content.Context;
+import android.graphics.PointF;
 import android.graphics.Rect;
 import android.os.Bundle;
 import android.os.Handler;
@@ -18,6 +19,8 @@
 import com.android.launcher3.dragndrop.DragOptions;
 import com.android.launcher3.util.Thunk;
 
+import java.util.ArrayList;
+
 public class WidgetHostViewLoader implements DragController.DragListener {
     private static final String TAG = "WidgetHostViewLoader";
     private static final boolean LOGD = false;
@@ -152,24 +155,28 @@
     }
 
     public static Bundle getDefaultOptionsForWidget(Context context, PendingAddWidgetInfo info) {
-        Rect rect = new Rect();
-        AppWidgetResizeFrame.getWidgetSizeRanges(context, info.spanX, info.spanY, rect);
+        ArrayList<PointF> sizes = AppWidgetResizeFrame
+                .getWidgetSizes(context, info.spanX, info.spanY);
+
         Rect padding = AppWidgetHostView.getDefaultPaddingForWidget(context,
                 info.componentName, null);
-
         float density = context.getResources().getDisplayMetrics().density;
-        int xPaddingDips = (int) ((padding.left + padding.right) / density);
-        int yPaddingDips = (int) ((padding.top + padding.bottom) / density);
+        float xPaddingDips = (padding.left + padding.right) / density;
+        float yPaddingDips = (padding.top + padding.bottom) / density;
+
+        for (PointF size : sizes) {
+            size.x = Math.max(0.f, size.x - xPaddingDips);
+            size.y = Math.max(0.f, size.y - yPaddingDips);
+        }
+
+        Rect rect = AppWidgetResizeFrame.getMinMaxSizes(sizes, null /* outRect */);
 
         Bundle options = new Bundle();
-        options.putInt(AppWidgetManager.OPTION_APPWIDGET_MIN_WIDTH,
-                rect.left - xPaddingDips);
-        options.putInt(AppWidgetManager.OPTION_APPWIDGET_MIN_HEIGHT,
-                rect.top - yPaddingDips);
-        options.putInt(AppWidgetManager.OPTION_APPWIDGET_MAX_WIDTH,
-                rect.right - xPaddingDips);
-        options.putInt(AppWidgetManager.OPTION_APPWIDGET_MAX_HEIGHT,
-                rect.bottom - yPaddingDips);
+        options.putInt(AppWidgetManager.OPTION_APPWIDGET_MIN_WIDTH, rect.left);
+        options.putInt(AppWidgetManager.OPTION_APPWIDGET_MIN_HEIGHT, rect.top);
+        options.putInt(AppWidgetManager.OPTION_APPWIDGET_MAX_WIDTH, rect.right);
+        options.putInt(AppWidgetManager.OPTION_APPWIDGET_MAX_HEIGHT, rect.bottom);
+        options.putParcelableArrayList(AppWidgetManager.OPTION_APPWIDGET_SIZES, sizes);
         return options;
     }
 }
diff --git a/src/com/android/launcher3/widget/WidgetManagerHelper.java b/src/com/android/launcher3/widget/WidgetManagerHelper.java
index c0c5c48..7248971 100644
--- a/src/com/android/launcher3/widget/WidgetManagerHelper.java
+++ b/src/com/android/launcher3/widget/WidgetManagerHelper.java
@@ -31,14 +31,11 @@
 import com.android.launcher3.model.WidgetsModel;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.pm.UserCache;
-import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
 import java.util.Collections;
 import java.util.List;
-import java.util.Map;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -122,15 +119,6 @@
                 appWidgetId).getBoolean(WIDGET_OPTION_RESTORE_COMPLETED);
     }
 
-    public static Map<ComponentKey, AppWidgetProviderInfo> getAllProvidersMap(Context context) {
-        if (WidgetsModel.GO_DISABLE_WIDGETS) {
-            return Collections.emptyMap();
-        }
-        return allWidgetsSteam(context).collect(
-                        Collectors.toMap(info -> new ComponentKey(info.provider, info.getProfile()),
-                        Function.identity()));
-    }
-
     private static Stream<AppWidgetProviderInfo> allWidgetsSteam(Context context) {
         AppWidgetManager awm = context.getSystemService(AppWidgetManager.class);
         return Stream.concat(
diff --git a/src/com/android/launcher3/widget/WidgetsBottomSheet.java b/src/com/android/launcher3/widget/WidgetsBottomSheet.java
index 223cda2..6abbf21 100644
--- a/src/com/android/launcher3/widget/WidgetsBottomSheet.java
+++ b/src/com/android/launcher3/widget/WidgetsBottomSheet.java
@@ -24,22 +24,25 @@
 import android.util.AttributeSet;
 import android.util.IntProperty;
 import android.util.Pair;
-import android.view.Gravity;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewGroup;
 import android.view.animation.Interpolator;
+import android.widget.ScrollView;
+import android.widget.TableLayout;
+import android.widget.TableRow;
 import android.widget.TextView;
 
+import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Insettable;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.R;
-import com.android.launcher3.ResourceUtils;
 import com.android.launcher3.anim.PendingAnimation;
 import com.android.launcher3.dragndrop.LivePreviewWidgetCell;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.widget.util.WidgetsTableUtils;
 
 import java.util.List;
 
@@ -65,6 +68,8 @@
     private static final int DEFAULT_CLOSE_DURATION = 200;
     private ItemInfo mOriginalItemInfo;
     private Rect mInsets;
+    private final int mMaxTableHeight;
+    private int mMaxHorizontalSpan = 4;
 
     public WidgetsBottomSheet(Context context, AttributeSet attrs) {
         this(context, attrs, 0);
@@ -75,12 +80,41 @@
         setWillNotDraw(false);
         mInsets = new Rect();
         mContent = this;
+        DeviceProfile deviceProfile = mLauncher.getDeviceProfile();
+        // Set the max table height to 2 / 3 of the grid height so that the bottom picker won't
+        // take over the entire view vertically.
+        mMaxTableHeight = deviceProfile.inv.numRows * 2 / 3  * deviceProfile.cellHeightPx;
+    }
+
+    @Override
+    protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
+        super.onMeasure(widthMeasureSpec, heightMeasureSpec);
+
+        int paddingPx = 2 * getResources().getDimensionPixelOffset(
+                R.dimen.widget_cell_horizontal_padding);
+        int maxHorizontalSpan = findViewById(R.id.widgets_table).getMeasuredWidth()
+                / (mLauncher.getDeviceProfile().cellWidthPx + paddingPx);
+
+        if (mMaxHorizontalSpan != maxHorizontalSpan) {
+            // Ensure the table layout is showing widgets in the right column after measure.
+            mMaxHorizontalSpan = maxHorizontalSpan;
+            onWidgetsBound();
+        }
     }
 
     @Override
     protected void onLayout(boolean changed, int l, int t, int r, int b) {
         super.onLayout(changed, l, t, r, b);
         setTranslationShift(mTranslationShift);
+
+        // Ensure the scroll view height is not larger than mMaxTableHeight, which is a value
+        // smaller than the entire screen height.
+        ScrollView widgetsTableScrollView = findViewById(R.id.widgets_table_scroll_view);
+        if (widgetsTableScrollView.getMeasuredHeight() > mMaxTableHeight) {
+            ViewGroup.LayoutParams layoutParams = widgetsTableScrollView.getLayoutParams();
+            layoutParams.height = mMaxTableHeight;
+            widgetsTableScrollView.setLayoutParams(layoutParams);
+        }
     }
 
     public void populateAndShow(ItemInfo itemInfo) {
@@ -101,39 +135,21 @@
                         mOriginalItemInfo.getTargetComponent().getPackageName(),
                         mOriginalItemInfo.user));
 
-        ViewGroup widgetRow = findViewById(R.id.widgets);
-        ViewGroup widgetCells = widgetRow.findViewById(R.id.widgets_cell_list);
+        TableLayout widgetsTable = findViewById(R.id.widgets_table);
+        widgetsTable.removeAllViews();
 
-        widgetCells.removeAllViews();
-
-        for (int i = 0; i < widgets.size(); i++) {
-            WidgetCell widget = addItemCell(widgetCells);
-            widget.applyFromCellItem(widgets.get(i), LauncherAppState.getInstance(mLauncher)
-                    .getWidgetCache());
-            widget.ensurePreview();
-            widget.setVisibility(View.VISIBLE);
-            if (i < widgets.size() - 1) {
-                addDivider(widgetCells);
-            }
-        }
-
-        if (widgets.size() == 1) {
-            // If there is only one widget, we want to center it instead of left-align.
-            WidgetsBottomSheet.LayoutParams params = (WidgetsBottomSheet.LayoutParams)
-                    widgetRow.getLayoutParams();
-            params.gravity = Gravity.CENTER_HORIZONTAL;
-        } else {
-            // Otherwise, add an empty view to the start as padding (but still scroll edge to edge).
-            View leftPaddingView = LayoutInflater.from(getContext()).inflate(
-                    R.layout.widget_list_divider, widgetRow, false);
-            leftPaddingView.getLayoutParams().width = ResourceUtils.pxFromDp(
-                    16, getResources().getDisplayMetrics());
-            widgetCells.addView(leftPaddingView, 0);
-        }
-    }
-
-    private void addDivider(ViewGroup parent) {
-        LayoutInflater.from(getContext()).inflate(R.layout.widget_list_divider, parent, true);
+        WidgetsTableUtils.groupWidgetItemsIntoTable(widgets, mMaxHorizontalSpan).forEach(row -> {
+            TableRow tableRow = new TableRow(getContext());
+            row.forEach(widgetItem -> {
+                WidgetCell widget = addItemCell(tableRow);
+                widget.setPreviewSize(widgetItem.spanX, widgetItem.spanY);
+                widget.applyFromCellItem(widgetItem, LauncherAppState.getInstance(mLauncher)
+                        .getWidgetCache());
+                widget.ensurePreview();
+                widget.setVisibility(View.VISIBLE);
+            });
+            widgetsTable.addView(tableRow);
+        });
     }
 
     protected WidgetCell addItemCell(ViewGroup parent) {
diff --git a/src/com/android/launcher3/widget/picker/SearchAndRecommendationsScrollController.java b/src/com/android/launcher3/widget/picker/SearchAndRecommendationsScrollController.java
new file mode 100644
index 0000000..a5ed20a
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/SearchAndRecommendationsScrollController.java
@@ -0,0 +1,162 @@
+/*
+ * Copyright (C) 2021 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.widget.picker;
+
+import android.view.View;
+import android.widget.RelativeLayout;
+
+import androidx.annotation.Nullable;
+import androidx.recyclerview.widget.RecyclerView;
+
+import com.android.launcher3.views.RecyclerViewFastScroller;
+import com.android.launcher3.widget.picker.WidgetsFullSheet.SearchAndRecommendationViewHolder;
+import com.android.launcher3.workprofile.PersonalWorkPagedView;
+
+/**
+ * A controller which measures & updates {@link WidgetsFullSheet}'s views padding, margin and
+ * vertical displacement upon scrolling.
+ */
+final class SearchAndRecommendationsScrollController implements
+        RecyclerViewFastScroller.OnFastScrollChangeListener {
+    private final boolean mHasWorkProfile;
+    private final SearchAndRecommendationViewHolder mViewHolder;
+    private final RecyclerView mPrimaryRecyclerView;
+
+    // The following are only non null if mHasWorkProfile is true.
+    @Nullable private final RecyclerView mWorkRecyclerView;
+    @Nullable private final View mPrimaryWorkTabsView;
+    @Nullable private final PersonalWorkPagedView mPrimaryWorkViewPager;
+
+    private int mMaxCollapsibleHeight = 0;
+
+    SearchAndRecommendationsScrollController(
+            boolean hasWorkProfile,
+            SearchAndRecommendationViewHolder viewHolder,
+            RecyclerView primaryRecyclerView,
+            @Nullable RecyclerView workRecyclerView,
+            @Nullable View personalWorkTabsView,
+            @Nullable PersonalWorkPagedView primaryWorkViewPager) {
+        mHasWorkProfile = hasWorkProfile;
+        mViewHolder = viewHolder;
+        mPrimaryRecyclerView = primaryRecyclerView;
+        mWorkRecyclerView = workRecyclerView;
+        mPrimaryWorkTabsView = personalWorkTabsView;
+        mPrimaryWorkViewPager = primaryWorkViewPager;
+    }
+
+    /**
+     * Updates the margin and padding of {@link WidgetsFullSheet} to accumulate collapsible views.
+     */
+    public void updateMarginAndPadding() {
+        // The maximum vertical distance, in pixels, until the last collapsible element is not
+        // visible from the screen when the user scrolls down the recycler view.
+        mMaxCollapsibleHeight = mViewHolder.mContainer.getPaddingTop()
+                + mViewHolder.mCollapseHandle.getMeasuredHeight()
+                + mViewHolder.mHeaderTitle.getMeasuredHeight();
+
+        int topContainerHeight = mViewHolder.mContainer.getMeasuredHeight();
+        if (mHasWorkProfile) {
+            // In a work profile setup, the full widget sheet contains the following views:
+            //           -------               -|
+            //           Widgets               -|---> LinearLayout for search & recommendations
+            //          Search bar             -|
+            //      Personal | Work
+            //           View Pager
+            //
+            // Views after the search & recommendations are not bound by RelativelyLayout param.
+            // To position them on the expected location, padding & margin are added to these views
+
+            // Tabs should have a padding of the height of the search & recommendations container.
+            mPrimaryWorkTabsView.setPadding(
+                    mPrimaryWorkTabsView.getPaddingLeft(),
+                    topContainerHeight,
+                    mPrimaryWorkTabsView.getPaddingRight(),
+                    mPrimaryWorkTabsView.getPaddingBottom());
+
+            // Instead of setting the top offset directly, we split the top offset into two values:
+            // 1. topOffsetAfterAllViewsCollapsed: this is the top offset after all collapsible
+            //    views are no longer visible on the screen.
+            //    This value is set as the margin for the view pager.
+            // 2. mMaxCollapsibleDistance
+            //    This value is set as the padding for the recycler views in order to work with
+            //    clipToPadding="false", which is an attribute for not showing top / bottom padding
+            //    when a recycler view has not reached the top or bottom of the list.
+            //    e.g. a list of 10 entries, only 3 entries are visible at a time.
+            //         case 1: recycler view is scrolled to the top. Top padding is visible/
+            //         (top padding)
+            //         item 1
+            //         item 2
+            //         item 3
+            //
+            //         case 2: recycler view is scrolled to the middle. No padding is visible.
+            //         item 4
+            //         item 5
+            //         item 6
+            //
+            //         case 3: recycler view is scrolled to the end. bottom padding is visible.
+            //         item 8
+            //         item 9
+            //         item 10
+            //         (bottom padding): not set in this case.
+            //
+            // When the views are first inflated, the sum of topOffsetAfterAllViewsCollapsed and
+            // mMaxCollapsibleDistance should equal to the top container height.
+            int tabsViewActualHeight =
+                    mPrimaryWorkTabsView.getMeasuredHeight() - mPrimaryWorkTabsView.getPaddingTop();
+            int topOffsetAfterAllViewsCollapsed =
+                    topContainerHeight + tabsViewActualHeight - mMaxCollapsibleHeight;
+
+            RelativeLayout.LayoutParams layoutParams =
+                    (RelativeLayout.LayoutParams) mPrimaryWorkViewPager.getLayoutParams();
+            layoutParams.setMargins(0, topOffsetAfterAllViewsCollapsed, 0, 0);
+            mPrimaryWorkViewPager.setLayoutParams(layoutParams);
+            mPrimaryWorkViewPager.requestLayout();
+
+            mPrimaryRecyclerView.setPadding(
+                    mPrimaryRecyclerView.getPaddingLeft(),
+                    mMaxCollapsibleHeight,
+                    mPrimaryRecyclerView.getPaddingRight(),
+                    mPrimaryRecyclerView.getPaddingBottom());
+            mWorkRecyclerView.setPadding(
+                    mWorkRecyclerView.getPaddingLeft(),
+                    mMaxCollapsibleHeight,
+                    mWorkRecyclerView.getPaddingRight(),
+                    mWorkRecyclerView.getPaddingBottom());
+        } else {
+            mPrimaryRecyclerView.setPadding(
+                    mPrimaryRecyclerView.getPaddingLeft(),
+                    topContainerHeight,
+                    mPrimaryRecyclerView.getPaddingRight(),
+                    mPrimaryRecyclerView.getPaddingBottom());
+        }
+    }
+
+    /**
+     * Changes the displacement of collapsible views (e.g. title & widget recommendations) and fixed
+     * views (e.g. recycler views, tabs) upon scrolling.
+     */
+    @Override
+    public void onThumbOffsetYChanged(int y) {
+        if (mMaxCollapsibleHeight > 0) {
+            int yDisplacement = Math.max(-y, -mMaxCollapsibleHeight);
+            mViewHolder.mHeaderTitle.setTranslationY(yDisplacement);
+            mViewHolder.mSearchBar.setTranslationY(yDisplacement);
+            if (mHasWorkProfile) {
+                mPrimaryWorkTabsView.setTranslationY(yDisplacement);
+            }
+        }
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 81cd73a..3a186c5 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -34,11 +34,14 @@
 import android.view.View;
 import android.view.animation.AnimationUtils;
 import android.view.animation.Interpolator;
+import android.widget.EditText;
+import android.widget.TextView;
 
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import androidx.recyclerview.widget.RecyclerView;
 
+import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Insettable;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppState;
@@ -60,7 +63,8 @@
  * Popup for showing the full list of available widgets
  */
 public class WidgetsFullSheet extends BaseWidgetSheet
-        implements Insettable, ProviderChangedListener, OnActivePageChangedListener {
+        implements Insettable, ProviderChangedListener, OnActivePageChangedListener,
+        WidgetsRecyclerView.HeaderViewDimensionsProvider {
 
     private static final long DEFAULT_OPEN_DURATION = 267;
     private static final long FADE_IN_DURATION = 150;
@@ -76,6 +80,10 @@
             mPrimaryWidgetsFilter.negate();
 
     @Nullable private PersonalWorkPagedView mViewPager;
+    private int mInitialTabsHeight = 0;
+    private View mTabsView;
+    private SearchAndRecommendationViewHolder mSearchAndRecommendationViewHolder;
+    private SearchAndRecommendationsScrollController mSearchAndRecommendationsScrollController;
 
     public WidgetsFullSheet(Context context, AttributeSet attrs, int defStyleAttr) {
         super(context, attrs, defStyleAttr);
@@ -97,8 +105,9 @@
         LayoutInflater layoutInflater = LayoutInflater.from(getContext());
         int contentLayoutRes = mHasWorkProfile ? R.layout.widgets_full_sheet_paged_view
                 : R.layout.widgets_full_sheet_recyclerview;
-        layoutInflater.inflate(contentLayoutRes,  springLayout, true);
+        layoutInflater.inflate(contentLayoutRes, springLayout, true);
 
+        RecyclerViewFastScroller fastScroller = findViewById(R.id.fast_scroller);
         if (mHasWorkProfile) {
             mViewPager = findViewById(R.id.widgets_view_pager);
             // Temporarily disable swipe gesture until widgets list horizontal scrollviews per
@@ -107,10 +116,12 @@
             mViewPager.initParentViews(this);
             mViewPager.getPageIndicator().setOnActivePageChangedListener(this);
             mViewPager.getPageIndicator().setActiveMarker(AdapterHolder.PRIMARY);
+            mTabsView = findViewById(R.id.tabs);
             findViewById(R.id.tab_personal)
                     .setOnClickListener((View view) -> mViewPager.snapToPage(0));
             findViewById(R.id.tab_work)
                     .setOnClickListener((View view) -> mViewPager.snapToPage(1));
+            fastScroller.setIsRecyclerViewFirstChildInParent(false);
             springLayout.addSpringView(R.id.primary_widgets_list_view);
             springLayout.addSpringView(R.id.work_widgets_list_view);
         } else {
@@ -118,12 +129,36 @@
             springLayout.addSpringView(R.id.primary_widgets_list_view);
         }
 
+        layoutInflater.inflate(R.layout.widgets_full_sheet_search_and_recommendations, springLayout,
+                true);
+        springLayout.addSpringView(R.id.search_and_recommendations_container);
+
+        mSearchAndRecommendationViewHolder = new SearchAndRecommendationViewHolder(
+                findViewById(R.id.search_and_recommendations_container));
+        mSearchAndRecommendationsScrollController = new SearchAndRecommendationsScrollController(
+                mHasWorkProfile,
+                mSearchAndRecommendationViewHolder,
+                findViewById(R.id.primary_widgets_list_view),
+                mHasWorkProfile ? findViewById(R.id.work_widgets_list_view) : null,
+                mTabsView,
+                mViewPager);
+        fastScroller.setOnFastScrollChangeListener(mSearchAndRecommendationsScrollController);
+
         onWidgetsBound();
     }
 
     @Override
     public void onActivePageChanged(int currentActivePage) {
         mAdapters.get(currentActivePage).mWidgetsRecyclerView.bindFastScrollbar();
+
+        reset();
+    }
+
+    private void reset() {
+        mAdapters.get(AdapterHolder.PRIMARY).mWidgetsRecyclerView.scrollToTop();
+        if (mHasWorkProfile) {
+            mAdapters.get(AdapterHolder.WORK).mWidgetsRecyclerView.scrollToTop();
+        }
     }
 
     @VisibleForTesting
@@ -181,20 +216,31 @@
 
     @Override
     protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
+        DeviceProfile deviceProfile = mLauncher.getDeviceProfile();
         int widthUsed;
         if (mInsets.bottom > 0) {
             widthUsed = mInsets.left + mInsets.right;
         } else {
-            Rect padding = mLauncher.getDeviceProfile().workspacePadding;
+            Rect padding = deviceProfile.workspacePadding;
             widthUsed = Math.max(padding.left + padding.right,
                     2 * (mInsets.left + mInsets.right));
         }
 
-        int heightUsed = mInsets.top + mLauncher.getDeviceProfile().edgeMarginPx;
+        int heightUsed = mInsets.top + deviceProfile.edgeMarginPx;
         measureChildWithMargins(mContent, widthMeasureSpec,
                 widthUsed, heightMeasureSpec, heightUsed);
         setMeasuredDimension(MeasureSpec.getSize(widthMeasureSpec),
                 MeasureSpec.getSize(heightMeasureSpec));
+
+        int paddingPx = 2 * getResources().getDimensionPixelOffset(
+                R.dimen.widget_cell_horizontal_padding);
+        int maxSpansPerRow = getMeasuredWidth() / (deviceProfile.cellWidthPx + paddingPx);
+        mAdapters.get(AdapterHolder.PRIMARY).mWidgetsListAdapter.setMaxHorizontalSpansPerRow(
+                maxSpansPerRow);
+        if (mHasWorkProfile) {
+            mAdapters.get(AdapterHolder.WORK).mWidgetsListAdapter.setMaxHorizontalSpansPerRow(
+                    maxSpansPerRow);
+        }
     }
 
     @Override
@@ -209,6 +255,12 @@
                 contentLeft + contentWidth, height);
 
         setTranslationShift(mTranslationShift);
+
+        if (mInitialTabsHeight == 0 && mTabsView != null) {
+            mInitialTabsHeight = mTabsView.getMeasuredHeight();
+        }
+
+        mSearchAndRecommendationsScrollController.updateMarginAndPadding();
     }
 
     @Override
@@ -314,6 +366,14 @@
         AccessibilityManagerCompat.sendStateEventToTest(getContext(), NORMAL_STATE_ORDINAL);
     }
 
+    @Override
+    public int getHeaderViewHeight() {
+        // No need to check work profile here because mInitialTabHeight is always 0 if there is no
+        // work profile.
+        return mInitialTabsHeight
+                + mSearchAndRecommendationViewHolder.mContainer.getMeasuredHeight();
+    }
+
     /** A holder class for holding adapters & their corresponding recycler view. */
     private final class AdapterHolder {
         static final int PRIMARY = 0;
@@ -343,9 +403,24 @@
         void setup(WidgetsRecyclerView recyclerView) {
             mWidgetsRecyclerView = recyclerView;
             mWidgetsRecyclerView.setAdapter(mWidgetsListAdapter);
+            mWidgetsRecyclerView.setHeaderViewDimensionsProvider(WidgetsFullSheet.this);
             mWidgetsRecyclerView.setEdgeEffectFactory(
                     ((TopRoundedCornerView) mContent).createEdgeEffectFactory());
             mWidgetsListAdapter.setApplyBitmapDeferred(false, mWidgetsRecyclerView);
         }
     }
+
+    final class SearchAndRecommendationViewHolder {
+        final View mContainer;
+        final View mCollapseHandle;
+        final EditText mSearchBar;
+        final TextView mHeaderTitle;
+
+        SearchAndRecommendationViewHolder(View searchAndRecommendationContainer) {
+            mContainer = searchAndRecommendationContainer;
+            mCollapseHandle = mContainer.findViewById(R.id.collapse_handle);
+            mSearchBar = mContainer.findViewById(R.id.widgets_search_bar);
+            mHeaderTitle = mContainer.findViewById(R.id.title);
+        }
+    }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
index 72b4a02..8b49d1e 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -19,10 +19,10 @@
 import android.util.Log;
 import android.util.SparseArray;
 import android.view.LayoutInflater;
-import android.view.View;
 import android.view.View.OnClickListener;
 import android.view.View.OnLongClickListener;
 import android.view.ViewGroup;
+import android.widget.TableRow;
 
 import androidx.annotation.Nullable;
 import androidx.recyclerview.widget.RecyclerView;
@@ -67,7 +67,7 @@
 
     private final WidgetsDiffReporter mDiffReporter;
     private final SparseArray<ViewHolderBinder> mViewHolderBinders = new SparseArray<>();
-    private final WidgetsListRowViewHolderBinder mWidgetsListRowViewHolderBinder;
+    private final WidgetsListTableViewHolderBinder mWidgetsListTableViewHolderBinder;
     private final WidgetListBaseRowEntryComparator mRowComparator =
             new WidgetListBaseRowEntryComparator();
 
@@ -84,9 +84,9 @@
             WidgetPreviewLoader widgetPreviewLoader, IconCache iconCache,
             OnClickListener iconClickListener, OnLongClickListener iconLongClickListener) {
         mDiffReporter = new WidgetsDiffReporter(iconCache, this);
-        mWidgetsListRowViewHolderBinder = new WidgetsListRowViewHolderBinder(context,
+        mWidgetsListTableViewHolderBinder = new WidgetsListTableViewHolderBinder(context,
                 layoutInflater, iconClickListener, iconLongClickListener, widgetPreviewLoader);
-        mViewHolderBinders.put(VIEW_TYPE_WIDGETS_LIST, mWidgetsListRowViewHolderBinder);
+        mViewHolderBinders.put(VIEW_TYPE_WIDGETS_LIST, mWidgetsListTableViewHolderBinder);
         mViewHolderBinders.put(VIEW_TYPE_WIDGETS_HEADER,
                 new WidgetsListHeaderViewHolderBinder(layoutInflater, this::onHeaderClicked));
     }
@@ -101,16 +101,16 @@
      * @see WidgetCell#setApplyBitmapDeferred(boolean)
      */
     public void setApplyBitmapDeferred(boolean isDeferred, RecyclerView rv) {
-        mWidgetsListRowViewHolderBinder.setApplyBitmapDeferred(isDeferred);
+        mWidgetsListTableViewHolderBinder.setApplyBitmapDeferred(isDeferred);
 
         for (int i = rv.getChildCount() - 1; i >= 0; i--) {
             ViewHolder viewHolder = rv.getChildViewHolder(rv.getChildAt(i));
             if (viewHolder.getItemViewType() == VIEW_TYPE_WIDGETS_LIST) {
                 WidgetsRowViewHolder holder = (WidgetsRowViewHolder) viewHolder;
-                for (int j = holder.cellContainer.getChildCount() - 1; j >= 0; j--) {
-                    View v = holder.cellContainer.getChildAt(j);
-                    if (v instanceof WidgetCell) {
-                        ((WidgetCell) v).setApplyBitmapDeferred(isDeferred);
+                for (int j = holder.mTableContainer.getChildCount() - 1; j >= 0; j--) {
+                    TableRow row =  (TableRow) holder.mTableContainer.getChildAt(j);
+                    for (int k = row.getChildCount() - 1; k >= 0; k--) {
+                        ((WidgetCell) row.getChildAt(k)).setApplyBitmapDeferred(isDeferred);
                     }
                 }
             }
@@ -204,6 +204,23 @@
         }
     }
 
+    /**
+     * Sets the max horizontal spans that are allowed for grouping more than one widgets in a table
+     * row.
+     *
+     * <p>If there is only one widget in a row, that widget horizontal span is allowed to exceed
+     * {@code maxHorizontalSpans}.
+     * <p>Let's say the max horizontal spans is set to 5. Widgets can be grouped in the same row if
+     * their total horizontal spans added don't exceed 5.
+     * Example 1: Row 1: 2x2, 2x3, 1x1. Total horizontal spans is 5. This is okay.
+     * Example 2: Row 1: 2x2, 4x3, 1x1. the total horizontal spans is 7. This is wrong.
+     *            4x3 and 1x1 should be moved to a new row.
+     * Example 3: Row 1: 6x4. This is okay because this is the only item in the row.
+     */
+    public void setMaxHorizontalSpansPerRow(int maxHorizontalSpans) {
+        mWidgetsListTableViewHolderBinder.setMaxSpansPerRow(maxHorizontalSpans);
+    }
+
     /** Comparator for sorting WidgetListRowEntry based on package title. */
     public static class WidgetListBaseRowEntryComparator implements
             Comparator<WidgetsListBaseEntry> {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java
deleted file mode 100644
index bd78777..0000000
--- a/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java
+++ /dev/null
@@ -1,149 +0,0 @@
-/*
- * Copyright (C) 2021 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.widget.picker;
-
-import android.content.Context;
-import android.util.Log;
-import android.view.LayoutInflater;
-import android.view.View;
-import android.view.View.OnClickListener;
-import android.view.View.OnLongClickListener;
-import android.view.ViewGroup;
-
-import com.android.launcher3.R;
-import com.android.launcher3.WidgetPreviewLoader;
-import com.android.launcher3.dragndrop.LivePreviewWidgetCell;
-import com.android.launcher3.model.WidgetItem;
-import com.android.launcher3.recyclerview.ViewHolderBinder;
-import com.android.launcher3.widget.WidgetCell;
-import com.android.launcher3.widget.model.WidgetsListContentEntry;
-
-import java.util.List;
-
-/**
- * Binds data from {@link WidgetsListContentEntry} to UI elements in {@link WidgetsRowViewHolder}.
- */
-public class WidgetsListRowViewHolderBinder
-        implements ViewHolderBinder<WidgetsListContentEntry, WidgetsRowViewHolder> {
-    private static final boolean DEBUG = false;
-    private static final String TAG = "WidgetsListRowViewHolderBinder";
-
-    private final LayoutInflater mLayoutInflater;
-    private final int mIndent;
-    private final OnClickListener mIconClickListener;
-    private final OnLongClickListener mIconLongClickListener;
-    private final WidgetPreviewLoader mWidgetPreviewLoader;
-    private boolean mApplyBitmapDeferred = false;
-
-    public WidgetsListRowViewHolderBinder(
-            Context context,
-            LayoutInflater layoutInflater,
-            OnClickListener iconClickListener,
-            OnLongClickListener iconLongClickListener,
-            WidgetPreviewLoader widgetPreviewLoader) {
-        mLayoutInflater = layoutInflater;
-        mIndent = context.getResources().getDimensionPixelSize(R.dimen.widget_section_indent);
-        mIconClickListener = iconClickListener;
-        mIconLongClickListener = iconLongClickListener;
-        mWidgetPreviewLoader = widgetPreviewLoader;
-    }
-
-    /**
-     * Defers applying bitmap on all the {@link WidgetCell} at
-     * {@link #bindViewHolder(WidgetsRowViewHolder, WidgetsListContentEntry)} if
-     * {@code applyBitmapDeferred} is {@code true}.
-     */
-    public void setApplyBitmapDeferred(boolean applyBitmapDeferred) {
-        mApplyBitmapDeferred = applyBitmapDeferred;
-    }
-
-    @Override
-    public WidgetsRowViewHolder newViewHolder(ViewGroup parent) {
-        if (DEBUG) {
-            Log.v(TAG, "\nonCreateViewHolder");
-        }
-
-        ViewGroup container = (ViewGroup) mLayoutInflater.inflate(
-                R.layout.widgets_scroll_container, parent, false);
-
-        // if the end padding is 0, then container view (horizontal scroll view) doesn't respect
-        // the end of the linear layout width + the start padding and doesn't allow scrolling.
-        container.findViewById(R.id.widgets_cell_list).setPaddingRelative(mIndent, 0, 1, 0);
-
-        return new WidgetsRowViewHolder(container);
-    }
-
-    @Override
-    public void bindViewHolder(WidgetsRowViewHolder holder, WidgetsListContentEntry entry) {
-        List<WidgetItem> infoList = entry.mWidgets;
-
-        ViewGroup row = holder.cellContainer;
-        if (DEBUG) {
-            Log.d(TAG, String.format("onBindViewHolder [widget#=%d, row.getChildCount=%d]",
-                    infoList.size(), row.getChildCount()));
-        }
-
-        // Add more views.
-        // if there are too many, hide them.
-        int expectedChildCount = infoList.size() + Math.max(0, infoList.size() - 1);
-        int childCount = row.getChildCount();
-
-        if (expectedChildCount > childCount) {
-            for (int i = childCount; i < expectedChildCount; i++) {
-                if ((i & 1) == 1) {
-                    // Add a divider for odd index
-                    mLayoutInflater.inflate(R.layout.widget_list_divider, row);
-                } else {
-                    // Add cell for even index
-                    LivePreviewWidgetCell widget = (LivePreviewWidgetCell) mLayoutInflater.inflate(
-                            R.layout.live_preview_widget_cell, row, false);
-
-                    // set up touch.
-                    widget.setOnClickListener(mIconClickListener);
-                    widget.setOnLongClickListener(mIconLongClickListener);
-                    row.addView(widget);
-                }
-            }
-        } else if (expectedChildCount < childCount) {
-            for (int i = expectedChildCount; i < childCount; i++) {
-                row.getChildAt(i).setVisibility(View.GONE);
-            }
-        }
-
-        // Bind the view in the widget horizontal tray region.
-        for (int i = 0; i < infoList.size(); i++) {
-            LivePreviewWidgetCell widget = (LivePreviewWidgetCell) row.getChildAt(2 * i);
-            widget.reset();
-            widget.applyFromCellItem(infoList.get(i), mWidgetPreviewLoader);
-            widget.setApplyBitmapDeferred(mApplyBitmapDeferred);
-            widget.ensurePreview();
-            widget.setVisibility(View.VISIBLE);
-
-            if (i > 0) {
-                row.getChildAt(2 * i - 1).setVisibility(View.VISIBLE);
-            }
-        }
-    }
-
-    @Override
-    public void unbindViewHolder(WidgetsRowViewHolder holder) {
-        int total = holder.cellContainer.getChildCount();
-        for (int i = 0; i < total; i += 2) {
-            WidgetCell widget = (WidgetCell) holder.cellContainer.getChildAt(i);
-            widget.clear();
-        }
-    }
-}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
new file mode 100644
index 0000000..9c8684a
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
@@ -0,0 +1,178 @@
+/*
+ * Copyright (C) 2021 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.widget.picker;
+
+import android.content.Context;
+import android.util.Log;
+import android.view.LayoutInflater;
+import android.view.View;
+import android.view.View.OnClickListener;
+import android.view.View.OnLongClickListener;
+import android.view.ViewGroup;
+import android.widget.TableLayout;
+import android.widget.TableRow;
+
+import com.android.launcher3.R;
+import com.android.launcher3.WidgetPreviewLoader;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.recyclerview.ViewHolderBinder;
+import com.android.launcher3.widget.WidgetCell;
+import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.util.WidgetsTableUtils;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Binds data from {@link WidgetsListContentEntry} to UI elements in {@link WidgetsRowViewHolder}.
+ */
+public final class WidgetsListTableViewHolderBinder
+        implements ViewHolderBinder<WidgetsListContentEntry, WidgetsRowViewHolder> {
+    private static final boolean DEBUG = false;
+    private static final String TAG = "WidgetsListRowViewHolderBinder";
+
+    private int mMaxSpansPerRow = 4;
+    private final LayoutInflater mLayoutInflater;
+    private final int mIndent;
+    private final OnClickListener mIconClickListener;
+    private final OnLongClickListener mIconLongClickListener;
+    private final WidgetPreviewLoader mWidgetPreviewLoader;
+    private boolean mApplyBitmapDeferred = false;
+
+    public WidgetsListTableViewHolderBinder(
+            Context context,
+            LayoutInflater layoutInflater,
+            OnClickListener iconClickListener,
+            OnLongClickListener iconLongClickListener,
+            WidgetPreviewLoader widgetPreviewLoader) {
+        mLayoutInflater = layoutInflater;
+        mIndent = context.getResources().getDimensionPixelSize(R.dimen.widget_section_indent);
+        mIconClickListener = iconClickListener;
+        mIconLongClickListener = iconLongClickListener;
+        mWidgetPreviewLoader = widgetPreviewLoader;
+    }
+
+    /**
+     * Defers applying bitmap on all the {@link WidgetCell} at
+     * {@link #bindViewHolder(WidgetsRowViewHolder, WidgetsListContentEntry)} if
+     * {@code applyBitmapDeferred} is {@code true}.
+     */
+    public void setApplyBitmapDeferred(boolean applyBitmapDeferred) {
+        mApplyBitmapDeferred = applyBitmapDeferred;
+    }
+
+    public void setMaxSpansPerRow(int maxSpansPerRow) {
+        mMaxSpansPerRow = maxSpansPerRow;
+    }
+
+    @Override
+    public WidgetsRowViewHolder newViewHolder(ViewGroup parent) {
+        if (DEBUG) {
+            Log.v(TAG, "\nonCreateViewHolder");
+        }
+
+        ViewGroup container = (ViewGroup) mLayoutInflater.inflate(
+                R.layout.widgets_table_container, parent, false);
+
+        // if the end padding is 0, then container view (horizontal scroll view) doesn't respect
+        // the end of the linear layout width + the start padding and doesn't allow scrolling.
+        container.findViewById(R.id.widgets_table).setPaddingRelative(mIndent, 0, 1, 0);
+
+        return new WidgetsRowViewHolder(container);
+    }
+
+    @Override
+    public void bindViewHolder(WidgetsRowViewHolder holder, WidgetsListContentEntry entry) {
+        TableLayout table = holder.mTableContainer;
+        if (DEBUG) {
+            Log.d(TAG, String.format("onBindViewHolder [widget#=%d, table.getChildCount=%d]",
+                    entry.mWidgets.size(), table.getChildCount()));
+        }
+
+        List<ArrayList<WidgetItem>> widgetItemsTable =
+                WidgetsTableUtils.groupWidgetItemsIntoTable(entry.mWidgets, mMaxSpansPerRow);
+        recycleTableBeforeBinding(table, widgetItemsTable);
+        // Bind the widget items.
+        for (int i = 0; i < widgetItemsTable.size(); i++) {
+            List<WidgetItem> widgetItemsPerRow = widgetItemsTable.get(i);
+            for (int j = 0; j < widgetItemsPerRow.size(); j++) {
+                TableRow row = (TableRow) table.getChildAt(i);
+                row.setVisibility(View.VISIBLE);
+                WidgetCell widget = (WidgetCell) row.getChildAt(j);
+                widget.clear();
+                WidgetItem widgetItem = widgetItemsPerRow.get(j);
+                widget.setPreviewSize(widgetItem.spanX, widgetItem.spanY);
+                widget.applyFromCellItem(widgetItem, mWidgetPreviewLoader);
+                widget.setApplyBitmapDeferred(mApplyBitmapDeferred);
+                widget.ensurePreview();
+                widget.setVisibility(View.VISIBLE);
+            }
+        }
+    }
+
+    /**
+     * Adds and hides table rows and columns from {@code table} to ensure there is sufficient room
+     * to display {@code widgetItemsTable}.
+     *
+     * <p>Instead of recreating all UI elements in {@code table}, this function recycles all
+     * existing UI elements. Instead of deleting excessive elements, it hides them.
+     */
+    private void recycleTableBeforeBinding(TableLayout table,
+            List<ArrayList<WidgetItem>> widgetItemsTable) {
+        // Hide extra table rows.
+        for (int i = widgetItemsTable.size(); i < table.getChildCount(); i++) {
+            table.getChildAt(i).setVisibility(View.GONE);
+        }
+
+        for (int i = 0; i < widgetItemsTable.size(); i++) {
+            List<WidgetItem> widgetItems = widgetItemsTable.get(i);
+            TableRow tableRow;
+            if (i < table.getChildCount()) {
+                tableRow = (TableRow) table.getChildAt(i);
+            } else {
+                tableRow = new TableRow(table.getContext());
+                table.addView(tableRow);
+            }
+            if (tableRow.getChildCount() > widgetItems.size()) {
+                for (int j = widgetItems.size(); j < tableRow.getChildCount(); j++) {
+                    tableRow.getChildAt(j).setVisibility(View.GONE);
+                }
+            } else {
+                for (int j = tableRow.getChildCount(); j < widgetItems.size(); j++) {
+                    WidgetCell widget = (WidgetCell) mLayoutInflater.inflate(
+                            R.layout.widget_cell, tableRow, false);
+                    // set up touch.
+                    widget.setOnClickListener(mIconClickListener);
+                    widget.setOnLongClickListener(mIconLongClickListener);
+                    tableRow.addView(widget);
+                }
+            }
+        }
+    }
+
+    @Override
+    public void unbindViewHolder(WidgetsRowViewHolder holder) {
+        int numOfRows = holder.mTableContainer.getChildCount();
+        for (int i = 0; i < numOfRows; i++) {
+            TableRow tableRow = (TableRow) holder.mTableContainer.getChildAt(i);
+            int numOfCols = tableRow.getChildCount();
+            for (int j = 0; j < numOfCols; j++) {
+                WidgetCell widget = (WidgetCell) tableRow.getChildAt(j);
+                widget.clear();
+            }
+        }
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java b/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java
index 52e9496..d65a809 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java
@@ -40,6 +40,7 @@
 
     private final Point mFastScrollerOffset = new Point();
     private boolean mTouchDownOnScroller;
+    private HeaderViewDimensionsProvider mHeaderViewDimensionsProvider;
 
     public WidgetsRecyclerView(Context context) {
         this(context, null);
@@ -135,8 +136,8 @@
     @Override
     protected int getAvailableScrollHeight() {
         View child = getChildAt(0);
-        return child.getMeasuredHeight() * mAdapter.getItemCount() - getScrollbarTrackHeight()
-                - mScrollbarTop;
+        return child.getMeasuredHeight() * mAdapter.getItemCount() + getScrollBarTop()
+                + getPaddingBottom() - mScrollbar.getHeight();
     }
 
     private boolean isModelNotReady() {
@@ -145,7 +146,9 @@
 
     @Override
     public int getScrollBarTop() {
-        return mScrollbarTop;
+        return mHeaderViewDimensionsProvider == null
+                ? mScrollbarTop
+                : mHeaderViewDimensionsProvider.getHeaderViewHeight() + mScrollbarTop;
     }
 
     @Override
@@ -171,4 +174,21 @@
     @Override
     public void onRequestDisallowInterceptTouchEvent(boolean disallowIntercept) {
     }
+
+    public void setHeaderViewDimensionsProvider(
+            HeaderViewDimensionsProvider headerViewDimensionsProvider) {
+        mHeaderViewDimensionsProvider = headerViewDimensionsProvider;
+    }
+
+    /**
+     * Provides dimensions of the header view that is shown at the top of a
+     * {@link WidgetsRecyclerView}.
+     */
+    public interface HeaderViewDimensionsProvider {
+        /**
+         * Returns the height, in pixels, of the header view that is shown at the top of a
+         * {@link WidgetsRecyclerView}.
+         */
+        int getHeaderViewHeight();
+    }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java b/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
index ae94584..aef1103 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
@@ -16,6 +16,7 @@
 package com.android.launcher3.widget.picker;
 
 import android.view.ViewGroup;
+import android.widget.TableLayout;
 
 import androidx.recyclerview.widget.RecyclerView.ViewHolder;
 
@@ -24,11 +25,11 @@
 /** A {@link ViewHolder} for showing widgets of an app in the full widget picker. */
 public final class WidgetsRowViewHolder extends ViewHolder {
 
-    public final ViewGroup cellContainer;
+    public final TableLayout mTableContainer;
 
     public WidgetsRowViewHolder(ViewGroup v) {
         super(v);
 
-        cellContainer = v.findViewById(R.id.widgets_cell_list);
+        mTableContainer = v.findViewById(R.id.widgets_table);
     }
 }
diff --git a/src/com/android/launcher3/widget/util/WidgetsTableUtils.java b/src/com/android/launcher3/widget/util/WidgetsTableUtils.java
new file mode 100644
index 0000000..e73d661
--- /dev/null
+++ b/src/com/android/launcher3/widget/util/WidgetsTableUtils.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright (C) 2021 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.widget.util;
+
+import com.android.launcher3.model.WidgetItem;
+
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+import java.util.stream.Collectors;
+
+/** An utility class which groups {@link WidgetItem}s into a table. */
+public final class WidgetsTableUtils {
+
+    /**
+     * Groups widgets in the following order:
+     * 1. Widgets always go before shortcuts.
+     * 2. Widgets with smaller horizontal spans will be shown first.
+     * 3. If widgets have the same horizontal spans, then widgets with a smaller vertical spans will
+     *    go first.
+     * 4. If both widgets have the same horizontal and vertical spans, they will use the same order
+     *    from the given {@code widgetItems}.
+     */
+    private static final Comparator<WidgetItem> WIDGET_SHORTCUT_COMPARATOR = (item, otherItem) -> {
+        if (item.widgetInfo != null && otherItem.widgetInfo == null) return -1;
+
+        if (item.widgetInfo == null && otherItem.widgetInfo != null) return 1;
+        if (item.spanX == otherItem.spanX) {
+            if (item.spanY == otherItem.spanY) return 0;
+            return item.spanY > otherItem.spanY ? 1 : -1;
+        }
+        return item.spanX > otherItem.spanX ? 1 : -1;
+    };
+
+
+    /**
+     * Groups widgets items into a 2D array which matches their appearance in a UI table.
+     *
+     * <p>Grouping:
+     * 1. Widgets and shortcuts never group together in the same row.
+     * 2. The ordered widgets are grouped together in the same row until their total horizontal
+     *    spans exceed the {@code maxSpansPerRow}.
+     * 3. The order shortcuts are grouped together in the same row until their total horizontal
+     *    spans exceed the {@code maxSpansPerRow}.
+     */
+    public static List<ArrayList<WidgetItem>> groupWidgetItemsIntoTable(
+            List<WidgetItem> widgetItems, final int maxSpansPerRow) {
+        List<WidgetItem> sortedWidgetItems = widgetItems.stream().sorted(WIDGET_SHORTCUT_COMPARATOR)
+                .collect(Collectors.toList());
+        List<ArrayList<WidgetItem>> widgetItemsTable = new ArrayList<>();
+        ArrayList<WidgetItem> widgetItemsAtRow = null;
+        for (WidgetItem widgetItem : sortedWidgetItems) {
+            if (widgetItemsAtRow == null) {
+                widgetItemsAtRow = new ArrayList<>();
+                widgetItemsTable.add(widgetItemsAtRow);
+            }
+            int numOfWidgetItems = widgetItemsAtRow.size();
+            int totalHorizontalSpan = widgetItemsAtRow.stream().map(item -> item.spanX)
+                    .reduce(/* default= */ 0, Integer::sum);
+            if (numOfWidgetItems == 0) {
+                widgetItemsAtRow.add(widgetItem);
+            } else if (widgetItem.spanX + totalHorizontalSpan <= maxSpansPerRow
+                    && widgetItem.hasSameType(widgetItemsAtRow.get(numOfWidgetItems - 1))) {
+                // Group items in the same row if
+                // 1. they are with the same type, i.e. a row can only have widgets or shortcuts but
+                //    never a mix of both.
+                // 2. the total number of horizontal spans are smaller than or equal to
+                //    MAX_SPAN_PER_ROW. If an item has a horizontal span > MAX_SPAN_PER_ROW, we just
+                //    place it in its own row regardless of the horizontal span limit.
+                widgetItemsAtRow.add(widgetItem);
+            } else {
+                widgetItemsAtRow = new ArrayList<>();
+                widgetItemsTable.add(widgetItemsAtRow);
+                widgetItemsAtRow.add(widgetItem);
+            }
+        }
+        return widgetItemsTable;
+    }
+}
diff --git a/tests/Android.mk b/tests/Android.mk
index 43d51fc..2c7d30a 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -36,6 +36,9 @@
 endif
 
 LOCAL_MODULE := ub-launcher-aosp-tapl
+LOCAL_LICENSE_KINDS := SPDX-license-identifier-Apache-2.0
+LOCAL_LICENSE_CONDITIONS := notice
+LOCAL_NOTICE_FILE := $(LOCAL_PATH)/../NOTICE
 LOCAL_SDK_VERSION := system_current
 
 include $(BUILD_STATIC_JAVA_LIBRARY)
diff --git a/tests/AndroidManifest-common.xml b/tests/AndroidManifest-common.xml
index f243f27..744dee0 100644
--- a/tests/AndroidManifest-common.xml
+++ b/tests/AndroidManifest-common.xml
@@ -98,7 +98,7 @@
         <activity
             android:name="com.android.launcher3.testcomponent.TestLauncherActivity"
             android:clearTaskOnLaunch="true"
-            android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout|smallestScreenSize"
+            android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout"
             android:enabled="false"
             android:label="Test launcher"
             android:launchMode="singleTask"
diff --git a/tests/src/com/android/launcher3/ui/WorkTabTest.java b/tests/src/com/android/launcher3/ui/WorkTabTest.java
index aef26ae..5e41d43d 100644
--- a/tests/src/com/android/launcher3/ui/WorkTabTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkTabTest.java
@@ -154,9 +154,12 @@
                 Log.d(TestProtocol.WORK_PROFILE_REMOVED, "Work tab not setup. Skipping test");
                 return false;
             }
-            return ((AllAppsPagedView) l.getAppsView().getContentView()).getCurrentPage()
-                    == WORK_PAGE && ((TextView) workEduView.findViewById(
-                    R.id.content_text)).getText().equals(
+            if (((AllAppsPagedView) l.getAppsView().getContentView()).getCurrentPage()
+                    != WORK_PAGE) {
+                Log.d(TestProtocol.WORK_PROFILE_REMOVED, "Work page not highlighted");
+                return false;
+            }
+            return ((TextView) workEduView.findViewById(R.id.content_text)).getText().equals(
                     l.getResources().getString(R.string.work_profile_edu_work_apps));
         }, 60000);
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/Launchable.java b/tests/tapl/com/android/launcher3/tapl/Launchable.java
index 3fc83ff..c4a566b 100644
--- a/tests/tapl/com/android/launcher3/tapl/Launchable.java
+++ b/tests/tapl/com/android/launcher3/tapl/Launchable.java
@@ -67,7 +67,7 @@
                 () -> "Launching an app didn't open a new window: " + label);
 
         mLauncher.assertTrue(
-                "App didn't start: " + label,
+                "App didn't start: " + label + " (" + selector + ")",
                 TestHelpers.wait(Until.hasObject(selector), LauncherInstrumentation.WAIT_TIME_MS));
         return new Background(mLauncher);
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 6afadfa..f279a82 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -365,9 +365,11 @@
 
             if (hasSystemUiObject("keyguard_status_view")) return "Phone is locked";
 
-            if (!mDevice.hasObject(By.textStartsWith(""))) return "Screen is empty";
+            if (!mDevice.wait(Until.hasObject(By.textStartsWith("")), WAIT_TIME_MS)) {
+                return "Screen is empty";
+            }
 
-            final String navigationModeError = getNavigationModeMismatchError();
+            final String navigationModeError = getNavigationModeMismatchError(true);
             if (navigationModeError != null) return navigationModeError;
         } catch (Throwable e) {
             Log.w(TAG, "getSystemAnomalyMessage failed", e);
@@ -535,17 +537,28 @@
         mExpectedRotation = expectedRotation;
     }
 
-    public String getNavigationModeMismatchError() {
+    public String getNavigationModeMismatchError(boolean waitForCorrectState) {
+        final int waitTime = waitForCorrectState ? WAIT_TIME_MS : 0;
         final NavigationModel navigationModel = getNavigationModel();
-        final boolean hasRecentsButton = hasSystemUiObject("recent_apps");
-        final boolean hasHomeButton = hasSystemUiObject("home");
-        if ((navigationModel == NavigationModel.THREE_BUTTON) != hasRecentsButton) {
-            return "Presence of recents button doesn't match the interaction mode, mode="
-                    + navigationModel.name() + ", hasRecents=" + hasRecentsButton;
+
+        if (navigationModel == NavigationModel.THREE_BUTTON) {
+            if (!mDevice.wait(Until.hasObject(By.res(SYSTEMUI_PACKAGE, "recent_apps")), waitTime)) {
+                return "Recents button not present in 3-button mode";
+            }
+        } else {
+            if (!mDevice.wait(Until.gone(By.res(SYSTEMUI_PACKAGE, "recent_apps")), waitTime)) {
+                return "Recents button is present in non-3-button mode";
+            }
         }
-        if ((navigationModel != NavigationModel.ZERO_BUTTON) != hasHomeButton) {
-            return "Presence of home button doesn't match the interaction mode, mode="
-                    + navigationModel.name() + ", hasHome=" + hasHomeButton;
+
+        if (navigationModel == NavigationModel.ZERO_BUTTON) {
+            if (!mDevice.wait(Until.gone(By.res(SYSTEMUI_PACKAGE, "home")), waitTime)) {
+                return "Home button is present in gestural mode";
+            }
+        } else {
+            if (!mDevice.wait(Until.hasObject(By.res(SYSTEMUI_PACKAGE, "home")), waitTime)) {
+                return "Home button not present in non-gestural mode";
+            }
         }
         return null;
     }
@@ -556,7 +569,7 @@
         assertEquals("Unexpected display rotation",
                 mExpectedRotation, mDevice.getDisplayRotation());
 
-        final String error = getNavigationModeMismatchError();
+        final String error = getNavigationModeMismatchError(true);
         assertTrue(error, error == null);
 
         log("verifyContainerType: " + containerType);
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index f95abdb..22f4d31 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -27,7 +27,6 @@
 import androidx.test.uiautomator.UiObject2;
 import androidx.test.uiautomator.Until;
 
-import com.android.launcher3.tapl.LauncherInstrumentation.GestureScope;
 import com.android.launcher3.testing.TestProtocol;
 
 import java.util.Collection;
@@ -107,43 +106,25 @@
                     fullWidgetsPicker.wait(Until.scrollable(true), WAIT_TIME_MS));
             final Point displaySize = mLauncher.getRealDisplaySize();
 
-            final UiObject2 widgetsContainer = findTestAppWidgetsScrollContainer();
+            final UiObject2 widgetsContainer = findTestAppWidgetsTableContainer();
             mLauncher.assertTrue("Can't locate widgets list for the test app: "
-                                    + mLauncher.getLauncherPackageName(),
+                            + mLauncher.getLauncherPackageName(),
                     widgetsContainer != null);
             final BySelector labelSelector = By.clazz("android.widget.TextView").text(labelText);
             int i = 0;
             for (; ; ) {
-                final Collection<UiObject2> cells = widgetsContainer.getChildren();
-                mLauncher.assertTrue("Widgets doesn't have 2 rows: ", cells.size() >= 2);
-                for (UiObject2 cell : cells) {
-                    final UiObject2 label = cell.findObject(labelSelector);
-                    // The logic below doesn't handle the case which a widget cell of the given
-                    // label is not yet visible on the horizontal scrolling container. This won't be
-                    // an issue once we get rid of the horizontal scrolling container.
-                    if (label == null) continue;
-
-                    final UiObject2 widget = cell;
-                    mLauncher.assertEquals(
-                            "View is not WidgetCell",
-                            "com.android.launcher3.widget.WidgetCell",
-                            widget.getClassName());
-
-                    int maxWidth = 0;
-                    for (UiObject2 sibling : widget.getParent().getChildren()) {
-                        maxWidth = Math.max(mLauncher.getVisibleBounds(sibling).width(), maxWidth);
-                    }
-
-                    if (mLauncher.getVisibleBounds(widget).bottom
-                            <= displaySize.y - mLauncher.getBottomGestureSize()) {
-                        int visibleDelta = maxWidth - mLauncher.getVisibleBounds(widget).width();
-                        if (visibleDelta > 0) {
-                            Rect parentBounds = mLauncher.getVisibleBounds(cell.getParent());
-                            mLauncher.linearGesture(parentBounds.centerX() + visibleDelta
-                                            + mLauncher.getTouchSlop(),
-                                    parentBounds.centerY(), parentBounds.centerX(),
-                                    parentBounds.centerY(), 10, true, GestureScope.INSIDE);
+                final Collection<UiObject2> tableRows = widgetsContainer.getChildren();
+                for (UiObject2 row : tableRows) {
+                    final Collection<UiObject2> widgetCells = row.getChildren();
+                    for (UiObject2 widget : widgetCells) {
+                        final UiObject2 label = widget.findObject(labelSelector);
+                        if (label == null) {
+                            continue;
                         }
+                        mLauncher.assertEquals(
+                                "View is not WidgetCell",
+                                "com.android.launcher3.widget.WidgetCell",
+                                widget.getClassName());
 
                         return new Widget(mLauncher, widget);
                     }
@@ -151,7 +132,7 @@
 
                 mLauncher.assertTrue("Too many attempts", ++i <= 40);
                 final int scroll = getWidgetsScroll();
-                mLauncher.scrollToLastVisibleRow(widgetsContainer, cells, 0);
+                mLauncher.scrollToLastVisibleRow(fullWidgetsPicker, tableRows, 0);
                 final int newScroll = getWidgetsScroll();
                 mLauncher.assertTrue(
                         "Scrolled in a wrong direction in Widgets: from " + scroll + " to "
@@ -162,13 +143,13 @@
     }
 
     /** Finds the widgets list of this test app from the collapsed full widgets picker. */
-    private UiObject2 findTestAppWidgetsScrollContainer() {
+    private UiObject2 findTestAppWidgetsTableContainer() {
         final BySelector headerSelector = By.res(mLauncher.getLauncherPackageName(),
                 "widgets_list_header");
         final BySelector targetAppSelector = By.clazz("android.widget.TextView").text(
                 mLauncher.getContext().getPackageName());
         final BySelector widgetsContainerSelector = By.res(mLauncher.getLauncherPackageName(),
-                "widgets_cell_list");
+                "widgets_table");
 
         boolean hasHeaderExpanded = false;
         for (int i = 0; i < 40; i++) {
@@ -196,14 +177,12 @@
                 // Look for a widgets list.
                 UiObject2 widgetsContainer = fullWidgetsPicker.findObject(widgetsContainerSelector);
                 if (widgetsContainer != null) {
-                    // Make sure the widgets list is fully visible on the screen.
-                    mLauncher.scrollToLastVisibleRow(fullWidgetsPicker,
-                            widgetsContainer.getChildren(), 0);
                     return widgetsContainer;
                 }
                 mLauncher.scrollToLastVisibleRow(fullWidgetsPicker, List.of(headerTitle), 0);
             } else {
-                mLauncher.scrollToLastVisibleRow(fullWidgetsPicker, header.getChildren(), 0);
+                mLauncher.scrollToLastVisibleRow(fullWidgetsPicker, fullWidgetsPicker.getChildren(),
+                        0);
             }
         }
 
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index f0e686f..d43e235 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -110,7 +110,8 @@
                     TestProtocol.REQUEST_HOME_TO_ALL_APPS_SWIPE_HEIGHT).
                     getInt(TestProtocol.TEST_INFO_RESPONSE_FIELD);
             LauncherInstrumentation.log(
-                    "switchToAllApps: swipeHeight = " + swipeHeight + ", slop = "
+                    "switchToAllApps: deviceHeight = " + deviceHeight + ", startY = " + startY
+                            + ", swipeHeight = " + swipeHeight + ", slop = "
                             + mLauncher.getTouchSlop());
 
             mLauncher.swipeToState(