Consider WidgetSections for suggestions in standalone picker

This cl address the problem for standalone picker (follow up to match
ag/27720721) to ensure widgets that aren't in section of their owning
package didn't appear in suggestions

Bug: 345520128
Test: Unit tests
Flag: EXEMPT bugfix
Change-Id: Ia0ef96c5be77db56b84c76ace498125d07f4be42

quickstep/src/com/android/launcher3/WidgetPickerActivity.java

@@ -43,17 +43,18 @@ import com.android.launcher3.model.WidgetPredictionsRequester;
 import com.android.launcher3.model.WidgetsModel;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.popup.PopupDataProvider;
-import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.widget.BaseWidgetSheet;
 import com.android.launcher3.widget.WidgetCell;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
-import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.model.WidgetsListContentEntry;
 import com.android.launcher3.widget.picker.WidgetsFullSheet;
 
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.function.Function;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
@@ -240,14 +241,16 @@ public class WidgetPickerActivity extends BaseActivity {
                     );
             bindWidgets(allWidgets);
             if (mUiSurface != null) {
-                Map<PackageUserKey, List<WidgetItem>> allWidgetsMap = allWidgets.stream()
-                        .filter(WidgetsListHeaderEntry.class::isInstance)
+                Map<ComponentKey, WidgetItem> allWidgetItems = allWidgets.stream()
+                        .filter(entry -> entry instanceof WidgetsListContentEntry)
+                        .flatMap(entry -> entry.mWidgets.stream())
+                        .distinct()
                         .collect(Collectors.toMap(
-                                entry -> PackageUserKey.fromPackageItemInfo(entry.mPkgItem),
-                                entry -> entry.mWidgets)
-                        );
+                                widget -> new ComponentKey(widget.componentName, widget.user),
+                                Function.identity()
+                        ));
                 mWidgetPredictionsRequester = new WidgetPredictionsRequester(app.getContext(),
-                        mUiSurface, allWidgetsMap);
+                        mUiSurface, allWidgetItems);
                 mWidgetPredictionsRequester.request(mAddedWidgets, this::bindRecommendedWidgets);
             }
         });

quickstep/src/com/android/launcher3/model/WidgetPredictionsRequester.java

@@ -40,7 +40,6 @@ import androidx.annotation.WorkerThread;
 
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.ComponentKey;
-import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.widget.PendingAddWidgetInfo;
 import com.android.launcher3.widget.picker.WidgetRecommendationCategoryProvider;
 
@@ -67,10 +66,10 @@ public class WidgetPredictionsRequester {
     @NonNull
     private final String mUiSurface;
     @NonNull
-    private final Map<PackageUserKey, List<WidgetItem>> mAllWidgets;
+    private final Map<ComponentKey, WidgetItem> mAllWidgets;
 
     public WidgetPredictionsRequester(Context context, @NonNull String uiSurface,
-            @NonNull Map<PackageUserKey, List<WidgetItem>> allWidgets) {
+            @NonNull Map<ComponentKey, WidgetItem> allWidgets) {
         mContext = context;
         mUiSurface = uiSurface;
         mAllWidgets = Collections.unmodifiableMap(allWidgets);
@@ -172,33 +171,19 @@ public class WidgetPredictionsRequester {
      */
     @VisibleForTesting
     static List<WidgetItem> filterPredictions(List<AppTarget> predictions,
-            Map<PackageUserKey, List<WidgetItem>> allWidgets, Predicate<WidgetItem> filter) {
+            Map<ComponentKey, WidgetItem> allWidgets, Predicate<WidgetItem> filter) {
         List<WidgetItem> servicePredictedItems = new ArrayList<>();
-        List<WidgetItem> localFilteredWidgets = new ArrayList<>();
 
         for (AppTarget prediction : predictions) {
-            List<WidgetItem> widgetsInPackage = allWidgets.get(
-                    new PackageUserKey(prediction.getPackageName(), prediction.getUser()));
-            if (widgetsInPackage == null || widgetsInPackage.isEmpty()) {
-                continue;
-            }
             String className = prediction.getClassName();
             if (!TextUtils.isEmpty(className)) {
-                WidgetItem item = widgetsInPackage.stream()
-                        .filter(w -> className.equals(w.componentName.getClassName()))
-                        .filter(filter)
-                        .findFirst().orElse(null);
-                if (item != null) {
-                    servicePredictedItems.add(item);
-                    continue;
+                WidgetItem widgetItem = allWidgets.get(
+                        new ComponentKey(new ComponentName(prediction.getPackageName(), className),
+                                prediction.getUser()));
+                if (widgetItem != null && filter.test(widgetItem)) {
+                    servicePredictedItems.add(widgetItem);
                 }
             }
-            // No widget was added by the service, try local filtering
-            widgetsInPackage.stream().filter(filter).findFirst()
-                    .ifPresent(localFilteredWidgets::add);
-        }
-        if (servicePredictedItems.isEmpty()) {
-            servicePredictedItems.addAll(localFilteredWidgets);
         }
 
         return servicePredictedItems;

quickstep/tests/src/com/android/launcher3/model/WidgetsPredictionsRequesterTest.kt

@@ -34,7 +34,7 @@ import com.android.launcher3.model.WidgetPredictionsRequester.buildBundleForPred
 import com.android.launcher3.model.WidgetPredictionsRequester.filterPredictions
 import com.android.launcher3.model.WidgetPredictionsRequester.notOnUiSurfaceFilter
 import com.android.launcher3.util.ActivityContextWrapper
-import com.android.launcher3.util.PackageUserKey
+import com.android.launcher3.util.ComponentKey
 import com.android.launcher3.util.WidgetUtils.createAppWidgetProviderInfo
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo
 import com.google.common.truth.Truth.assertThat
@@ -62,7 +62,7 @@ class WidgetsPredictionsRequesterTest {
     private lateinit var widgetItem1b: WidgetItem
     private lateinit var widgetItem2: WidgetItem
 
-    private lateinit var allWidgets: Map<PackageUserKey, List<WidgetItem>>
+    private lateinit var allWidgets: Map<ComponentKey, WidgetItem>
 
     @Mock private lateinit var iconCache: IconCache
 
@@ -93,9 +93,9 @@ class WidgetsPredictionsRequesterTest {
 
         allWidgets =
             mapOf(
-                PackageUserKey(APP_1_PACKAGE_NAME, mUserHandle) to
-                    listOf(widgetItem1a, widgetItem1b),
-                PackageUserKey(APP_2_PACKAGE_NAME, mUserHandle) to listOf(widgetItem2),
+                ComponentKey(widgetItem1a.componentName, widgetItem1a.user) to widgetItem1a,
+                ComponentKey(widgetItem1b.componentName, widgetItem1b.user) to widgetItem1b,
+                ComponentKey(widgetItem2.componentName, widgetItem2.user) to widgetItem2,
             )
     }
 
@@ -156,7 +156,7 @@ class WidgetsPredictionsRequesterTest {
     }
 
     @Test
-    fun filterPredictions_appPredictions_returnsWidgetFromPackage() {
+    fun filterPredictions_appPredictions_returnsEmptyList() {
         val widgetsAlreadyOnSurface = arrayListOf(widget1bInfo)
         val filter: Predicate<WidgetItem> = notOnUiSurfaceFilter(widgetsAlreadyOnSurface)
 
@@ -176,8 +176,7 @@ class WidgetsPredictionsRequesterTest {
                 ),
             )
 
-        assertThat(filterPredictions(predictions, allWidgets, filter))
-            .containsExactly(widgetItem1a, widgetItem2)
+        assertThat(filterPredictions(predictions, allWidgets, filter)).isEmpty()
     }
 
     private fun createWidgetItem(