Retain add button on rotation for two-pane and bottom sheet

Reselects the WidgetCell that was selected when reloading the sheet on
rotation. WidgetsFullSheet is excluded because it does not retain the
open header on rotation.

Bug: 331429554
Test: manual, see screencast
Flag: ACONFIG com.android.launcher3.enable_widget_tap_to_add NEXTFOOD
Change-Id: Id3d21f44b1dc525e144296f513f5a460fc51474c

src/com/android/launcher3/Utilities.java

@@ -70,6 +70,7 @@ import android.util.TypedValue;
 import android.view.MotionEvent;
 import android.view.View;
 import android.view.ViewConfiguration;
+import android.view.ViewGroup;
 import android.view.animation.Interpolator;
 
 import androidx.annotation.ChecksSdkIntAtLeast;
@@ -104,6 +105,7 @@ import java.lang.reflect.Method;
 import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
+import java.util.function.Predicate;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -835,4 +837,27 @@ public final class Utilities {
                 // No-Op
         }
     }
+
+    /**
+     * Does a depth-first search through the View hierarchy starting at root, to find a view that
+     * matches the predicate. Returns null if no View was found. View has a findViewByPredicate
+     * member function but it is currently a @hide API.
+     */
+    @Nullable
+    public static <T extends View> T findViewByPredicate(@NonNull View root,
+            @NonNull Predicate<View> predicate) {
+        if (predicate.test(root)) {
+            return (T) root;
+        }
+        if (root instanceof ViewGroup parent) {
+            int count = parent.getChildCount();
+            for (int i = 0; i < count; i++) {
+                View view = findViewByPredicate(parent.getChildAt(i), predicate);
+                if (view != null) {
+                    return (T) view;
+                }
+            }
+        }
+        return null;
+    }
 }

src/com/android/launcher3/widget/BaseWidgetSheet.java

@@ -42,6 +42,7 @@ import com.android.launcher3.Insettable;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.PendingAddItemInfo;
 import com.android.launcher3.R;
+import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.popup.PopupDataProvider;
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.shared.TestProtocol;
@@ -74,6 +75,7 @@ public abstract class BaseWidgetSheet extends AbstractSlideInView<BaseActivity>
     private boolean mDisableNavBarScrim = false;
 
     @Nullable private WidgetCell mWidgetCellWithAddButton = null;
+    @Nullable private WidgetItem mLastSelectedWidgetItem = null;
 
     public BaseWidgetSheet(Context context, AttributeSet attrs, int defStyleAttr) {
         super(context, attrs, defStyleAttr);
@@ -161,6 +163,11 @@ public abstract class BaseWidgetSheet extends AbstractSlideInView<BaseActivity>
             }
 
             mWidgetCellWithAddButton = mWidgetCellWithAddButton != wc ? wc : null;
+            if (mWidgetCellWithAddButton != null) {
+                mLastSelectedWidgetItem = mWidgetCellWithAddButton.getWidgetItem();
+            } else {
+                mLastSelectedWidgetItem = null;
+            }
         } else {
             mActivityContext.getItemOnClickListener().onClick(wc);
         }
@@ -236,6 +243,14 @@ public abstract class BaseWidgetSheet extends AbstractSlideInView<BaseActivity>
         return 0;
     }
 
+    /**
+     * Returns the component of the widget that is currently showing an add button, if any.
+     */
+    @Nullable
+    protected WidgetItem getLastSelectedWidgetItem() {
+        return mLastSelectedWidgetItem;
+    }
+
     @Override
     public boolean onLongClick(View v) {
         TestLogging.recordEvent(TestProtocol.SEQUENCE_MAIN, "Widgets.onLongClick");

src/com/android/launcher3/widget/WidgetCell.java

@@ -503,6 +503,15 @@ public class WidgetCell extends LinearLayout {
         super.onMeasure(widthMeasureSpec, heightMeasureSpec);
     }
 
+    @Override
+    protected void onLayout(boolean changed, int l, int t, int r, int b) {
+        super.onLayout(changed, l, t, r, b);
+
+        if (changed && isShowingAddButton()) {
+            post(this::setupIconOrTextButton);
+        }
+    }
+
     /**
      * Loads a high resolution package icon to show next to the widget title.
      */
@@ -627,4 +636,19 @@ public class WidgetCell extends LinearLayout {
         set.playSequentially(hideAnim, showAnim);
         set.start();
     }
+
+    /**
+     * Returns true if this WidgetCell is displaying the same item as info.
+     */
+    public boolean matchesItem(WidgetItem info) {
+        if (info == null || mItem == null) return false;
+        if (info.widgetInfo != null && mItem.widgetInfo != null) {
+            return info.widgetInfo.getUser().equals(mItem.widgetInfo.getUser())
+                    && info.widgetInfo.getComponent().equals(mItem.widgetInfo.getComponent());
+        } else if (info.activityInfo != null && mItem.activityInfo != null) {
+            return info.activityInfo.getUser().equals(mItem.activityInfo.getUser())
+                    && info.activityInfo.getComponent().equals(mItem.activityInfo.getComponent());
+        }
+        return false;
+    }
 }

src/com/android/launcher3/widget/WidgetsBottomSheet.java

@@ -142,6 +142,9 @@ public class WidgetsBottomSheet extends BaseWidgetSheet {
                     row.forEach(widgetItem -> {
                         WidgetCell widget = addItemCell(tableRow);
                         widget.applyFromCellItem(widgetItem);
+                        if (widget.matchesItem(getLastSelectedWidgetItem())) {
+                            widget.callOnClick();
+                        }
                     });
                     widgetsTable.addView(tableRow);
                 });

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

@@ -29,6 +29,7 @@ import android.util.AttributeSet;
 import android.view.LayoutInflater;
 import android.view.MotionEvent;
 import android.view.View;
+import android.view.ViewGroup;
 import android.view.ViewParent;
 import android.widget.FrameLayout;
 import android.widget.LinearLayout;
@@ -40,6 +41,7 @@ import androidx.annotation.Px;
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
+import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.recyclerview.ViewHolderBinder;
 import com.android.launcher3.util.PackageUserKey;
@@ -194,15 +196,21 @@ public class WidgetsTwoPaneSheet extends WidgetsFullSheet {
                 layoutParams.width = 0;
             }
             layoutParams.weight = layoutParams.width == 0 ? 0.33F : 0;
-            leftPane.setLayoutParams(layoutParams);
-            requestApplyInsets();
-            if (mSelectedHeader != null) {
-                if (mSelectedHeader.equals(mSuggestedWidgetsPackageUserKey)) {
-                    mSuggestedWidgetsHeader.callOnClick();
-                } else {
-                    getHeaderChangeListener().onHeaderChanged(mSelectedHeader);
+
+            post(() -> {
+                // The following calls all trigger requestLayout, so we post them to avoid
+                // calling requestLayout during a layout pass. This also fixes the related warnings
+                // in logcat.
+                leftPane.setLayoutParams(layoutParams);
+                requestApplyInsets();
+                if (mSelectedHeader != null) {
+                    if (mSelectedHeader.equals(mSuggestedWidgetsPackageUserKey)) {
+                        mSuggestedWidgetsHeader.callOnClick();
+                    } else {
+                        getHeaderChangeListener().onHeaderChanged(mSelectedHeader);
+                    }
                 }
-            }
+            });
         }
     }
 
@@ -222,6 +230,9 @@ public class WidgetsTwoPaneSheet extends WidgetsFullSheet {
         if (mSuggestedWidgetsContainer == null && mRecommendedWidgetsCount > 0) {
             setupSuggestedWidgets(LayoutInflater.from(getContext()));
             mSuggestedWidgetsHeader.callOnClick();
+        } else if (mSelectedHeader.equals(mSuggestedWidgetsPackageUserKey)) {
+            // Reselect widget if we are reloading recommendations while it is currently showing.
+            selectWidgetCell(mWidgetRecommendationsContainer, getLastSelectedWidgetItem());
         }
     }
 
@@ -269,6 +280,16 @@ public class WidgetsTwoPaneSheet extends WidgetsFullSheet {
             mRightPaneScrollView.setScrollY(0);
             mRightPane.setAccessibilityPaneTitle(suggestionsRightPaneTitle);
             mSuggestedWidgetsPackageUserKey = PackageUserKey.fromPackageItemInfo(packageItemInfo);
+            final boolean isChangingHeaders =
+                    !mSelectedHeader.equals(mSuggestedWidgetsPackageUserKey);
+            if (isChangingHeaders)  {
+                // If switching from another header, unselect any WidgetCells. This is necessary
+                // because we do not clear/recycle the WidgetCells in the recommendations container
+                // when the header is clicked, only when onRecommendationsBound is called. That
+                // means a WidgetCell in the recommendations container may still be selected from
+                // the last time the recommendations were shown.
+                unselectWidgetCell(mWidgetRecommendationsContainer, getLastSelectedWidgetItem());
+            }
             mSelectedHeader = mSuggestedWidgetsPackageUserKey;
         });
         mSuggestedWidgetsContainer.addView(mSuggestedWidgetsHeader);
@@ -357,6 +378,8 @@ public class WidgetsTwoPaneSheet extends WidgetsFullSheet {
         return new HeaderChangeListener() {
             @Override
             public void onHeaderChanged(@NonNull PackageUserKey selectedHeader) {
+                final boolean isSameHeader = mSelectedHeader != null
+                        && mSelectedHeader.equals(selectedHeader);
                 mSelectedHeader = selectedHeader;
                 WidgetsListContentEntry contentEntry = mActivityContext.getPopupDataProvider()
                         .getSelectedAppWidgets(selectedHeader);
@@ -384,11 +407,20 @@ public class WidgetsTwoPaneSheet extends WidgetsFullSheet {
                         contentEntryToBind,
                         ViewHolderBinder.POSITION_FIRST | ViewHolderBinder.POSITION_LAST,
                         Collections.EMPTY_LIST);
+                if (isSameHeader) {
+                    // Reselect the last selected widget if we are reloading the same header.
+                    selectWidgetCell(widgetsRowViewHolder.tableContainer,
+                            getLastSelectedWidgetItem());
+                }
                 widgetsRowViewHolder.mDataCallback = data -> {
                     mWidgetsListTableViewHolderBinder.bindViewHolder(widgetsRowViewHolder,
                             contentEntryToBind,
                             ViewHolderBinder.POSITION_FIRST | ViewHolderBinder.POSITION_LAST,
                             Collections.singletonList(data));
+                    if (isSameHeader) {
+                        selectWidgetCell(widgetsRowViewHolder.tableContainer,
+                                getLastSelectedWidgetItem());
+                    }
                 };
                 mRightPane.removeAllViews();
                 mRightPane.addView(widgetsRowViewHolder.itemView);
@@ -401,6 +433,24 @@ public class WidgetsTwoPaneSheet extends WidgetsFullSheet {
         };
     }
 
+    private static void selectWidgetCell(ViewGroup parent, WidgetItem item) {
+        if (parent == null || item == null) return;
+        WidgetCell cell = Utilities.findViewByPredicate(parent, v -> v instanceof WidgetCell wc
+                && wc.matchesItem(item));
+        if (cell != null && !cell.isShowingAddButton()) {
+            cell.callOnClick();
+        }
+    }
+
+    private static void unselectWidgetCell(ViewGroup parent, WidgetItem item) {
+        if (parent == null || item == null) return;
+        WidgetCell cell = Utilities.findViewByPredicate(parent, v -> v instanceof WidgetCell wc
+                && wc.matchesItem(item));
+        if (cell != null && cell.isShowingAddButton()) {
+            cell.hideAddButton(/* animate= */ false);
+        }
+    }
+
     @Override
     public void setInsets(Rect insets) {
         super.setInsets(insets);