Update references that read popup data provider to use picker provider

Covers the UI references related to widget picker or its shortcut.

Test: Existing tests
Flag: EXEMPT minor refactoring
Bug: 353347512
Change-Id: Ifed1b7600ca9836df147949f3740ca1e89acd04a
diff --git a/src/com/android/launcher3/popup/SystemShortcut.java b/src/com/android/launcher3/popup/SystemShortcut.java
index 83e9810..f7e1168 100644
--- a/src/com/android/launcher3/popup/SystemShortcut.java
+++ b/src/com/android/launcher3/popup/SystemShortcut.java
@@ -5,6 +5,7 @@
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_APP_INFO_TAP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_DONT_SUGGEST_APP_TAP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_WIDGETS_TAP;
+import static com.android.launcher3.widget.picker.model.data.WidgetPickerDataUtils.findAllWidgetsForPackageUser;
 
 import android.app.ActivityOptions;
 import android.content.ComponentName;
@@ -28,7 +29,6 @@
 import com.android.launcher3.SecondaryDropTarget;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.allapps.PrivateProfileManager;
-import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.pm.UserCache;
@@ -39,9 +39,9 @@
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.widget.WidgetsBottomSheet;
+import com.android.launcher3.widget.picker.model.data.WidgetPickerData;
 
 import java.util.Arrays;
-import java.util.List;
 
 /**
  * Represents a system shortcut for a given app. The shortcut should have a label and icon, and an
@@ -107,11 +107,12 @@
     }
 
     public static final Factory<ActivityContext> WIDGETS = (context, itemInfo, originalView) -> {
-        if (itemInfo.getTargetComponent() == null) return null;
-        final List<WidgetItem> widgets =
-                context.getPopupDataProvider().getWidgetsForPackageUser(new PackageUserKey(
-                        itemInfo.getTargetComponent().getPackageName(), itemInfo.user));
-        if (widgets.isEmpty()) {
+        final PackageUserKey packageUserKey = PackageUserKey.fromItemInfo(itemInfo);
+        if (packageUserKey == null) return null;
+
+        final WidgetPickerData data = context.getWidgetPickerDataProvider().get();
+        if (findAllWidgetsForPackageUser(data, packageUserKey).isEmpty()) {
+            // hides widget picker shortcut if there are no widgets for the package.
             return null;
         }
         return new Widgets(context, itemInfo, originalView);
diff --git a/src/com/android/launcher3/widget/WidgetsBottomSheet.java b/src/com/android/launcher3/widget/WidgetsBottomSheet.java
index 57c065a..ddbd291 100644
--- a/src/com/android/launcher3/widget/WidgetsBottomSheet.java
+++ b/src/com/android/launcher3/widget/WidgetsBottomSheet.java
@@ -17,6 +17,7 @@
 package com.android.launcher3.widget;
 
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_BOTTOM_WIDGETS_TRAY;
+import static com.android.launcher3.widget.picker.model.data.WidgetPickerDataUtils.findAllWidgetsForPackageUser;
 
 import android.content.Context;
 import android.graphics.Rect;
@@ -40,6 +41,7 @@
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.widget.picker.model.data.WidgetPickerData;
 import com.android.launcher3.widget.util.WidgetsTableUtils;
 
 import java.util.List;
@@ -124,10 +126,10 @@
 
     @Override
     public void onWidgetsBound() {
-        List<WidgetItem> widgets = mActivityContext.getPopupDataProvider().getWidgetsForPackageUser(
-                new PackageUserKey(
-                        mOriginalItemInfo.getTargetComponent().getPackageName(),
-                        mOriginalItemInfo.user));
+        final WidgetPickerData data = mActivityContext.getWidgetPickerDataProvider().get();
+        final PackageUserKey packageUserKey = PackageUserKey.fromItemInfo(mOriginalItemInfo);
+        List<WidgetItem> widgets = packageUserKey != null ? findAllWidgetsForPackageUser(data,
+                packageUserKey) : List.of();
 
         TableLayout widgetsTable = findViewById(R.id.widgets_table);
         widgetsTable.removeAllViews();
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 21b7be4..2af8e6f 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -74,6 +74,7 @@
 import com.android.launcher3.workprofile.PersonalWorkSlidingTabStrip.OnActivePageChangedListener;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -471,7 +472,7 @@
      * Returns all displayable widgets.
      */
     protected List<WidgetsListBaseEntry> getWidgetsToDisplay() {
-        return mActivityContext.getPopupDataProvider().getAllWidgets();
+        return mActivityContext.getWidgetPickerDataProvider().get().getAllWidgets();
     }
 
     @Override
@@ -572,12 +573,11 @@
         if (mIsInSearchMode) {
             return;
         }
-
         if (enableCategorizedWidgetSuggestions()) {
             // We avoid applying new recommendations when some are already displayed.
             if (mRecommendedWidgetsMap.isEmpty()) {
                 mRecommendedWidgetsMap =
-                        mActivityContext.getPopupDataProvider().getCategorizedRecommendedWidgets();
+                        mActivityContext.getWidgetPickerDataProvider().get().getRecommendations();
             }
             mRecommendedWidgetsCount = mWidgetRecommendationsView.setRecommendations(
                     mRecommendedWidgetsMap,
@@ -589,17 +589,20 @@
             );
         } else {
             if (mRecommendedWidgets.isEmpty()) {
-                mRecommendedWidgets =
-                        mActivityContext.getPopupDataProvider().getRecommendedWidgets();
+                mRecommendedWidgets = mActivityContext.getWidgetPickerDataProvider().get()
+                        .getRecommendations()
+                        .values().stream()
+                        .flatMap(Collection::stream).toList();
+                mRecommendedWidgetsCount = mWidgetRecommendationsView.setRecommendations(
+                        mRecommendedWidgets,
+                        mDeviceProfile,
+                        /* availableHeight= */ getMaxAvailableHeightForRecommendations(),
+                        /* availableWidth= */ mMaxSpanPerRow,
+                        /* cellPadding= */ mWidgetCellHorizontalPadding
+                );
             }
-            mRecommendedWidgetsCount = mWidgetRecommendationsView.setRecommendations(
-                    mRecommendedWidgets,
-                    mDeviceProfile,
-                    /* availableHeight= */ getMaxAvailableHeightForRecommendations(),
-                    /* availableWidth= */ mMaxSpanPerRow,
-                    /* cellPadding= */ mWidgetCellHorizontalPadding
-            );
         }
+
         mWidgetRecommendationsContainer.setVisibility(
                 mRecommendedWidgetsCount > 0 ? VISIBLE : GONE);
     }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
index c4c755a..c2cd903 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
@@ -22,6 +22,7 @@
 import static com.android.launcher3.UtilitiesKt.modifyAttributesOnViewTree;
 import static com.android.launcher3.UtilitiesKt.restoreAttributesOnViewTree;
 import static com.android.launcher3.widget.picker.WidgetsListItemAnimator.WIDGET_LIST_ITEM_APPEARANCE_DELAY;
+import static com.android.launcher3.widget.picker.model.data.WidgetPickerDataUtils.findContentEntryForPackageUser;
 
 import android.content.Context;
 import android.graphics.Rect;
@@ -287,9 +288,9 @@
     @Override
     protected List<WidgetsListBaseEntry> getWidgetsToDisplay() {
         List<WidgetsListBaseEntry> allWidgets =
-                mActivityContext.getPopupDataProvider().getAllWidgets();
+                mActivityContext.getWidgetPickerDataProvider().get().getAllWidgets();
         List<WidgetsListBaseEntry> defaultWidgets =
-                mActivityContext.getPopupDataProvider().getDefaultWidgets();
+                mActivityContext.getWidgetPickerDataProvider().get().getDefaultWidgets();
 
         if (allWidgets.isEmpty() || defaultWidgets.isEmpty()) {
             // no menu if there are no default widgets to show
@@ -359,7 +360,7 @@
         WidgetsListHeaderEntry widgetsListHeaderEntry = WidgetsListHeaderEntry.create(
                         packageItemInfo,
                         /*titleSectionName=*/ suggestionsHeaderTitle,
-                        /*items=*/ mActivityContext.getPopupDataProvider().getRecommendedWidgets(),
+                        /*items=*/ List.of(), // not necessary
                         /*visibleWidgetsCount=*/ 0)
                 .withWidgetListShown();
 
@@ -509,11 +510,11 @@
                 final boolean isUserClick = mSelectedHeader != null
                         && !getAccessibilityInitialFocusView().isAccessibilityFocused();
                 mSelectedHeader = selectedHeader;
-                WidgetsListContentEntry contentEntry =
-                        mActivityContext.getPopupDataProvider().getSelectedAppWidgets(
-                                selectedHeader, /*useDefault=*/
-                                (mWidgetOptionsMenuState != null
-                                        && !mWidgetOptionsMenuState.showAllWidgets));
+                final boolean showDefaultWidgets = mWidgetOptionsMenuState != null
+                        && !mWidgetOptionsMenuState.showAllWidgets;
+                WidgetsListContentEntry contentEntry = findContentEntryForPackageUser(
+                        mActivityContext.getWidgetPickerDataProvider().get(),
+                        selectedHeader, showDefaultWidgets);
 
                 if (contentEntry == null || mRightPane == null) {
                     return;
diff --git a/tests/src/com/android/launcher3/popup/SystemShortcutTest.java b/tests/src/com/android/launcher3/popup/SystemShortcutTest.java
index 98b6b4b..dcfcad5 100644
--- a/tests/src/com/android/launcher3/popup/SystemShortcutTest.java
+++ b/tests/src/com/android/launcher3/popup/SystemShortcutTest.java
@@ -63,6 +63,8 @@
 import com.android.launcher3.util.TestSandboxModelContextWrapper;
 import com.android.launcher3.util.UserIconInfo;
 import com.android.launcher3.views.BaseDragLayer;
+import com.android.launcher3.widget.picker.model.WidgetPickerDataProvider;
+import com.android.launcher3.widget.picker.model.data.WidgetPickerData;
 
 import org.junit.After;
 import org.junit.Assert;
@@ -73,8 +75,6 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.ArrayList;
-
 @SmallTest
 @RunWith(LauncherMultivalentJUnit.class)
 public class SystemShortcutTest {
@@ -86,7 +86,7 @@
     private TestSandboxModelContextWrapper mTestContext;
     private final SandboxModelContext mSandboxContext = new SandboxModelContext();
     private PrivateProfileManager mPrivateProfileManager;
-    private PopupDataProvider mPopupDataProvider;
+    private WidgetPickerDataProvider mWidgetPickerDataProvider;
     private AppInfo mAppInfo;
     @Mock UserCache mUserCache;
     @Mock ApiWrapper mApiWrapper;
@@ -119,8 +119,8 @@
         spyOn(mPrivateProfileManager);
         when(mPrivateProfileManager.getProfileUser()).thenReturn(PRIVATE_HANDLE);
 
-        mPopupDataProvider = mTestContext.getPopupDataProvider();
-        spyOn(mPopupDataProvider);
+        mWidgetPickerDataProvider = mTestContext.getWidgetPickerDataProvider();
+        spyOn(mWidgetPickerDataProvider);
     }
 
     @After
@@ -141,7 +141,7 @@
         mAppInfo = new AppInfo();
         mAppInfo.componentName = new ComponentName(mTestContext, getClass());
         assertNotNull(mAppInfo.getTargetComponent());
-        doReturn(new ArrayList<>()).when(mPopupDataProvider).getWidgetsForPackageUser(any());
+        doReturn(new WidgetPickerData()).when(mWidgetPickerDataProvider).get();
         spyOn(mAppInfo);
         SystemShortcut systemShortcut = SystemShortcut.WIDGETS
                 .getShortcut(mTestContext, mAppInfo, mView);