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/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index deb1388..248fa46 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();
@@ -1748,7 +1748,7 @@
         if (alpha > 0) {
             setVisibility(VISIBLE);
         } else if (!mFreezeViewVisibility) {
-            setVisibility(GONE);
+            setVisibility(INVISIBLE);
         }
     }
 
@@ -1760,7 +1760,7 @@
         if (mFreezeViewVisibility != freezeViewVisibility) {
             mFreezeViewVisibility = freezeViewVisibility;
             if (!mFreezeViewVisibility) {
-                setVisibility(mContentAlpha > 0 ? VISIBLE : GONE);
+                setVisibility(mContentAlpha > 0 ? VISIBLE : INVISIBLE);
             }
         }
     }
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..ffa239a
--- /dev/null
+++ b/res/layout/widgets_table_container.xml
@@ -0,0 +1,21 @@
+<?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:background="?android:attr/colorPrimaryDark" />
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/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/allapps/AllAppsContainerView.java b/src/com/android/launcher3/allapps/AllAppsContainerView.java
index 746bfba..ec2a6d5 100644
--- a/src/com/android/launcher3/allapps/AllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsContainerView.java
@@ -757,6 +757,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/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/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index ef7bdf4..1d0a466 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -158,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.");
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index bf420d9..7745e90 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -392,6 +392,9 @@
 
         @UiEvent(doc = "User selected from a selection row of Slice.")
         LAUNCHER_SLICE_SELECTION_ACTION(707),
+
+        @UiEvent(doc = "Launcher entered into AllApps state with device search enabled.")
+        LAUNCHER_ALLAPPS_ENTRY_WITH_DEVICE_SEARCH(720),
         ;
 
         // ADD MORE
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/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 81cd73a..39953b1 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -39,6 +39,7 @@
 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;
@@ -181,20 +182,30 @@
 
     @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 maxSpansPerRow = getMeasuredWidth() / (deviceProfile.cellWidthPx
+                + deviceProfile.workspaceCellPaddingXPx);
+        mAdapters.get(AdapterHolder.PRIMARY).mWidgetsListAdapter.setMaxHorizontalSpansPerRow(
+                maxSpansPerRow);
+        if (mHasWorkProfile) {
+            mAdapters.get(AdapterHolder.WORK).mWidgetsListAdapter.setMaxHorizontalSpansPerRow(
+                    maxSpansPerRow);
+        }
     }
 
     @Override
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..2355700
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
@@ -0,0 +1,176 @@
+/*
+ * 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);
+                WidgetItem widgetItem = widgetItemsPerRow.get(j);
+                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/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/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);
             }
         }