Update the widget picker UI to show the "view all" button

Screenshot test in follow up to include the button.

Bug: 356127021
Flag: com.android.launcher3.enable_tiered_widgets_by_default_in_picker
Test: Manual
Change-Id: Iac847361a0fa4a4b57d37e802ce6a7a5d8bb17d7
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index e5cd76a..2550ebb 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -175,6 +175,10 @@
         @UiEvent(doc = "User searched for a widget in the widget picker.")
         LAUNCHER_WIDGETSTRAY_SEARCHED(819),
 
+        @UiEvent(doc = "User clicked on view all button to expand the displayed list in the "
+                + "widget picker.")
+        LAUNCHER_WIDGETSTRAY_EXPAND_PRESS(1978),
+
         @UiEvent(doc = "A dragged item is dropped on 'Remove' button in the target bar")
         LAUNCHER_ITEM_DROPPED_ON_REMOVE(465),
 
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 1860977..2f64ab1 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -16,12 +16,16 @@
 package com.android.launcher3.widget.picker;
 
 import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
+import static com.android.launcher3.Flags.enableTieredWidgetsByDefaultInPicker;
 import static com.android.launcher3.Flags.enableUnfoldedTwoPanePicker;
 import static com.android.launcher3.allapps.ActivityAllAppsContainerView.AdapterHolder.SEARCH;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_WIDGETSTRAY_EXPAND_PRESS;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_WIDGETSTRAY_SEARCHED;
 import static com.android.launcher3.testing.shared.TestProtocol.NORMAL_STATE_ORDINAL;
 import static com.android.launcher3.views.RecyclerViewFastScroller.FastScrollerLocation.WIDGET_SCROLLER;
 
+import static java.util.Collections.emptyList;
+
 import android.animation.Animator;
 import android.content.Context;
 import android.content.res.Resources;
@@ -68,6 +72,7 @@
 import com.android.launcher3.widget.BaseWidgetSheet;
 import com.android.launcher3.widget.WidgetCell;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.picker.model.data.WidgetPickerData;
 import com.android.launcher3.widget.picker.search.SearchModeListener;
 import com.android.launcher3.widget.picker.search.WidgetsSearchBar;
 import com.android.launcher3.widget.picker.search.WidgetsSearchBar.WidgetsSearchDataProvider;
@@ -87,7 +92,8 @@
  */
 public class WidgetsFullSheet extends BaseWidgetSheet
         implements OnActivePageChangedListener,
-        WidgetsRecyclerView.HeaderViewDimensionsProvider, SearchModeListener {
+        WidgetsRecyclerView.HeaderViewDimensionsProvider, SearchModeListener,
+        WidgetsListAdapter.ExpandButtonClickListener {
 
     private static final long FADE_IN_DURATION = 150;
 
@@ -257,7 +263,13 @@
         mSearchBar.initialize(new WidgetsSearchDataProvider() {
             @Override
             public List<WidgetsListBaseEntry> getWidgets() {
-                return getWidgetsToDisplay();
+                if (enableTieredWidgetsByDefaultInPicker()) {
+                    // search all
+                    return mActivityContext.getWidgetPickerDataProvider().get().getAllWidgets();
+                } else {
+                    // Can be removed when inlining enableTieredWidgetsByDefaultInPicker flag
+                    return getWidgetsToDisplay();
+                }
             }
         }, /* searchModeListener= */ this);
     }
@@ -482,6 +494,9 @@
     /**
      * Returns all displayable widgets.
      */
+    // Used by the two pane sheet to show 3-dot menu to toggle between default lists and all lists
+    // when enableTieredWidgetsByDefaultInPicker is OFF. This code path and the 3-dot menu can be
+    // safely deleted when it's alternative "enableTieredWidgetsByDefaultInPicker" flag is inlined.
     protected List<WidgetsListBaseEntry> getWidgetsToDisplay() {
         return mActivityContext.getWidgetPickerDataProvider().get().getAllWidgets();
     }
@@ -491,16 +506,27 @@
         if (mIsInSearchMode) {
             return;
         }
-        List<WidgetsListBaseEntry> widgets = getWidgetsToDisplay();
+        List<WidgetsListBaseEntry> widgets;
+        List<WidgetsListBaseEntry> defaultWidgets = emptyList();
+
+        if (enableTieredWidgetsByDefaultInPicker()) {
+            WidgetPickerData dataProvider =
+                    mActivityContext.getWidgetPickerDataProvider().get();
+            widgets = dataProvider.getAllWidgets();
+            defaultWidgets = dataProvider.getDefaultWidgets();
+        } else {
+            // This code path can be deleted once enableTieredWidgetsByDefaultInPicker is inlined.
+            widgets = getWidgetsToDisplay();
+        }
 
         AdapterHolder primaryUserAdapterHolder = mAdapters.get(AdapterHolder.PRIMARY);
-        primaryUserAdapterHolder.mWidgetsListAdapter.setWidgets(widgets);
+        primaryUserAdapterHolder.mWidgetsListAdapter.setWidgets(widgets, defaultWidgets);
 
         if (mHasWorkProfile) {
             mViewPager.setVisibility(VISIBLE);
             mTabBar.setVisibility(VISIBLE);
             AdapterHolder workUserAdapterHolder = mAdapters.get(AdapterHolder.WORK);
-            workUserAdapterHolder.mWidgetsListAdapter.setWidgets(widgets);
+            workUserAdapterHolder.mWidgetsListAdapter.setWidgets(widgets, defaultWidgets);
             onActivePageChanged(mViewPager.getCurrentPage());
         } else {
             onActivePageChanged(0);
@@ -518,6 +544,16 @@
     }
 
     @Override
+    public void onWidgetsListExpandButtonClick(View v) {
+        AdapterHolder currentAdapterHolder = mAdapters.get(getCurrentAdapterHolderType());
+        currentAdapterHolder.mWidgetsListAdapter.useExpandedList();
+        onWidgetsBound();
+        currentAdapterHolder.mWidgetsRecyclerView.announceForAccessibility(
+                mActivityContext.getString(R.string.widgets_list_expanded));
+        mActivityContext.getStatsLogManager().logger().log(LAUNCHER_WIDGETSTRAY_EXPAND_PRESS);
+    }
+
+    @Override
     public void enterSearchMode(boolean shouldLog) {
         if (mIsInSearchMode) return;
         setViewVisibilityBasedOnSearch(/*isInSearchMode= */ true);
@@ -832,7 +868,7 @@
                 + marginLayoutParams.topMargin;
     }
 
-    private int getCurrentAdapterHolderType() {
+    protected int getCurrentAdapterHolderType() {
         if (mIsInSearchMode) {
             return SEARCH;
         } else if (mViewPager != null) {
@@ -861,6 +897,7 @@
             WidgetsFullSheet sheet = show(BaseActivity.fromContext(getContext()), false);
             sheet.restoreRecommendations(mRecommendedWidgets, mRecommendedWidgetsMap);
             sheet.restoreHierarchyState(widgetsState);
+            sheet.restoreAdapterStates(mAdapters);
             sheet.restorePreviousAdapterHolderType(getCurrentAdapterHolderType());
         } else if (!isTwoPane()) {
             reset();
@@ -876,6 +913,17 @@
         mRecommendedWidgetsMap = recommendedWidgetsMap;
     }
 
+    private void restoreAdapterStates(SparseArray<AdapterHolder> adapters) {
+        if (adapters.contains(AdapterHolder.WORK)) {
+            mAdapters.get(AdapterHolder.WORK).mWidgetsListAdapter.restoreState(
+                    adapters.get(AdapterHolder.WORK).mWidgetsListAdapter);
+        }
+        mAdapters.get(AdapterHolder.PRIMARY).mWidgetsListAdapter.restoreState(
+                adapters.get(AdapterHolder.PRIMARY).mWidgetsListAdapter);
+        mAdapters.get(AdapterHolder.SEARCH).mWidgetsListAdapter.restoreState(
+                adapters.get(AdapterHolder.SEARCH).mWidgetsListAdapter);
+    }
+
     /**
      * Indicates if layout should be re-created on device profile change - so that a different
      * layout can be displayed.
@@ -1045,6 +1093,7 @@
                     this::getEmptySpaceHeight,
                     /* iconClickListener= */ WidgetsFullSheet.this,
                     /* iconLongClickListener= */ WidgetsFullSheet.this,
+                    /* expandButtonClickListener= */ WidgetsFullSheet.this,
                     isTwoPane());
             mWidgetsListAdapter.setHasStableIds(true);
             switch (mAdapterType) {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
index 3d3a669..3c67538 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -49,6 +49,7 @@
 import com.android.launcher3.widget.model.WidgetListSpaceEntry;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListExpandActionEntry;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 import com.android.launcher3.widget.util.WidgetSizes;
 
@@ -82,6 +83,7 @@
     public static final int VIEW_TYPE_WIDGETS_SPACE = R.id.view_type_widgets_space;
     public static final int VIEW_TYPE_WIDGETS_LIST = R.id.view_type_widgets_list;
     public static final int VIEW_TYPE_WIDGETS_HEADER = R.id.view_type_widgets_header;
+    public static final int VIEW_TYPE_WIDGETS_EXPAND = R.id.view_type_widgets_list_expand;
 
     private final Context mContext;
     private final SparseArray<ViewHolderBinder> mViewHolderBinders = new SparseArray<>();
@@ -90,7 +92,9 @@
     @Nullable private WidgetsTwoPaneSheet.HeaderChangeListener mHeaderChangeListener;
 
     private final List<WidgetsListBaseEntry> mAllEntries = new ArrayList<>();
+    private final List<WidgetsListBaseEntry> mAllDefaultEntries = new ArrayList<>();
     private ArrayList<WidgetsListBaseEntry> mVisibleEntries = new ArrayList<>();
+
     @Nullable private PackageUserKey mWidgetsContentVisiblePackageUserKey = null;
 
     private Predicate<WidgetsListBaseEntry> mHeaderAndSelectedContentFilter = entry ->
@@ -102,9 +106,12 @@
     @Nullable private PackageUserKey mPendingClickHeader;
     @Px private int mMaxHorizontalSpan;
 
+    private boolean mShowOnlyDefaultList = true;
+
     public WidgetsListAdapter(Context context, LayoutInflater layoutInflater,
             IntSupplier emptySpaceHeightProvider, OnClickListener iconClickListener,
             OnLongClickListener iconLongClickListener,
+            ExpandButtonClickListener expandButtonClickListener,
             boolean isTwoPane) {
         mContext = context;
         mMaxHorizontalSpan = WidgetSizes.getWidgetSizePx(
@@ -123,6 +130,16 @@
         mViewHolderBinders.put(
                 VIEW_TYPE_WIDGETS_SPACE,
                 new WidgetsSpaceViewHolderBinder(emptySpaceHeightProvider));
+        mViewHolderBinders.put(VIEW_TYPE_WIDGETS_EXPAND,
+                new WidgetsListExpandActionViewHolderBinder(layoutInflater,
+                        expandButtonClickListener::onWidgetsListExpandButtonClick));
+    }
+
+    /**
+     * Copies state info from another adapter.
+     */
+    public void restoreState(WidgetsListAdapter adapter) {
+        mShowOnlyDefaultList = adapter.mShowOnlyDefaultList;
     }
 
     public void setHeaderChangeListener(WidgetsTwoPaneSheet.HeaderChangeListener
@@ -168,10 +185,21 @@
     }
 
     /** Updates the widget list based on {@code tempEntries}. */
-    public void setWidgets(List<WidgetsListBaseEntry> tempEntries) {
+    public void setWidgets(List<WidgetsListBaseEntry> tempEntries,
+            List<WidgetsListBaseEntry> tempDefaultEntries) {
         mAllEntries.clear();
         mAllEntries.add(new WidgetListSpaceEntry());
         tempEntries.stream().sorted(mRowComparator).forEach(mAllEntries::add);
+
+        mAllDefaultEntries.clear();
+
+        if (mShowOnlyDefaultList && !tempDefaultEntries.isEmpty()) {
+            mAllDefaultEntries.add(new WidgetListSpaceEntry());
+            tempDefaultEntries.stream().sorted(mRowComparator).forEach(mAllDefaultEntries::add);
+            // Include view all action when default entries exist.
+            mAllDefaultEntries.add(new WidgetsListExpandActionEntry());
+        }
+
         updateVisibleEntries();
     }
 
@@ -179,7 +207,8 @@
     public void setWidgetsOnSearch(List<WidgetsListBaseEntry> searchResults) {
         // Forget the expanded package every time widget list is refreshed in search mode.
         mWidgetsContentVisiblePackageUserKey = null;
-        setWidgets(searchResults);
+        mShowOnlyDefaultList = false;
+        setWidgets(searchResults, /*tempDefaultEntries=*/ List.of());
     }
 
     private void updateVisibleEntries() {
@@ -190,10 +219,11 @@
         OptionalInt topForPackageUserKey =
                 getOffsetForPosition(previousPositionForPackageUserKey);
 
-        List<WidgetsListBaseEntry> newVisibleEntries = mAllEntries.stream()
+        List<WidgetsListBaseEntry> newVisibleEntries = getAllEntries().stream()
                 .filter(entry -> (((mFilter == null || mFilter.test(entry))
                         && mHeaderAndSelectedContentFilter.test(entry))
-                        || entry instanceof WidgetListSpaceEntry)
+                        || entry instanceof WidgetListSpaceEntry
+                        || entry instanceof WidgetsListExpandActionEntry)
                         && (mHeaderChangeListener == null
                         || !(entry instanceof WidgetsListContentEntry)))
                 .map(entry -> {
@@ -227,6 +257,11 @@
         }
     }
 
+    private List<WidgetsListBaseEntry> getAllEntries() {
+        return (mShowOnlyDefaultList && !mAllDefaultEntries.isEmpty()) ? mAllDefaultEntries
+                : mAllEntries;
+    }
+
     /** Returns whether {@code entry} matches {@code key}. */
     private static boolean isHeaderForPackageUserKey(
             @NonNull WidgetsListBaseEntry entry, @Nullable PackageUserKey key) {
@@ -262,7 +297,13 @@
 
         // The first entry has an empty space, count from second entries.
         int listPos = (pos > 1) ? POSITION_DEFAULT : POSITION_FIRST;
-        if (pos == (getItemCount() - 1)) {
+        int lastIndex = getItemCount() - 1;
+        // Last index may be the view all entry
+        int actualLastItemIndex = (mVisibleEntries.get(
+                lastIndex) instanceof WidgetsListExpandActionEntry) ? getItemCount() - 2
+                : getItemCount() - 1;
+
+        if (pos == (actualLastItemIndex)) {
             listPos |= POSITION_LAST;
         }
         viewHolderBinder.bindViewHolder(holder, mVisibleEntries.get(pos), listPos, payloads);
@@ -319,6 +360,8 @@
             return VIEW_TYPE_WIDGETS_HEADER;
         } else if (entry instanceof WidgetListSpaceEntry) {
             return VIEW_TYPE_WIDGETS_SPACE;
+        } else if (entry instanceof WidgetsListExpandActionEntry) {
+            return VIEW_TYPE_WIDGETS_EXPAND;
         }
         throw new UnsupportedOperationException("ViewHolderBinder not found for " + entry);
     }
@@ -412,6 +455,23 @@
         updateVisibleEntries();
     }
 
+    /**
+     * Returns the widget content {@link WidgetsListContentEntry} for a selected header.
+     */
+    public WidgetsListContentEntry getContentEntry(PackageUserKey selectedHeader) {
+        return getAllEntries().stream().filter(entry -> entry instanceof WidgetsListContentEntry)
+                .map(entry -> (WidgetsListContentEntry) entry)
+                .filter(entry -> PackageUserKey.fromPackageItemInfo(entry.mPkgItem).equals(
+                        selectedHeader)).findFirst().orElse(null);
+    }
+
+    /**
+     * Sets adapter to use expanded list when updating widgets.
+     */
+    public void useExpandedList() {
+        mShowOnlyDefaultList = false;
+    }
+
     /** Comparator for sorting WidgetListRowEntry based on package title. */
     public static class WidgetListBaseRowEntryComparator implements
             Comparator<WidgetsListBaseEntry> {
@@ -430,4 +490,10 @@
             return 1;
         }
     }
+
+    /** Callback interface for the interaction with the expand button */
+    public interface ExpandButtonClickListener {
+        /** Called when user clicks the button at end of widget apps list to expand it. */
+        void onWidgetsListExpandButtonClick(View view);
+    }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
index f4b99a0..f9bd5f1 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
@@ -16,6 +16,7 @@
 package com.android.launcher3.widget.picker;
 
 import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
+import static com.android.launcher3.Flags.enableTieredWidgetsByDefaultInPicker;
 import static com.android.launcher3.Flags.enableUnfoldedTwoPanePicker;
 import static com.android.launcher3.UtilitiesKt.CLIP_CHILDREN_FALSE_MODIFIER;
 import static com.android.launcher3.UtilitiesKt.CLIP_TO_PADDING_FALSE_MODIFIER;
@@ -147,7 +148,9 @@
         mHeaderDescription = mContent.findViewById(R.id.widget_picker_description);
 
         mWidgetOptionsMenu = mContent.findViewById(R.id.widget_picker_widget_options_menu);
-        setupWidgetOptionsMenu();
+        if (!enableTieredWidgetsByDefaultInPicker()) {
+            setupWidgetOptionsMenu();
+        }
 
         mRightPane = mContent.findViewById(R.id.right_pane);
         mRightPaneScrollView = mContent.findViewById(R.id.right_pane_scroll_view);
@@ -286,6 +289,9 @@
         }
     }
 
+    // Used by the two pane sheet to show 3-dot menu to toggle between default lists and all lists
+    // when enableTieredWidgetsByDefaultInPicker is OFF. This code path and the 3-dot menu can be
+    // safely deleted when it's alternative "enableTieredWidgetsByDefaultInPicker" flag is inlined.
     @Override
     protected List<WidgetsListBaseEntry> getWidgetsToDisplay() {
         List<WidgetsListBaseEntry> allWidgets =
@@ -319,6 +325,15 @@
     }
 
     @Override
+    public void onWidgetsListExpandButtonClick(View v) {
+        super.onWidgetsListExpandButtonClick(v);
+        // Refresh right pane with updated data for the selected header.
+        if (mSelectedHeader != null && mSelectedHeader != mSuggestedWidgetsPackageUserKey) {
+            getHeaderChangeListener().onHeaderChanged(mSelectedHeader);
+        }
+    }
+
+    @Override
     public void onRecommendedWidgetsBound() {
         super.onRecommendedWidgetsBound();
 
@@ -511,11 +526,20 @@
                         && !mOpenCloseAnimation.getAnimationPlayer().isRunning()
                         && !getAccessibilityInitialFocusView().isAccessibilityFocused();
                 mSelectedHeader = selectedHeader;
-                final boolean showDefaultWidgets = mWidgetOptionsMenuState != null
-                        && !mWidgetOptionsMenuState.showAllWidgets;
-                WidgetsListContentEntry contentEntry = findContentEntryForPackageUser(
-                        mActivityContext.getWidgetPickerDataProvider().get(),
-                        selectedHeader, showDefaultWidgets);
+
+                WidgetsListContentEntry contentEntry;
+                if (enableTieredWidgetsByDefaultInPicker()) {
+                    contentEntry = mAdapters.get(
+                            getCurrentAdapterHolderType()).mWidgetsListAdapter.getContentEntry(
+                            selectedHeader);
+                } else { // Can be deleted when inlining the "enableTieredWidgetsByDefaultInPicker"
+                    // flag
+                    final boolean showDefaultWidgets = mWidgetOptionsMenuState != null
+                            && !mWidgetOptionsMenuState.showAllWidgets;
+                    contentEntry = findContentEntryForPackageUser(
+                            mActivityContext.getWidgetPickerDataProvider().get(),
+                            selectedHeader, showDefaultWidgets);
+                }
 
                 if (contentEntry == null || mRightPane == null) {
                     return;