Maintain selected or last header on clicking "Show all" button

Bug: 356127021
Test: Manual and tapl test
Flag: com.android.launcher3.enable_tiered_widgets_by_default_in_picker
Change-Id: I762695f9999e5ea4c722dbedde6d7efabd3ab0cb

res/layout/widgets_list_expand_button.xml

@@ -15,6 +15,7 @@
   -->
 
 <Button xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/widget_list_expand_button"
     style="@style/Button.Rounded.Colored"
     android:layout_width="wrap_content"
     android:layout_height="wrap_content"

src/com/android/launcher3/views/StickyHeaderLayout.java

@@ -120,7 +120,19 @@ public class StickyHeaderLayout extends LinearLayout implements
     }
 
     private float getCurrentScroll() {
-        return mScrollOffset + (mCurrentEmptySpaceView == null ? 0 : mCurrentEmptySpaceView.getY());
+        float scroll;
+        if (mCurrentRecyclerView.getVisibility() != VISIBLE) {
+            // When no list is displayed, assume no scroll.
+            scroll = 0f;
+        } else if (mCurrentEmptySpaceView != null) {
+            // Otherwise use empty space view as reference to position.
+            scroll = mCurrentEmptySpaceView.getY();
+        } else {
+            // If there is no empty space view, but the list is visible, we are scrolled away
+            // completely, so assume all non-sticky children should also be scrolled away.
+            scroll = -mHeaderHeight;
+        }
+        return mScrollOffset + scroll;
     }
 
     @Override

src/com/android/launcher3/widget/picker/WidgetsFullSheet.java

@@ -605,9 +605,12 @@ public class WidgetsFullSheet extends BaseWidgetSheet
             mAdapters.get(AdapterHolder.SEARCH).mWidgetsRecyclerView.setVisibility(GONE);
             mAdapters.get(getCurrentAdapterHolderType()).mWidgetsRecyclerView.setVisibility(
                     VISIBLE);
-            // Visibility of recommended widgets, recycler views and headers are handled in methods
-            // below.
-            post(this::onRecommendedWidgetsBound);
+            if (mRecommendedWidgetsCount > 0) {
+                // Display recommendations immediately, if present, so that other parts of sticky
+                // header (e.g. personal / work tabs) don't flash in interim.
+                mWidgetRecommendationsContainer.setVisibility(VISIBLE);
+            }
+            // Visibility of recycler views and headers are handled in methods below.
             onWidgetsBound();
         }
     }

src/com/android/launcher3/widget/picker/WidgetsListAdapter.java

@@ -103,7 +103,7 @@ public class WidgetsListAdapter extends Adapter<ViewHolder> implements OnHeaderC
                             .equals(mWidgetsContentVisiblePackageUserKey);
     @Nullable private Predicate<WidgetsListBaseEntry> mFilter = null;
     @Nullable private RecyclerView mRecyclerView;
-    @Nullable private PackageUserKey mPendingClickHeader;
+    @Nullable private PackageUserKey mHeaderPositionToMaintain;
     @Px private int mMaxHorizontalSpan;
 
     private boolean mShowOnlyDefaultList = true;
@@ -215,7 +215,7 @@ public class WidgetsListAdapter extends Adapter<ViewHolder> implements OnHeaderC
         // Get the current top of the header with the matching key before adjusting the visible
         // entries.
         OptionalInt previousPositionForPackageUserKey =
-                getPositionForPackageUserKey(mPendingClickHeader);
+                getPositionForPackageUserKey(mHeaderPositionToMaintain);
         OptionalInt topForPackageUserKey =
                 getOffsetForPosition(previousPositionForPackageUserKey);
 
@@ -247,13 +247,15 @@ public class WidgetsListAdapter extends Adapter<ViewHolder> implements OnHeaderC
         mVisibleEntries.addAll(newVisibleEntries);
         diffResult.dispatchUpdatesTo(this);
 
-        if (mPendingClickHeader != null) {
+        if (mHeaderPositionToMaintain != null && mRecyclerView != null) {
             // Get the position for the clicked header after adjusting the visible entries. The
             // position may have changed if another header had previously been expanded.
             OptionalInt positionForPackageUserKey =
-                    getPositionForPackageUserKey(mPendingClickHeader);
-            scrollToPositionAndMaintainOffset(positionForPackageUserKey, topForPackageUserKey);
-            mPendingClickHeader = null;
+                    getPositionForPackageUserKey(mHeaderPositionToMaintain);
+            // Post scroll updates to be applied after diff updates.
+            mRecyclerView.post(() -> scrollToPositionAndMaintainOffset(positionForPackageUserKey,
+                    topForPackageUserKey));
+            mHeaderPositionToMaintain = null;
         }
     }
 
@@ -384,7 +386,7 @@ public class WidgetsListAdapter extends Adapter<ViewHolder> implements OnHeaderC
 
         // Store the header that was clicked so that its position will be maintained the next time
         // we update the entries.
-        mPendingClickHeader = packageUserKey;
+        mHeaderPositionToMaintain = packageUserKey;
 
         updateVisibleEntries();
 
@@ -470,6 +472,16 @@ public class WidgetsListAdapter extends Adapter<ViewHolder> implements OnHeaderC
      */
     public void useExpandedList() {
         mShowOnlyDefaultList = false;
+        if (mWidgetsContentVisiblePackageUserKey != null) {
+            // Maintain selected header for the next update that expands the list.
+            mHeaderPositionToMaintain = mWidgetsContentVisiblePackageUserKey;
+        } else if (mVisibleEntries.size() > 2) {
+            // Maintain last visible header shown above expand button since there was no selected
+            // header.
+            mHeaderPositionToMaintain = PackageUserKey.fromPackageItemInfo(
+                    mVisibleEntries.get(mVisibleEntries.size() - 2).mPkgItem);
+        }
+
     }
 
     /** Comparator for sorting WidgetListRowEntry based on package title. */