diff --git a/res/drawable/ic_expand_less.xml b/res/drawable/ic_expand_less.xml
new file mode 100644
index 0000000..8360cee
--- /dev/null
+++ b/res/drawable/ic_expand_less.xml
@@ -0,0 +1,25 @@
+<?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.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="24"
+    android:viewportHeight="24"
+    android:tint="?android:attr/textColorHint">
+    <path
+        android:fillColor="#FF000000"
+        android:pathData="M18.59,16.41L20,15l-8,-8 -8,8 1.41,1.41L12,9.83"/>
+</vector>
diff --git a/res/drawable/ic_expand_more.xml b/res/drawable/ic_expand_more.xml
new file mode 100644
index 0000000..49e24f6
--- /dev/null
+++ b/res/drawable/ic_expand_more.xml
@@ -0,0 +1,25 @@
+<?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.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="24"
+    android:viewportHeight="24"
+    android:tint="?android:attr/textColorHint">
+    <path
+        android:fillColor="#FF000000"
+        android:pathData="M5.41,7.59L4,9l8,8 8,-8 -1.41,-1.41L12,14.17"/>
+</vector>
diff --git a/res/drawable/widgets_tray_expand_button.xml b/res/drawable/widgets_tray_expand_button.xml
new file mode 100644
index 0000000..8316e0f
--- /dev/null
+++ b/res/drawable/widgets_tray_expand_button.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.
+-->
+<selector xmlns:android="http://schemas.android.com/apk/res/android">
+    <item android:state_checked="true"
+        android:drawable="@drawable/ic_expand_less" />
+    <item android:state_checked="false"
+        android:drawable="@drawable/ic_expand_more" />
+</selector>
diff --git a/res/layout/widgets_list_row_header.xml b/res/layout/widgets_list_row_header.xml
new file mode 100644
index 0000000..faff10c
--- /dev/null
+++ b/res/layout/widgets_list_row_header.xml
@@ -0,0 +1,71 @@
+<?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.
+-->
+<com.android.launcher3.widget.picker.WidgetsListHeader xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:tools="http://schemas.android.com/tools"
+    android:id="@+id/widgets_list_header"
+    android:layout_width="match_parent"
+    android:layout_height="wrap_content"
+    android:background="?android:attr/selectableItemBackground"
+    android:paddingVertical="20dp"
+    android:orientation="horizontal">
+
+    <ImageView
+        android:id="@+id/app_icon"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:layout_marginHorizontal="16dp"
+        android:importantForAccessibility="no"
+        tools:src="@drawable/ic_corp"/>
+
+    <LinearLayout
+        android:layout_width="0dp"
+        android:layout_height="wrap_content"
+        android:layout_gravity="center_vertical"
+        android:layout_weight="1"
+        android:orientation="vertical"
+        android:focusable="true"
+        android:descendantFocusability="afterDescendants">
+
+        <TextView
+            android:id="@+id/app_title"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:layout_gravity="center_vertical"
+            android:textColor="?android:attr/textColorPrimary"
+            android:textSize="16sp"
+            tools:text="App name" />
+
+        <TextView
+            android:id="@+id/app_subtitle"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            tools:text="n widgets" />
+
+    </LinearLayout>
+
+    <!-- This checkbox is not clickable. The outermost LinearLayout is responsible to handle all
+         click event and update the checkbox state. -->
+    <CheckBox
+        android:id="@+id/toggle"
+        android:layout_marginHorizontal="16dp"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:layout_gravity="center_vertical"
+        android:layout_alignParentEnd="true"
+        android:clickable="false"
+        android:button="@drawable/widgets_tray_expand_button"/>
+
+</com.android.launcher3.widget.picker.WidgetsListHeader>
\ No newline at end of file
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index e593fb4..b19ea22 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -185,4 +185,8 @@
         <attr name="android:name" />
         <attr name="android:id" />
     </declare-styleable>
+
+    <declare-styleable name="WidgetsListRowHeader">
+        <attr name="appIconSize" format="dimension" />
+    </declare-styleable>
 </resources>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 447c9ac..c30019b 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -54,6 +54,11 @@
     <string name="add_item_request_drag_hint">Touch &amp; hold to place manually</string>
     <!-- Button label to automatically add icon on home screen [CHAR_LIMIT=50] -->
     <string name="place_automatically">Add automatically</string>
+    <!-- Label for showing the number of widgets an app has in the full widgets picker. [CHAR_LIMIT=25] -->
+    <plurals name="widgets_tray_subtitle">
+        <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>
 
     <!-- All Apps -->
     <!-- Search bar text in the apps view. [CHAR_LIMIT=50] -->
diff --git a/robolectric_tests/config/robolectric.properties b/robolectric_tests/config/robolectric.properties
index 0ac997f..4e811f3 100644
--- a/robolectric_tests/config/robolectric.properties
+++ b/robolectric_tests/config/robolectric.properties
@@ -1,4 +1,4 @@
-sdk=30
+sdk=29
 shadows= \
     com.android.launcher3.shadows.LShadowAppPredictionManager \
     com.android.launcher3.shadows.LShadowAppWidgetManager \
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
new file mode 100644
index 0000000..04797a6
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
@@ -0,0 +1,269 @@
+/*
+ * 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 static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyZeroInteractions;
+import static org.robolectric.Shadows.shadowOf;
+
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+import android.graphics.Bitmap;
+import android.os.UserHandle;
+
+import androidx.recyclerview.widget.RecyclerView;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.icons.BitmapInfo;
+import com.android.launcher3.icons.ComponentWithLabel;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.picker.WidgetsListAdapter.WidgetListBaseRowEntryComparator;
+
+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 WidgetsDiffReporterTest {
+    private static final String TEST_PACKAGE_PREFIX = "com.google.test";
+    private static final WidgetListBaseRowEntryComparator COMPARATOR =
+            new WidgetListBaseRowEntryComparator();
+
+    @Mock private IconCache mIconCache;
+    @Mock private RecyclerView.Adapter mAdapter;
+
+    private InvariantDeviceProfile mTestProfile;
+    private WidgetsDiffReporter mWidgetsDiffReporter;
+    private Context mContext;
+    private WidgetsListHeaderEntry mHeaderA;
+    private WidgetsListHeaderEntry mHeaderB;
+    private WidgetsListHeaderEntry mHeaderC;
+    private WidgetsListHeaderEntry mHeaderD;
+    private WidgetsListHeaderEntry mHeaderE;
+    private WidgetsListContentEntry mContentC;
+    private WidgetsListContentEntry mContentE;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile.numRows = 5;
+        mTestProfile.numColumns = 5;
+
+        doAnswer(invocation -> ((ComponentWithLabel) invocation.getArgument(0))
+                .getComponent().getPackageName())
+                .when(mIconCache).getTitleNoCache(any());
+
+        mContext = RuntimeEnvironment.application;
+        mWidgetsDiffReporter = new WidgetsDiffReporter(mIconCache, mAdapter);
+        mHeaderA = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "A",
+                /* appName= */ "A", /* numOfWidgets= */ 3);
+        mHeaderB = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "B",
+                /* appName= */ "B", /* numOfWidgets= */ 3);
+        mHeaderC = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "C",
+                /* appName= */ "C", /* numOfWidgets= */ 3);
+        mContentC = createWidgetsContentEntry(TEST_PACKAGE_PREFIX + "C",
+                /* appName= */ "C", /* numOfWidgets= */ 3);
+        mHeaderD = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "D",
+                /* appName= */ "D", /* numOfWidgets= */ 3);
+        mHeaderE = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "E",
+                /* appName= */ "E", /* numOfWidgets= */ 3);
+        mContentE = createWidgetsContentEntry(TEST_PACKAGE_PREFIX + "E",
+                /* appName= */ "E", /* numOfWidgets= */ 3);
+    }
+
+    @Test
+    public void listNotChanged_shouldNotInvokeAnyCallbacks() {
+        // GIVEN the current list has app headers [A, B, C].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mHeaderC));
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, currentList, COMPARATOR);
+
+        // THEN there is no adaptor callback.
+        verifyZeroInteractions(mAdapter);
+        // THEN the current list contains the same entries.
+        assertThat(currentList).containsExactly(mHeaderA, mHeaderB, mHeaderC);
+    }
+
+    @Test
+    public void headersOnly_emptyListToNonEmpty_shouldInvokeNotifyDataSetChanged() {
+        // GIVEN the current list has app headers [A, B, C].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>();
+
+        List<WidgetsListBaseEntry> newList = List.of(
+                createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "A", "A", 3),
+                createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "B", "B", 3),
+                createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "C", "C", 3));
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN notifyDataSetChanged is called
+        verify(mAdapter).notifyDataSetChanged();
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
+    @Test
+    public void headersOnly_nonEmptyToEmptyList_shouldInvokeNotifyDataSetChanged() {
+        // GIVEN the current list has app headers [A, B, C].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mHeaderC));
+        // GIVEN the new list is empty.
+        List<WidgetsListBaseEntry> newList = List.of();
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN notifyDataSetChanged is called.
+        verify(mAdapter).notifyDataSetChanged();
+        // THEN the current list isEmpty.
+        assertThat(currentList).isEmpty();
+    }
+
+    @Test
+    public void headersOnly_itemAddedAndRemovedInTheNewList_shouldInvokeCorrectCallbacks() {
+        // GIVEN the current list has app headers [A, B, D].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mHeaderD));
+        // GIVEN the new list has app headers [A, C, E].
+        List<WidgetsListBaseEntry> newList = List.of(mHeaderA, mHeaderC, mHeaderE);
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN "B" is removed from position 1.
+        verify(mAdapter).notifyItemRemoved(/* position= */ 1);
+        // THEN "D" is removed from position 2.
+        verify(mAdapter).notifyItemRemoved(/* position= */ 2);
+        // THEN "C" is inserted at position 1.
+        verify(mAdapter).notifyItemInserted(/* position= */ 1);
+        // THEN "E" is inserted at position 2.
+        verify(mAdapter).notifyItemInserted(/* position= */ 2);
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
+    @Test
+    public void headersContentsMix_itemAddedAndRemovedInTheNewList_shouldInvokeCorrectCallbacks() {
+        // GIVEN the current list has app headers [A, B, E content].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mContentE));
+        // GIVEN the new list has app headers [A, C content, D].
+        List<WidgetsListBaseEntry> newList = List.of(mHeaderA, mContentC, mHeaderD);
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN "B" is removed from position 1.
+        verify(mAdapter).notifyItemRemoved(/* position= */ 1);
+        // THEN "C content" is inserted at position 1.
+        verify(mAdapter).notifyItemInserted(/* position= */ 1);
+        // THEN "D" is inserted at position 2.
+        verify(mAdapter).notifyItemInserted(/* position= */ 2);
+        // THEN "E content" is removed from position 3.
+        verify(mAdapter).notifyItemRemoved(/* position= */ 3);
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
+    @Test
+    public void headersContentsMix_userInteractWithHeader_shouldInvokeCorrectCallbacks() {
+        // GIVEN the current list has app headers [A, B, E content].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mContentE));
+        // GIVEN the new list has app headers [A, B, E content].
+        List<WidgetsListBaseEntry> newList = List.of(mHeaderA, mHeaderB, mContentE);
+        // GIVEN the user has interacted with B.
+        mHeaderB.setIsWidgetListShown(true);
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN notify "B" has been changed.
+        verify(mAdapter).notifyItemChanged(/* position= */ 1);
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
+
+    private WidgetsListHeaderEntry createWidgetsHeaderEntry(String packageName, String appName,
+            int numOfWidgets) {
+        List<WidgetItem> widgetItems = generateWidgetItems(packageName, numOfWidgets);
+        PackageItemInfo pInfo = createPackageItemInfo(packageName, appName,
+                widgetItems.get(0).user);
+
+        return new WidgetsListHeaderEntry(pInfo, /* titleSectionName= */ "", widgetItems);
+    }
+
+    private WidgetsListContentEntry createWidgetsContentEntry(String packageName, String appName,
+            int numOfWidgets) {
+        List<WidgetItem> widgetItems = generateWidgetItems(packageName, numOfWidgets);
+        PackageItemInfo pInfo = createPackageItemInfo(packageName, appName,
+                widgetItems.get(0).user);
+
+        return new WidgetsListContentEntry(pInfo, /* titleSectionName= */ "", widgetItems);
+    }
+
+    private PackageItemInfo createPackageItemInfo(String packageName, String appName,
+            UserHandle userHandle) {
+        PackageItemInfo pInfo = new PackageItemInfo(packageName);
+        pInfo.title = appName;
+        pInfo.user = userHandle;
+        pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
+        return pInfo;
+    }
+
+    private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
+        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+        ArrayList<WidgetItem> widgetItems = new ArrayList<>();
+        for (int i = 0; i < numOfWidgets; i++) {
+            ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
+            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
+            widgetInfo.provider = cn;
+            ReflectionHelpers.setField(widgetInfo, "providerInfo",
+                    packageManager.addReceiverIfNotPresent(cn));
+
+            WidgetItem widgetItem = new WidgetItem(
+                    LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
+                    mTestProfile, mIconCache);
+            widgetItems.add(widgetItem);
+        }
+        return widgetItems;
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
index 9bea2fb..e94b253 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
@@ -40,11 +40,13 @@
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
@@ -56,9 +58,7 @@
 
 @RunWith(RobolectricTestRunner.class)
 public final class WidgetsListAdapterTest {
-
-    private static final String TEST_PACKAGE_1 = "com.google.test.1";
-    private static final String TEST_PACKAGE_2 = "com.google.test.2";
+    private static final String TEST_PACKAGE_PLACEHOLDER = "com.google.test";
 
     @Mock private LayoutInflater mMockLayoutInflater;
     @Mock private WidgetPreviewLoader mMockWidgetCache;
@@ -117,37 +117,76 @@
     }
 
     @Test
-    public void setWidgets_sameApp_moreWidgets_shouldNotifyItemChangedWithWidgetItemInfoDiff() {
-        // GIVEN the adapter was first populated with test package 1 & test package 2.
-        WidgetsListBaseEntry testPackage1With2WidgetsListEntry =
-                generateSampleAppWithWidgets(TEST_PACKAGE_1, /* numOfWidgets= */ 2);
-        WidgetsListBaseEntry testPackage2With2WidgetsListEntry =
-                generateSampleAppWithWidgets(TEST_PACKAGE_2, /* numOfWidgets= */ 2);
-        mAdapter.setWidgets(
-                List.of(testPackage1With2WidgetsListEntry, testPackage2With2WidgetsListEntry));
+    public void headerClick_expanded_shouldNotifyItemChange() {
+        // GIVEN a list of widgets entries:
+        // [com.google.test0, com.google.test0 content,
+        //  com.google.test1, com.google.test1 content,
+        //  com.google.test2, com.google.test2 content]
+        // The visible widgets entries: [com.google.test0, com.google.test1, com.google.test2].
+        mAdapter.setWidgets(generateSampleMap(3));
 
-        // WHEN the adapter is updated with the same list of apps but test package 2 has 3 widgets
+        // WHEN com.google.test.1 header is expanded.
+        mAdapter.onHeaderClicked(/* isExpanded= */ true, TEST_PACKAGE_PLACEHOLDER + 1);
+
+        // THEN the visible entries list becomes:
+        // [com.google.test0, com.google.test1, com.google.test1 content, com.google.test2]
+        // com.google.test.1 content is inserted into position 2.
+        verify(mListener).onItemRangeInserted(eq(2), eq(1));
+    }
+
+    @Test
+    public void setWidgets_expandedApp_moreWidgets_shouldNotifyItemChangedWithWidgetItemInfoDiff() {
+        // GIVEN the adapter was first populated with com.google.test0 & com.google.test1. Each app
+        // has one widget.
+        ArrayList<WidgetsListBaseEntry> allEntries = generateSampleMap(2);
+        mAdapter.setWidgets(allEntries);
+        // GIVEN test com.google.test1 is expanded.
+        // Visible entries in the adapter are:
+        // [com.google.test0, com.google.test1, com.google.test1 content]
+        mAdapter.onHeaderClicked(/* isExpanded= */ true, TEST_PACKAGE_PLACEHOLDER + 1);
+        Mockito.reset(mListener);
+
+        // WHEN the adapter is updated with the same list of apps but com.google.test1 has 2 widgets
         // now.
-        WidgetsListBaseEntry testPackage1With3WidgetsListEntry =
-                generateSampleAppWithWidgets(TEST_PACKAGE_2, /* numOfWidgets= */ 2);
-        mAdapter.setWidgets(
-                List.of(testPackage1With2WidgetsListEntry, testPackage1With3WidgetsListEntry));
+        WidgetsListContentEntry testPackage1ContentEntry =
+                (WidgetsListContentEntry) allEntries.get(3);
+        WidgetItem widgetItem = testPackage1ContentEntry.mWidgets.get(0);
+        WidgetsListContentEntry newTestPackage1ContentEntry = new WidgetsListContentEntry(
+                testPackage1ContentEntry.mPkgItem,
+                testPackage1ContentEntry.mTitleSectionName, List.of(widgetItem, widgetItem));
+        allEntries.set(3, newTestPackage1ContentEntry);
+        mAdapter.setWidgets(allEntries);
 
-        // THEN the onItemRangeChanged is invoked.
-        verify(mListener).onItemRangeChanged(eq(1), eq(1), isNull());
+        // THEN the onItemRangeChanged is invoked for "com.google.test1 content" at index 2.
+        verify(mListener).onItemRangeChanged(eq(2), eq(1), isNull());
     }
 
     @Test
     public void setWidgets_hodgepodge_shouldInvokeExpectedDataObserverCallbacks() {
+        // GIVEN a widgets entry list:
+        // Index:  0|   1      | 2|      3   | 4|     5    | 6|     7    | 8|     9    |
+        //        [A, A content, B, B content, C, C content, D, D content, E, E content]
         List<WidgetsListBaseEntry> allAppsWithWidgets = generateSampleMap(5);
-        // GIVEN the current widgets list consist of [A, B, E].
+        // GIVEN the current widgets list consist of [A, A content, B, B content, E, E content].
+        // GIVEN the visible widgets list consist of [A, B, E]
         List<WidgetsListBaseEntry> currentList = List.of(
-                allAppsWithWidgets.get(0), allAppsWithWidgets.get(1), allAppsWithWidgets.get(4));
+                // A & A content
+                allAppsWithWidgets.get(0), allAppsWithWidgets.get(1),
+                // B & B content
+                allAppsWithWidgets.get(2), allAppsWithWidgets.get(3),
+                // E & E content
+                allAppsWithWidgets.get(8), allAppsWithWidgets.get(9));
         mAdapter.setWidgets(currentList);
 
-        // WHEN the widgets list is updated to [A, C, D].
+        // WHEN the widgets list is updated to [A, A content, C, C content, D, D content].
+        // WHEN the visible widgets list is updated to [A, C, D].
         List<WidgetsListBaseEntry> newList = List.of(
-                allAppsWithWidgets.get(0), allAppsWithWidgets.get(2), allAppsWithWidgets.get(3));
+                // A & A content
+                allAppsWithWidgets.get(0), allAppsWithWidgets.get(1),
+                // C & C content
+                allAppsWithWidgets.get(4), allAppsWithWidgets.get(5),
+                // D & D content
+                allAppsWithWidgets.get(6), allAppsWithWidgets.get(7));
         mAdapter.setWidgets(newList);
 
         // Computation logic                           | [Intermediate list during computation]
@@ -162,15 +201,23 @@
     }
 
     /**
-     * Helper method to generate the sample widget model map that can be used for the tests
-     * @param num the number of WidgetItem the map should contain
+     * Generates a list of sample widget entries.
+     *
+     * <p>Each sample app has 1 widget only. An app is represented by 2 entries,
+     * {@link WidgetsListHeaderEntry} & {@link WidgetsListContentEntry}. Only
+     * {@link WidgetsListHeaderEntry} is always visible in the {@link WidgetsListAdapter}.
+     * {@link WidgetsListContentEntry} is only shown upon clicking the corresponding app's
+     * {@link WidgetsListHeaderEntry}. Only at most one {@link WidgetsListContentEntry} is shown at
+     * a time.
+     *
+     * @param num the number of apps that have widgets.
      */
     private ArrayList<WidgetsListBaseEntry> generateSampleMap(int num) {
         ArrayList<WidgetsListBaseEntry> result = new ArrayList<>();
         if (num <= 0) return result;
 
         for (int i = 0; i < num; i++) {
-            String packageName = "com.placeholder.apk" + i;
+            String packageName = TEST_PACKAGE_PLACEHOLDER + i;
 
             List<WidgetItem> widgetItems = generateWidgetItems(packageName, /* numOfWidgets= */ 1);
 
@@ -179,23 +226,13 @@
             pInfo.user = widgetItems.get(0).user;
             pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
 
+            result.add(new WidgetsListHeaderEntry(pInfo, /* titleSectionName= */ "", widgetItems));
             result.add(new WidgetsListContentEntry(pInfo, /* titleSectionName= */ "", widgetItems));
         }
 
         return result;
     }
 
-    private WidgetsListBaseEntry generateSampleAppWithWidgets(String packageName,
-            int numOfWidgets) {
-        PackageItemInfo appInfo = new PackageItemInfo(packageName);
-        appInfo.title = appInfo.packageName;
-        appInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
-
-        return new WidgetsListContentEntry(appInfo,
-                /* titleSectionName= */ "",
-                generateWidgetItems(packageName, numOfWidgets));
-    }
-
     private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
         ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
         ArrayList<WidgetItem> widgetItems = new ArrayList<>();
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
new file mode 100644
index 0000000..ae5b9a5
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
@@ -0,0 +1,173 @@
+/*
+ * 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 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.graphics.Bitmap;
+import android.view.LayoutInflater;
+import android.view.View;
+import android.widget.FrameLayout;
+import android.widget.TextView;
+
+import androidx.annotation.Nullable;
+
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.R;
+import com.android.launcher3.icons.BitmapInfo;
+import com.android.launcher3.icons.ComponentWithLabel;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.testing.TestActivity;
+import com.android.launcher3.widget.WidgetCell;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.picker.WidgetsListHeaderViewHolderBinder.OnHeaderClickListener;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.android.controller.ActivityController;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@RunWith(RobolectricTestRunner.class)
+public final class WidgetsListHeaderViewHolderBinderTest {
+    private static final String TEST_PACKAGE = "com.google.test";
+    private static final String APP_NAME = "Test app";
+
+    private Context mContext;
+    private WidgetsListHeaderViewHolderBinder mViewHolderBinder;
+    private InvariantDeviceProfile mTestProfile;
+    // Replace ActivityController with ActivityScenario, which is the recommended way for activity
+    // testing.
+    private ActivityController<TestActivity> mActivityController;
+    private TestActivity mTestActivity;
+    private FakeOnHeaderClickListener mFakeOnHeaderClickListener = new FakeOnHeaderClickListener();
+
+    @Mock
+    private IconCache mIconCache;
+    @Mock
+    private DeviceProfile mDeviceProfile;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mContext = RuntimeEnvironment.application;
+        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile.numRows = 5;
+        mTestProfile.numColumns = 5;
+
+        mActivityController = Robolectric.buildActivity(TestActivity.class);
+        mTestActivity = mActivityController.setup().get();
+        mTestActivity.setDeviceProfile(mDeviceProfile);
+
+        doAnswer(invocation -> {
+            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            return componentWithLabel.getComponent().getShortClassName();
+        }).when(mIconCache).getTitleNoCache(any());
+
+        mViewHolderBinder = new WidgetsListHeaderViewHolderBinder(
+                LayoutInflater.from(mTestActivity),
+                mFakeOnHeaderClickListener);
+    }
+
+    @After
+    public void tearDown() {
+        mActivityController.destroy();
+    }
+
+    @Test
+    public void bindViewHolder_appWith3Widgets_shouldShowTheCorrectAppNameAndSubtitle() {
+        WidgetsListHeaderHolder viewHolder = mViewHolderBinder.newViewHolder(
+                new FrameLayout(mTestActivity));
+        WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
+        WidgetsListHeaderEntry entry = generateSampleAppHeader(
+                APP_NAME,
+                TEST_PACKAGE,
+                /* numOfWidgets= */ 3);
+        mViewHolderBinder.bindViewHolder(viewHolder, entry);
+
+        TextView appTitle = widgetsListHeader.findViewById(R.id.app_title);
+        TextView appSubtitle = widgetsListHeader.findViewById(R.id.app_subtitle);
+        assertThat(appTitle.getText()).isEqualTo(APP_NAME);
+        assertThat(appSubtitle.getText()).isEqualTo("3 widgets");
+    }
+
+    private WidgetsListHeaderEntry generateSampleAppHeader(String appName, String packageName,
+            int numOfWidgets) {
+        PackageItemInfo appInfo = new PackageItemInfo(packageName);
+        appInfo.title = appName;
+        appInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
+
+        return new WidgetsListHeaderEntry(appInfo,
+                /* titleSectionName= */ "",
+                generateWidgetItems(packageName, numOfWidgets));
+    }
+
+    private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
+        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+        ArrayList<WidgetItem> widgetItems = new ArrayList<>();
+        for (int i = 0; i < numOfWidgets; i++) {
+            ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
+            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
+            widgetInfo.provider = cn;
+            ReflectionHelpers.setField(widgetInfo, "providerInfo",
+                    packageManager.addReceiverIfNotPresent(cn));
+
+            widgetItems.add(new WidgetItem(
+                    LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
+                    mTestProfile, mIconCache));
+        }
+        return widgetItems;
+    }
+
+    private void assertWidgetCellWithLabel(View view, String label) {
+        assertThat(view).isInstanceOf(WidgetCell.class);
+        TextView widgetLabel = (TextView) view.findViewById(R.id.widget_name);
+        assertThat(widgetLabel.getText()).isEqualTo(label);
+    }
+
+    private final class FakeOnHeaderClickListener implements OnHeaderClickListener {
+
+        boolean mShowWidgets = false;
+        @Nullable  String mHeaderClickedPackage = null;
+
+        @Override
+        public void onHeaderClicked(boolean showWidgets, String packageName) {
+            mShowWidgets = showWidgets;
+            mHeaderClickedPackage = packageName;
+        }
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java
index 4e9e227..ec9fde3 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java
@@ -119,19 +119,6 @@
     }
 
     @Test
-    public void bindViewHolder_appWith3Widgets_shouldMatchAppTitle() {
-        WidgetsRowViewHolder viewHolder = mViewHolderBinder.newViewHolder(
-                new FrameLayout(mTestActivity));
-        WidgetsListContentEntry entry = generateSampleAppWithWidgets(
-                APP_NAME,
-                TEST_PACKAGE,
-                /* numOfWidgets= */ 3);
-        mViewHolderBinder.bindViewHolder(viewHolder, entry);
-
-        assertThat(viewHolder.title.getText()).isEqualTo(APP_NAME);
-    }
-
-    @Test
     public void bindViewHolder_appWith3Widgets_shouldHave3Widgets() {
         WidgetsRowViewHolder viewHolder = mViewHolderBinder.newViewHolder(
                 new FrameLayout(mTestActivity));
diff --git a/src/com/android/launcher3/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index 21297c9..cea8cd6 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -24,7 +24,6 @@
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.animation.ObjectAnimator;
-import android.animation.ValueAnimator;
 import android.content.Context;
 import android.content.res.ColorStateList;
 import android.content.res.TypedArray;
@@ -34,8 +33,6 @@
 import android.graphics.Paint;
 import android.graphics.Path;
 import android.graphics.PointF;
-import android.graphics.PorterDuff.Mode;
-import android.graphics.PorterDuffColorFilter;
 import android.graphics.Rect;
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
@@ -52,7 +49,6 @@
 
 import androidx.annotation.Nullable;
 import androidx.annotation.UiThread;
-import androidx.core.graphics.ColorUtils;
 
 import com.android.launcher3.Launcher.OnResumeCallback;
 import com.android.launcher3.accessibility.LauncherAccessibilityDelegate;
@@ -798,7 +794,7 @@
         if (mIcon != null
                 && mIcon instanceof PlaceHolderIconDrawable
                 && iconUpdateAnimationEnabled()) {
-            animateIconUpdate((PlaceHolderIconDrawable) mIcon, icon);
+            ((PlaceHolderIconDrawable) mIcon).animateIconUpdate(icon);
         }
 
         mDisableRelayout = false;
@@ -950,28 +946,6 @@
         }
     }
 
-    private static void animateIconUpdate(PlaceHolderIconDrawable oldIcon, Drawable newIcon) {
-        int placeholderColor = oldIcon.mPaint.getColor();
-        int originalAlpha = Color.alpha(placeholderColor);
-
-        ValueAnimator iconUpdateAnimation = ValueAnimator.ofInt(originalAlpha, 0);
-        iconUpdateAnimation.setDuration(ICON_UPDATE_ANIMATION_DURATION);
-        iconUpdateAnimation.addUpdateListener(valueAnimator -> {
-            int newAlpha = (int) valueAnimator.getAnimatedValue();
-            int newColor = ColorUtils.setAlphaComponent(placeholderColor, newAlpha);
-
-            newIcon.setColorFilter(new PorterDuffColorFilter(newColor, Mode.SRC_ATOP));
-        });
-        iconUpdateAnimation.addListener(new AnimatorListenerAdapter() {
-            @Override
-            public void onAnimationEnd(Animator animation) {
-                newIcon.setColorFilter(null);
-            }
-        });
-        iconUpdateAnimation.start();
-    }
-
-
     @Override
     public void decorate(int color) {
         mHighlightColor = color;
diff --git a/src/com/android/launcher3/graphics/PlaceHolderIconDrawable.java b/src/com/android/launcher3/graphics/PlaceHolderIconDrawable.java
index d347e8f..b6d25c4 100644
--- a/src/com/android/launcher3/graphics/PlaceHolderIconDrawable.java
+++ b/src/com/android/launcher3/graphics/PlaceHolderIconDrawable.java
@@ -19,10 +19,19 @@
 
 import static com.android.launcher3.graphics.IconShape.getShapePath;
 
+import android.animation.Animator;
+import android.animation.AnimatorListenerAdapter;
+import android.animation.ValueAnimator;
 import android.content.Context;
 import android.graphics.Canvas;
+import android.graphics.Color;
 import android.graphics.Path;
+import android.graphics.PorterDuff;
+import android.graphics.PorterDuffColorFilter;
 import android.graphics.Rect;
+import android.graphics.drawable.Drawable;
+
+import androidx.core.graphics.ColorUtils;
 
 import com.android.launcher3.FastBitmapDrawable;
 import com.android.launcher3.R;
@@ -53,4 +62,27 @@
         canvas.drawPath(mProgressPath, mPaint);
         canvas.restoreToCount(saveCount);
     }
+
+    /** Updates this placeholder to {@code newIcon} with animation. */
+    public void animateIconUpdate(Drawable newIcon) {
+        int placeholderColor = mPaint.getColor();
+        int originalAlpha = Color.alpha(placeholderColor);
+
+        ValueAnimator iconUpdateAnimation = ValueAnimator.ofInt(originalAlpha, 0);
+        iconUpdateAnimation.setDuration(375);
+        iconUpdateAnimation.addUpdateListener(valueAnimator -> {
+            int newAlpha = (int) valueAnimator.getAnimatedValue();
+            int newColor = ColorUtils.setAlphaComponent(placeholderColor, newAlpha);
+
+            newIcon.setColorFilter(new PorterDuffColorFilter(newColor, PorterDuff.Mode.SRC_ATOP));
+        });
+        iconUpdateAnimation.addListener(new AnimatorListenerAdapter() {
+            @Override
+            public void onAnimationEnd(Animator animation) {
+                newIcon.setColorFilter(null);
+            }
+        });
+        iconUpdateAnimation.start();
+    }
+
 }
diff --git a/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java b/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java
index 10ea7db..09517e1 100644
--- a/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java
+++ b/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java
@@ -16,9 +16,15 @@
 
 package com.android.launcher3.widget.model;
 
+import static java.lang.annotation.RetentionPolicy.SOURCE;
+
+import androidx.annotation.IntDef;
+
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.PackageItemInfo;
 
+import java.lang.annotation.Retention;
+
 /** Holder class to store the package information of an entry shown in the widgets list. */
 public abstract class WidgetsListBaseEntry {
     public final PackageItemInfo mPkgItem;
@@ -33,4 +39,22 @@
         mPkgItem = pkgItem;
         mTitleSectionName = titleSectionName;
     }
+
+    /**
+     * Returns the ranking of this entry in the
+     * {@link com.android.launcher3.widget.picker.WidgetsListAdapter}.
+     *
+     * <p>Entries with smaller value should be shown first. See
+     * {@link com.android.launcher3.widget.picker.WidgetsDiffReporter} for more details.
+     */
+    @Rank
+    public abstract int getRank();
+
+    @Retention(SOURCE)
+    @IntDef({RANK_WIDGETS_LIST_HEADER, RANK_WIDGETS_LIST_CONTENT})
+    public @interface Rank {
+    }
+
+    public static final int RANK_WIDGETS_LIST_HEADER = 1;
+    public static final int RANK_WIDGETS_LIST_CONTENT = 2;
 }
diff --git a/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java b/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
index 407f194..b0cb8c7 100644
--- a/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
+++ b/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
@@ -41,4 +41,10 @@
     public String toString() {
         return mPkgItem.packageName + ":" + mWidgets.size();
     }
+
+    @Override
+    @Rank
+    public int getRank() {
+        return RANK_WIDGETS_LIST_CONTENT;
+    }
 }
diff --git a/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java b/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java
new file mode 100644
index 0000000..6899647
--- /dev/null
+++ b/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java
@@ -0,0 +1,64 @@
+/*
+ * 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.model;
+
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.data.PackageItemInfo;
+
+import java.util.Collection;
+
+/** An information holder for an app which has widgets or/and shortcuts. */
+public final class WidgetsListHeaderEntry extends WidgetsListBaseEntry {
+
+    public final int widgetsCount;
+    public final int shortcutsCount;
+
+    private boolean mIsWidgetListShown = false;
+    private boolean mHasEntryUpdated = false;
+
+    public WidgetsListHeaderEntry(PackageItemInfo pkgItem, String titleSectionName,
+            Collection<WidgetItem> items) {
+        super(pkgItem, titleSectionName);
+        widgetsCount = (int) items.stream().filter(item -> item.widgetInfo != null).count();
+        shortcutsCount = Math.max(0, items.size() - widgetsCount);
+    }
+
+    /** Sets if the widgets list associated with this header is shown. */
+    public void setIsWidgetListShown(boolean isWidgetListShown) {
+        if (mIsWidgetListShown != isWidgetListShown) {
+            this.mIsWidgetListShown = isWidgetListShown;
+            mHasEntryUpdated = true;
+        } else {
+            mHasEntryUpdated = false;
+        }
+    }
+
+    /** Returns {@code true} if the widgets list associated with this header is shown. */
+    public boolean isWidgetListShown() {
+        return mIsWidgetListShown;
+    }
+
+    /** Returns {@code true} if this entry has been updated due to user interactions. */
+    public boolean hasEntryUpdated() {
+        return mHasEntryUpdated;
+    }
+
+    @Override
+    @Rank
+    public int getRank() {
+        return RANK_WIDGETS_LIST_HEADER;
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java b/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
index 398d9ba..dbd1bdf 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
@@ -24,10 +24,12 @@
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 import com.android.launcher3.widget.picker.WidgetsListAdapter.WidgetListBaseRowEntryComparator;
 
 import java.util.ArrayList;
 import java.util.Iterator;
+import java.util.List;
 
 /**
  * Do diff on widget's tray list items and call the {@link RecyclerView.Adapter}
@@ -50,7 +52,7 @@
      * relevant {@link androidx.recyclerview.widget.RecyclerView.RecyclerViewDataObserver} methods.
      */
     public void process(ArrayList<WidgetsListBaseEntry> currentEntries,
-            ArrayList<WidgetsListBaseEntry> newEntries,
+            List<WidgetsListBaseEntry> newEntries,
             WidgetListBaseRowEntryComparator comparator) {
         if (DEBUG) {
             Log.d(TAG, "process oldEntries#=" + currentEntries.size()
@@ -78,7 +80,7 @@
         WidgetsListBaseEntry newRowEntry = newIter.next();
 
         do {
-            int diff = comparePackageName(orgRowEntry, newRowEntry, comparator);
+            int diff = compareAppNameAndType(orgRowEntry, newRowEntry, comparator);
             if (DEBUG) {
                 Log.d(TAG, String.format("diff=%d orgRowEntry (%s) newRowEntry (%s)",
                         diff, orgRowEntry != null ? orgRowEntry.toString() : null,
@@ -106,11 +108,13 @@
                 mListener.notifyItemInserted(index);
 
             } else {
-                // same package name but,
+                // same app name & type but,
                 // did the icon, title, etc, change?
+                // or did the header view changed due to user interactions?
                 // or did the widget size and desc, span, etc change?
                 if (!isSamePackageItemInfo(orgRowEntry.mPkgItem, newRowEntry.mPkgItem)
-                        || !areWidgetsEqual(orgRowEntry, newRowEntry)) {
+                        || hasHeaderUpdated(newRowEntry)
+                        || hasWidgetsListChanged(orgRowEntry, newRowEntry)) {
                     index = currentEntries.indexOf(orgRowEntry);
                     currentEntries.set(index, newRowEntry);
                     mListener.notifyItemChanged(index);
@@ -126,10 +130,13 @@
     }
 
     /**
-     * Compare package name using the same comparator as in {@link WidgetsListAdapter}.
-     * Also handle null row pointers.
+     * Compares the app name and then entry type for the given {@link WidgetsListBaseEntry}s.
+     *
+     * @Return 0 if both entries' order is the same. Negative integer if {@code newRowEntry} should
+     *         order before {@code orgRowEntry}. Positive integer if {@code orgRowEntry} should
+     *         order before {@code newRowEntry}.
      */
-    private int comparePackageName(WidgetsListBaseEntry curRow, WidgetsListBaseEntry newRow,
+    private int compareAppNameAndType(WidgetsListBaseEntry curRow, WidgetsListBaseEntry newRow,
             WidgetListBaseRowEntryComparator comparator) {
         if (curRow == null && newRow == null) {
             throw new IllegalStateException(
@@ -141,10 +148,18 @@
         } else if (curRow != null && newRow == null) {
             return -1; // old row needs to be deleted
         }
-        return comparator.compare(curRow, newRow);
+        int diff = comparator.compare(curRow, newRow);
+        if (diff == 0) {
+            return newRow.getRank() - curRow.getRank();
+        }
+        return diff;
     }
 
-    private boolean areWidgetsEqual(WidgetsListBaseEntry curRow,
+    /**
+     * Returns {@code true} if both {@code curRow} & {@code newRow} are
+     * {@link WidgetsListContentEntry}s with a different list of widgets.
+     */
+    private boolean hasWidgetsListChanged(WidgetsListBaseEntry curRow,
             WidgetsListBaseEntry newRow) {
         if (!(curRow instanceof WidgetsListContentEntry)
                 || !(newRow instanceof WidgetsListContentEntry)) {
@@ -152,7 +167,19 @@
         }
         WidgetsListContentEntry orgRowEntry = (WidgetsListContentEntry) curRow;
         WidgetsListContentEntry newRowEntry = (WidgetsListContentEntry) newRow;
-        return orgRowEntry.mWidgets.equals(newRowEntry.mWidgets);
+        return !orgRowEntry.mWidgets.equals(newRowEntry.mWidgets);
+    }
+
+    /**
+     * Returns {@code true} if {@code newRow} is {@link WidgetsListHeaderEntry} and its content has
+     * been changed due to user interactions.
+     */
+    private boolean hasHeaderUpdated(WidgetsListBaseEntry newRow) {
+        if (!(newRow instanceof WidgetsListHeaderEntry)) {
+            return false;
+        }
+        WidgetsListHeaderEntry newRowEntry = (WidgetsListHeaderEntry) newRow;
+        return newRowEntry.hasEntryUpdated();
     }
 
     private boolean isSamePackageItemInfo(PackageItemInfo curInfo, PackageItemInfo newInfo) {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
index 9d30842..5ec7f3b 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -24,6 +24,7 @@
 import android.view.View.OnLongClickListener;
 import android.view.ViewGroup;
 
+import androidx.annotation.Nullable;
 import androidx.recyclerview.widget.RecyclerView;
 import androidx.recyclerview.widget.RecyclerView.Adapter;
 import androidx.recyclerview.widget.RecyclerView.ViewHolder;
@@ -36,32 +37,42 @@
 import com.android.launcher3.widget.WidgetCell;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.picker.WidgetsListHeaderViewHolderBinder.OnHeaderClickListener;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.Comparator;
 import java.util.List;
+import java.util.stream.Collectors;
 
 /**
- * List view adapter for the widget tray.
+ * Recycler view adapter for the widget tray.
  *
- * <p>Memory vs. Performance:
- * The less number of types of views are inserted into a {@link RecyclerView}, the more recycling
- * happens and less memory is consumed.
+ * <p>This adapter supports view binding of subclasses of {@link WidgetsListBaseEntry}. There are 2
+ * subclasses: {@link WidgetsListHeader} & {@link WidgetsListContentEntry}.
+ * {@link WidgetsListHeader} entries are always visible in the recycler view. At most one
+ * {@link WidgetsListContentEntry} is shown in the recycler view at any time. Clicking a
+ * {@link WidgetsListHeader} will result in expanding / collapsing a corresponding
+ * {@link WidgetsListContentEntry} of the same app.
  */
-public class WidgetsListAdapter extends Adapter<ViewHolder> {
+public class WidgetsListAdapter extends Adapter<ViewHolder> implements OnHeaderClickListener {
 
     private static final String TAG = "WidgetsListAdapter";
     private static final boolean DEBUG = false;
 
     /** Uniquely identifies widgets list view type within the app. */
     private static final int VIEW_TYPE_WIDGETS_LIST = R.layout.widgets_list_row_view;
+    private static final int VIEW_TYPE_WIDGETS_HEADER = R.layout.widgets_list_row_header;
 
     private final WidgetsDiffReporter mDiffReporter;
     private final SparseArray<ViewHolderBinder> mViewHolderBinders = new SparseArray<>();
     private final WidgetsListRowViewHolderBinder mWidgetsListRowViewHolderBinder;
+    private final WidgetListBaseRowEntryComparator mRowComparator =
+            new WidgetListBaseRowEntryComparator();
 
-    private ArrayList<WidgetsListBaseEntry> mEntries = new ArrayList<>();
+    private List<WidgetsListBaseEntry> mAllEntries = new ArrayList<>();
+    private ArrayList<WidgetsListBaseEntry> mVisibleEntries = new ArrayList<>();
+    @Nullable private String mWidgetsContentVisiblePackage = null;
 
     public WidgetsListAdapter(Context context, LayoutInflater layoutInflater,
             WidgetPreviewLoader widgetPreviewLoader, IconCache iconCache,
@@ -70,6 +81,8 @@
         mWidgetsListRowViewHolderBinder = new WidgetsListRowViewHolderBinder(context,
                 layoutInflater, iconClickListener, iconLongClickListener, widgetPreviewLoader);
         mViewHolderBinders.put(VIEW_TYPE_WIDGETS_LIST, mWidgetsListRowViewHolderBinder);
+        mViewHolderBinders.put(VIEW_TYPE_WIDGETS_HEADER,
+                new WidgetsListHeaderViewHolderBinder(layoutInflater, this::onHeaderClicked));
     }
 
     /**
@@ -96,26 +109,39 @@
 
     @Override
     public int getItemCount() {
-        return mEntries.size();
+        return mVisibleEntries.size();
     }
 
     /** Gets the section name for {@link com.android.launcher3.views.RecyclerViewFastScroller}. */
     public String getSectionName(int pos) {
-        return mEntries.get(pos).mTitleSectionName;
+        return mVisibleEntries.get(pos).mTitleSectionName;
     }
 
     /** Updates the widget list. */
     public void setWidgets(List<WidgetsListBaseEntry> tempEntries) {
-        ArrayList<WidgetsListBaseEntry> newEntries = new ArrayList<>(tempEntries);
-        WidgetListBaseRowEntryComparator rowComparator = new WidgetListBaseRowEntryComparator();
-        Collections.sort(newEntries, rowComparator);
-        mDiffReporter.process(mEntries, newEntries, rowComparator);
+        mAllEntries = tempEntries.stream().sorted(mRowComparator)
+                .collect(Collectors.toList());
+        updateVisibleEntries();
+    }
+
+    private void updateVisibleEntries() {
+        mAllEntries.forEach(entry -> {
+            if (entry instanceof WidgetsListHeaderEntry) {
+                ((WidgetsListHeaderEntry) entry).setIsWidgetListShown(
+                        entry.mPkgItem.packageName.equals(mWidgetsContentVisiblePackage));
+            }
+        });
+        List<WidgetsListBaseEntry> newVisibleEntries = mAllEntries.stream()
+                .filter(entry -> entry instanceof WidgetsListHeaderEntry
+                        || entry.mPkgItem.packageName.equals(mWidgetsContentVisiblePackage))
+                .collect(Collectors.toList());
+        mDiffReporter.process(mVisibleEntries, newVisibleEntries, mRowComparator);
     }
 
     @Override
     public void onBindViewHolder(ViewHolder holder, int pos) {
         ViewHolderBinder viewHolderBinder = mViewHolderBinders.get(getItemViewType(pos));
-        viewHolderBinder.bindViewHolder(holder, mEntries.get(pos));
+        viewHolderBinder.bindViewHolder(holder, mVisibleEntries.get(pos));
     }
 
     @Override
@@ -148,13 +174,26 @@
 
     @Override
     public int getItemViewType(int pos) {
-        WidgetsListBaseEntry entry = mEntries.get(pos);
+        WidgetsListBaseEntry entry = mVisibleEntries.get(pos);
         if (entry instanceof WidgetsListContentEntry) {
             return VIEW_TYPE_WIDGETS_LIST;
+        } else if (entry instanceof WidgetsListHeaderEntry) {
+            return VIEW_TYPE_WIDGETS_HEADER;
         }
         throw new UnsupportedOperationException("ViewHolderBinder not found for " + entry);
     }
 
+    @Override
+    public void onHeaderClicked(boolean showWidgets, String expandedPackage) {
+        if (showWidgets) {
+            mWidgetsContentVisiblePackage = expandedPackage;
+            updateVisibleEntries();
+        } else if (expandedPackage.equals(mWidgetsContentVisiblePackage)) {
+            mWidgetsContentVisiblePackage = null;
+            updateVisibleEntries();
+        }
+    }
+
     /** Comparator for sorting WidgetListRowEntry based on package title. */
     public static class WidgetListBaseRowEntryComparator implements
             Comparator<WidgetsListBaseEntry> {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeader.java b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
new file mode 100644
index 0000000..823fb7b
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
@@ -0,0 +1,205 @@
+/*
+ * 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 static com.android.launcher3.FastBitmapDrawable.newIcon;
+
+import android.content.Context;
+import android.content.res.Resources;
+import android.content.res.TypedArray;
+import android.graphics.drawable.Drawable;
+import android.util.AttributeSet;
+import android.widget.CheckBox;
+import android.widget.ImageView;
+import android.widget.LinearLayout;
+import android.widget.TextView;
+
+import androidx.annotation.Nullable;
+import androidx.annotation.UiThread;
+
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.FastBitmapDrawable;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.R;
+import com.android.launcher3.graphics.PlaceHolderIconDrawable;
+import com.android.launcher3.icons.IconCache.ItemInfoUpdateReceiver;
+import com.android.launcher3.icons.cache.HandlerRunnable;
+import com.android.launcher3.model.data.ItemInfoWithIcon;
+import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+
+/**
+ * A UI represents a header of an app shown in the full widgets tray.
+ *
+ * It is a {@link LinearLayout} which contains an app icon, an app name, a subtitle and a checkbox
+ * which indicates if the widgets content view underneath this header should be shown.
+ */
+public final class WidgetsListHeader extends LinearLayout implements ItemInfoUpdateReceiver {
+
+    private boolean mEnableIconUpdateAnimation = false;
+
+    @Nullable private HandlerRunnable mIconLoadRequest;
+    @Nullable private Drawable mIconDrawable;
+    private final int mIconSize;
+
+    private ImageView mAppIcon;
+    private TextView mTitle;
+    private TextView mSubtitle;
+
+    private CheckBox mExpandToggle;
+    private boolean mIsExpanded = false;
+
+    public WidgetsListHeader(Context context) {
+        this(context, /* attrs= */ null);
+    }
+
+    public WidgetsListHeader(Context context, @Nullable AttributeSet attrs) {
+        this(context, attrs, /* defStyle= */ 0);
+    }
+
+    public WidgetsListHeader(Context context, @Nullable AttributeSet attrs, int defStyleAttr) {
+        super(context, attrs, defStyleAttr);
+
+        ActivityContext activity = ActivityContext.lookupContext(context);
+        DeviceProfile grid = activity.getDeviceProfile();
+        TypedArray a = context.obtainStyledAttributes(attrs,
+                R.styleable.WidgetsListRowHeader, defStyleAttr, /* defStyleRes= */ 0);
+        mIconSize = a.getDimensionPixelSize(R.styleable.WidgetsListRowHeader_appIconSize,
+                grid.iconSizePx);
+    }
+
+    @Override
+    protected void onFinishInflate() {
+        super.onFinishInflate();
+        mAppIcon = findViewById(R.id.app_icon);
+        mTitle = findViewById(R.id.app_title);
+        mSubtitle = findViewById(R.id.app_subtitle);
+        mExpandToggle = findViewById(R.id.toggle);
+    }
+
+    /**
+     * Sets a {@link OnExpansionChangeListener} to get a callback when this app widgets section
+     * expands / collapses.
+     */
+    @UiThread
+    public void setOnExpandChangeListener(
+            @Nullable OnExpansionChangeListener onExpandChangeListener) {
+        // Use the entire touch area of this view to expand / collapse an app widgets section.
+        setOnClickListener(view -> {
+            setExpanded(!mIsExpanded);
+            onExpandChangeListener.onExpansionChange(mIsExpanded);
+        });
+    }
+
+    /** Sets the expand toggle to expand / collapse. */
+    @UiThread
+    public void setExpanded(boolean isExpanded) {
+        this.mIsExpanded = isExpanded;
+        mExpandToggle.setChecked(isExpanded);
+    }
+
+    /** Apply app icon, labels and tag using a generic {@link WidgetsListHeaderEntry}. */
+    @UiThread
+    public void applyFromItemInfoWithIcon(WidgetsListHeaderEntry entry) {
+        applyIconAndLabel(entry);
+    }
+
+    @UiThread
+    private void applyIconAndLabel(WidgetsListHeaderEntry entry) {
+        PackageItemInfo info = entry.mPkgItem;
+        setIcon(info);
+        setTitles(entry);
+        setExpanded(entry.isWidgetListShown());
+
+        super.setTag(info);
+
+        verifyHighRes();
+    }
+
+    private void setIcon(PackageItemInfo info) {
+        FastBitmapDrawable icon = newIcon(getContext(), info);
+        applyDrawables(icon);
+        mIconDrawable = icon;
+        if (mIconDrawable != null) {
+            mIconDrawable.setVisible(
+                    /* visible= */ getWindowVisibility() == VISIBLE && isShown(),
+                    /* restart= */ false);
+        }
+    }
+
+    private void applyDrawables(Drawable icon) {
+        icon.setBounds(0, 0, mIconSize, mIconSize);
+
+        mAppIcon.setImageDrawable(icon);
+
+        // If the current icon is a placeholder color, animate its update.
+        if (mIconDrawable != null
+                && mIconDrawable instanceof PlaceHolderIconDrawable
+                && mEnableIconUpdateAnimation) {
+            ((PlaceHolderIconDrawable) mIconDrawable).animateIconUpdate(icon);
+        }
+    }
+
+    private void setTitles(WidgetsListHeaderEntry entry) {
+        mTitle.setText(entry.mPkgItem.title);
+
+        if (entry.widgetsCount > 0) {
+            Resources resources = getContext().getResources();
+            mSubtitle.setText(resources.getQuantityString(R.plurals.widgets_tray_subtitle,
+                    entry.widgetsCount, entry.widgetsCount));
+            mSubtitle.setVisibility(VISIBLE);
+        } else {
+            mSubtitle.setVisibility(GONE);
+        }
+    }
+
+    @Override
+    public void reapplyItemInfo(ItemInfoWithIcon info) {
+        if (getTag() == info) {
+            mIconLoadRequest = null;
+            mEnableIconUpdateAnimation = true;
+
+            // Optimization: Starting in N, pre-uploads the bitmap to RenderThread.
+            info.bitmap.icon.prepareToDraw();
+
+            setIcon((PackageItemInfo) info);
+
+            mEnableIconUpdateAnimation = false;
+        }
+    }
+
+    /** Verifies that the current icon is high-res otherwise posts a request to load the icon. */
+    public void verifyHighRes() {
+        if (mIconLoadRequest != null) {
+            mIconLoadRequest.cancel();
+            mIconLoadRequest = null;
+        }
+        if (getTag() instanceof ItemInfoWithIcon) {
+            ItemInfoWithIcon info = (ItemInfoWithIcon) getTag();
+            if (info.usingLowResIcon()) {
+                mIconLoadRequest = LauncherAppState.getInstance(getContext()).getIconCache()
+                        .updateIconInBackground(this, info);
+            }
+        }
+    }
+
+    /** A listener for the widget section expansion / collapse events. */
+    public interface OnExpansionChangeListener {
+        /** Notifies that the widget section is expanded or collapsed. */
+        void onExpansionChange(boolean isExpanded);
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeaderHolder.java b/src/com/android/launcher3/widget/picker/WidgetsListHeaderHolder.java
new file mode 100644
index 0000000..d4e1b1c
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeaderHolder.java
@@ -0,0 +1,32 @@
+/*
+ * 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 androidx.recyclerview.widget.RecyclerView.ViewHolder;
+
+/**
+ * A {@link ViewHolder} for {@link WidgetsListHeader} of an app, which renders the app icon, the app
+ * name, label and a button for showing / hiding widgets.
+ */
+public final class WidgetsListHeaderHolder extends ViewHolder {
+    final WidgetsListHeader mWidgetsListHeader;
+
+    public WidgetsListHeaderHolder(WidgetsListHeader view) {
+        super(view);
+
+        mWidgetsListHeader = view;
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
new file mode 100644
index 0000000..ed53e6f
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
@@ -0,0 +1,61 @@
+/*
+ * 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.LayoutInflater;
+import android.view.ViewGroup;
+
+import com.android.launcher3.R;
+import com.android.launcher3.recyclerview.ViewHolderBinder;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+
+/**
+ * Binds data from {@link WidgetsListHeaderEntry} to UI elements in {@link WidgetsListHeaderHolder}.
+ */
+public final class WidgetsListHeaderViewHolderBinder implements
+        ViewHolderBinder<WidgetsListHeaderEntry, WidgetsListHeaderHolder> {
+    private final LayoutInflater mLayoutInflater;
+    private final OnHeaderClickListener mOnHeaderClickListener;
+
+    public WidgetsListHeaderViewHolderBinder(LayoutInflater layoutInflater,
+            OnHeaderClickListener onHeaderClickListener) {
+        mLayoutInflater = layoutInflater;
+        mOnHeaderClickListener = onHeaderClickListener;
+    }
+
+    @Override
+    public WidgetsListHeaderHolder newViewHolder(ViewGroup parent) {
+        WidgetsListHeader header = (WidgetsListHeader) mLayoutInflater.inflate(
+                R.layout.widgets_list_row_header, parent, false);
+
+        return new WidgetsListHeaderHolder(header);
+    }
+
+    @Override
+    public void bindViewHolder(WidgetsListHeaderHolder viewHolder, WidgetsListHeaderEntry data) {
+        WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
+        widgetsListHeader.applyFromItemInfoWithIcon(data);
+        widgetsListHeader.setExpanded(data.isWidgetListShown());
+        widgetsListHeader.setOnExpandChangeListener(isExpanded ->
+                mOnHeaderClickListener.onHeaderClicked(isExpanded, data.mPkgItem.packageName));
+    }
+
+    /** A listener to be invoked when {@link WidgetsListHeader} is clicked. */
+    public interface OnHeaderClickListener {
+        /** Calls when {@link WidgetsListHeader} is clicked to show / hide widgets for a package. */
+        void onHeaderClicked(boolean showWidgets, String packageName);
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java
index 22a8d00..cec6b80 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java
@@ -76,7 +76,7 @@
         }
 
         ViewGroup container = (ViewGroup) mLayoutInflater.inflate(
-                R.layout.widgets_list_row_view, parent, false);
+                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.
@@ -122,9 +122,6 @@
             }
         }
 
-        // Bind the views in the application info section.
-        holder.title.applyFromItemInfoWithIcon(entry.mPkgItem);
-
         // Bind the view in the widget horizontal tray region.
         for (int i = 0; i < infoList.size(); i++) {
             WidgetCell widget = (WidgetCell) row.getChildAt(2 * i);
diff --git a/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java b/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
index 9be079e..ae94584 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
@@ -19,20 +19,16 @@
 
 import androidx.recyclerview.widget.RecyclerView.ViewHolder;
 
-import com.android.launcher3.BubbleTextView;
 import com.android.launcher3.R;
 
-/** A {@link ViewHolder} for a row in the full widget picker. */
+/** 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 BubbleTextView title;
 
     public WidgetsRowViewHolder(ViewGroup v) {
         super(v);
 
         cellContainer = v.findViewById(R.id.widgets_cell_list);
-        title = v.findViewById(R.id.section);
-        title.setAccessibilityDelegate(null);
     }
 }
diff --git a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
index f27922b..30c9b5f 100644
--- a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
+++ b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
@@ -31,6 +31,7 @@
 import com.android.launcher3.widget.WidgetManagerHelper;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 import com.android.launcher3.widget.picker.WidgetsDiffReporter;
 
 import java.util.ArrayList;
@@ -73,11 +74,11 @@
 
         for (Map.Entry<PackageItemInfo, List<WidgetItem>> entry : mWidgetsList.entrySet()) {
             PackageItemInfo pkgItem = entry.getKey();
+            List<WidgetItem> widgetItems = entry.getValue();
             String sectionName = (pkgItem.title == null) ? "" :
                     indexer.computeSectionName(pkgItem.title);
-            WidgetsListContentEntry row =
-                    new WidgetsListContentEntry(pkgItem, sectionName, entry.getValue());
-            result.add(row);
+            result.add(new WidgetsListHeaderEntry(pkgItem, sectionName, widgetItems));
+            result.add(new WidgetsListContentEntry(pkgItem, sectionName, widgetItems));
         }
         return result;
     }
diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
index 9d4ccff..737f891 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -92,9 +92,8 @@
 
         // Drag widget to homescreen
         WidgetConfigStartupMonitor monitor = new WidgetConfigStartupMonitor();
-        widgets.
-                getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager())).
-                dragToWorkspace(true, false);
+        widgets.getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager()))
+                .dragToWorkspace(true, false);
         // Widget id for which the config activity was opened
         mWidgetId = monitor.getWidgetId();
 
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index 49af616..f95abdb 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -31,6 +31,7 @@
 import com.android.launcher3.testing.TestProtocol;
 
 import java.util.Collection;
+import java.util.List;
 
 /**
  * All widgets container.
@@ -101,22 +102,28 @@
         try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck();
              LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
                      "getting widget " + labelText + " in widgets list")) {
-            final UiObject2 widgetsContainer = verifyActiveContainer();
+            final UiObject2 fullWidgetsPicker = verifyActiveContainer();
             mLauncher.assertTrue("Widgets container didn't become scrollable",
-                    widgetsContainer.wait(Until.scrollable(true), WAIT_TIME_MS));
+                    fullWidgetsPicker.wait(Until.scrollable(true), WAIT_TIME_MS));
             final Point displaySize = mLauncher.getRealDisplaySize();
-            final BySelector labelSelector = By.clazz("android.widget.TextView").text(labelText);
 
+            final UiObject2 widgetsContainer = findTestAppWidgetsScrollContainer();
+            mLauncher.assertTrue("Can't locate widgets list for the test app: "
+                                    + mLauncher.getLauncherPackageName(),
+                    widgetsContainer != null);
+            final BySelector labelSelector = By.clazz("android.widget.TextView").text(labelText);
             int i = 0;
             for (; ; ) {
-                final Collection<UiObject2> cells = mLauncher.getObjectsInContainer(
-                        widgetsContainer, "widgets_scroll_container");
-                mLauncher.assertTrue("Widgets doesn't have 2 rows", cells.size() >= 2);
+                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 = label.getParent().getParent();
+                    final UiObject2 widget = cell;
                     mLauncher.assertEquals(
                             "View is not WidgetCell",
                             "com.android.launcher3.widget.WidgetCell",
@@ -131,7 +138,7 @@
                             <= displaySize.y - mLauncher.getBottomGestureSize()) {
                         int visibleDelta = maxWidth - mLauncher.getVisibleBounds(widget).width();
                         if (visibleDelta > 0) {
-                            Rect parentBounds = mLauncher.getVisibleBounds(cell);
+                            Rect parentBounds = mLauncher.getVisibleBounds(cell.getParent());
                             mLauncher.linearGesture(parentBounds.centerX() + visibleDelta
                                             + mLauncher.getTouchSlop(),
                                     parentBounds.centerY(), parentBounds.centerX(),
@@ -153,4 +160,53 @@
             }
         }
     }
+
+    /** Finds the widgets list of this test app from the collapsed full widgets picker. */
+    private UiObject2 findTestAppWidgetsScrollContainer() {
+        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");
+
+        boolean hasHeaderExpanded = false;
+        for (int i = 0; i < 40; i++) {
+            UiObject2 fullWidgetsPicker = verifyActiveContainer();
+
+            UiObject2 header = fullWidgetsPicker.findObject(headerSelector);
+            mLauncher.assertTrue("Can't find a widget header", header != null);
+
+            // Look for a header that has the test app name.
+            UiObject2 headerTitle = fullWidgetsPicker.findObject(targetAppSelector);
+            if (headerTitle != null) {
+                // If we find the header and it has not been expanded, let's click it to see the
+                // widgets list.
+                if (!hasHeaderExpanded) {
+                    hasHeaderExpanded = true;
+                    mLauncher.clickLauncherObject(headerTitle);
+                    // After clicking the header, the recyclerview has been updated. Let's refresh
+                    // the container UIObject2.
+                    fullWidgetsPicker = verifyActiveContainer();
+                    // Refresh headerTitle because the first instance is stale after
+                    // verifyActiveContainer call.
+                    headerTitle = fullWidgetsPicker.findObject(targetAppSelector);
+                }
+
+                // 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);
+            }
+        }
+
+        return null;
+    }
 }