Update widget predictor to apply prediction filter

When enough widgets are not passing the filter, additional randomly
 selected widgets are added.

The count to decide whether to add more is a configuration, so that,
if some OEMs don't want any suggestions, can override the value to 0.

Bug: 356127021
Flag: com.android.launcher3.enable_tiered_widgets_by_default_in_picker
Test: Unit tests
Change-Id: Iffa8619149a1a4b468d367fc7bbee381be59469d
diff --git a/quickstep/res/values/config.xml b/quickstep/res/values/config.xml
index 5c80575..f3c9467 100644
--- a/quickstep/res/values/config.xml
+++ b/quickstep/res/values/config.xml
@@ -52,6 +52,11 @@
 
     <integer name="max_depth_blur_radius">23</integer>
 
+    <!-- If predicted widgets from prediction service are less than this number, additional
+    eligible widgets may be added locally by launcher. When set to 0, no widgets will be added
+    locally. -->
+    <integer name="widget_predictions_min_count">6</integer>
+
     <!-- Accessibility actions -->
     <item type="id" name="action_move_to_top_or_left" />
     <item type="id" name="action_move_to_bottom_or_right" />
diff --git a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
index 0395d32..9d9054e 100644
--- a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
+++ b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
@@ -16,8 +16,12 @@
 package com.android.launcher3.model;
 
 import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
+import static com.android.launcher3.Flags.enableTieredWidgetsByDefaultInPicker;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
 
+import static java.util.stream.Collectors.groupingBy;
+import static java.util.stream.Collectors.toMap;
+
 import android.app.prediction.AppTarget;
 import android.content.ComponentName;
 import android.content.Context;
@@ -26,6 +30,7 @@
 import androidx.annotation.NonNull;
 
 import com.android.launcher3.LauncherModel.ModelUpdateTask;
+import com.android.launcher3.R;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.QuickstepModelDelegate.PredictorState;
 import com.android.launcher3.model.data.ItemInfo;
@@ -34,8 +39,10 @@
 import com.android.launcher3.widget.picker.WidgetRecommendationCategoryProvider;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Random;
 import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -60,33 +67,72 @@
     @Override
     public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
             @NonNull AllAppsList apps) {
+        Predicate<WidgetItem> predictedWidgetsFilter = enableTieredWidgetsByDefaultInPicker()
+                ? dataModel.widgetsModel.getPredictedWidgetsFilter() : null;
         Set<ComponentKey> widgetsInWorkspace = dataModel.appWidgets.stream().map(
                 widget -> new ComponentKey(widget.providerName, widget.user)).collect(
                 Collectors.toSet());
-        Predicate<WidgetItem> notOnWorkspace = w -> !widgetsInWorkspace.contains(w);
-        Map<ComponentKey, WidgetItem> allWidgets =
-                dataModel.widgetsModel.getWidgetsByComponentKey();
+
+        // Widgets (excluding shortcuts & already added widgets) that belong to apps eligible for
+        // being in predictions.
+        Map<ComponentKey, WidgetItem> allEligibleWidgets =
+                dataModel.widgetsModel.getWidgetsByComponentKey()
+                        .entrySet()
+                        .stream()
+                        .filter(entry -> entry.getValue().widgetInfo != null
+                                && !widgetsInWorkspace.contains(entry.getValue())
+                                && (predictedWidgetsFilter == null
+                                || predictedWidgetsFilter.test(entry.getValue()))
+                        ).collect(toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+        Context context = taskController.getApp().getContext();
 
         List<WidgetItem> servicePredictedItems = new ArrayList<>();
+        List<String> addedWidgetApps = new ArrayList<>();
 
         for (AppTarget app : mTargets) {
             ComponentKey componentKey = new ComponentKey(
                     new ComponentName(app.getPackageName(), app.getClassName()), app.getUser());
-            WidgetItem widget = allWidgets.get(componentKey);
-            if (widget == null) {
+            WidgetItem widget = allEligibleWidgets.get(componentKey);
+            if (widget == null) { // widget not eligible.
                 continue;
             }
             String className = app.getClassName();
             if (!TextUtils.isEmpty(className)) {
-                if (notOnWorkspace.test(widget)) {
-                    servicePredictedItems.add(widget);
-                }
+                servicePredictedItems.add(widget);
+                addedWidgetApps.add(componentKey.componentName.getPackageName());
+            }
+        }
+
+        int minPredictionCount = context.getResources().getInteger(
+                R.integer.widget_predictions_min_count);
+        if (enableTieredWidgetsByDefaultInPicker()
+                && servicePredictedItems.size() < minPredictionCount) {
+            // Eligible apps that aren't already part of predictions.
+            Map<String, List<WidgetItem>> eligibleWidgetsByApp =
+                    allEligibleWidgets.values().stream()
+                            .filter(w -> !addedWidgetApps.contains(
+                                    w.componentName.getPackageName()))
+                            .collect(groupingBy(w -> w.componentName.getPackageName()));
+
+            // Randomize available apps list
+            List<String> appPackages = new ArrayList<>(eligibleWidgetsByApp.keySet());
+            Collections.shuffle(appPackages);
+
+            int widgetsToAdd = minPredictionCount - servicePredictedItems.size();
+            for (String appPackage : appPackages) {
+                if (widgetsToAdd <= 0) break;
+
+                List<WidgetItem> widgetsForApp = eligibleWidgetsByApp.get(appPackage);
+                int index = new Random().nextInt(widgetsForApp.size());
+                // Add a random widget from the app.
+                servicePredictedItems.add(widgetsForApp.get(index));
+                widgetsToAdd--;
             }
         }
 
         List<ItemInfo> items;
         if (enableCategorizedWidgetSuggestions()) {
-            Context context = taskController.getApp().getContext();
             WidgetRecommendationCategoryProvider categoryProvider =
                     WidgetRecommendationCategoryProvider.newInstance(context);
             items = servicePredictedItems.stream()
diff --git a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
index 7b57c81..c53c177 100644
--- a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
+++ b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -33,6 +33,7 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 
 import android.app.prediction.AppTarget;
 import android.app.prediction.AppTargetId;
@@ -42,6 +43,8 @@
 import android.content.pm.ApplicationInfo;
 import android.content.pm.LauncherApps;
 import android.os.UserHandle;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.text.TextUtils;
 
@@ -62,9 +65,13 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
 
 import java.util.Arrays;
 import java.util.List;
+import java.util.Set;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
 @SmallTest
@@ -72,6 +79,9 @@
 public final class WidgetsPredicationUpdateTaskTest {
 
     @Rule
+    public final MockitoRule mocks = MockitoJUnit.rule();
+
+    @Rule
     public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
     private AppWidgetProviderInfo mApp1Provider1;
@@ -145,6 +155,7 @@
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_TIERED_WIDGETS_BY_DEFAULT_IN_PICKER) // Flag off
     public void widgetsRecommendationRan_shouldOnlyReturnNotAddedWidgetsInAppPredictionOrder() {
         // Run on model executor so that no other task runs in the middle.
         runOnExecutorSync(MODEL_EXECUTOR, () -> {
@@ -184,6 +195,7 @@
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_TIERED_WIDGETS_BY_DEFAULT_IN_PICKER) // Flag off
     public void widgetsRecommendationRan_shouldReturnEmptyWidgetsWhenEmpty() {
         runOnExecutorSync(MODEL_EXECUTOR, () -> {
 
@@ -213,6 +225,50 @@
         });
     }
 
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_TIERED_WIDGETS_BY_DEFAULT_IN_PICKER)
+    public void widgetsRecommendationRan_keepsWidgetsNotOnWorkspace_addsWidgetsFromEligibleApps() {
+        runOnExecutorSync(MODEL_EXECUTOR, () -> {
+            WidgetsFilterDataProvider spiedFilterProvider = spy(
+                    mModelHelper.getModel().getWidgetsFilterDataProvider());
+            doAnswer(i -> new Predicate<WidgetItem>() {
+                @Override
+                public boolean test(WidgetItem widgetItem) {
+                    // app5's widget is already on workspace, but, app2 is not.
+                    // And app4's second widget is also not on workspace.
+                    return Set.of("app5", "app2", "app4").contains(
+                            widgetItem.componentName.getPackageName());
+                }
+            }).when(spiedFilterProvider).getPredictedWidgetsFilter();
+            mModelHelper.getBgDataModel().widgetsModel.updateWidgetFilters(spiedFilterProvider);
+            // App5's widget that's already on workspace.
+            AppTarget widget1 = new AppTarget(new AppTargetId("app5"), "app5", "provider1",
+                    mUserHandle);
+            // App4's widget eligible and not on workspace.
+            AppTarget widget2 = new AppTarget(new AppTargetId("app4"), "app4", "provider2",
+                    mUserHandle);
+
+            mCallback.mRecommendedWidgets = null;
+            mModelHelper.getModel().enqueueModelUpdateTask(
+                    newWidgetsPredicationTask(List.of(widget1, widget2)));
+            runOnExecutorSync(MAIN_EXECUTOR, () -> {
+            });
+
+            List<PendingAddWidgetInfo> recommendedWidgets = mCallback.mRecommendedWidgets.items
+                    .stream()
+                    .map(itemInfo -> (PendingAddWidgetInfo) itemInfo)
+                    .collect(Collectors.toList());
+            assertThat(recommendedWidgets).hasSize(2);
+            List<ComponentName> componentNames = recommendedWidgets.stream().map(
+                    w -> w.componentName).toList();
+            assertThat(componentNames).containsExactly(
+                    // Locally added, not on workspace, eligible app per filter
+                    mApp2Provider1.provider,
+                    // From prediction service, not on workspace, eligible app per filter
+                    mApp4Provider2.provider);
+        });
+    }
+
     private void assertWidgetInfo(
             LauncherAppWidgetProviderInfo actual, AppWidgetProviderInfo expected) {
         assertThat(actual.provider).isEqualTo(expected.provider);