Merge "Updating widget prediction handling:" into tm-qpr-dev
diff --git a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
index 7a483a8..1beabf1 100644
--- a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
+++ b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
@@ -18,22 +18,22 @@
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
 
 import android.app.prediction.AppTarget;
-import android.content.ComponentName;
 import android.text.TextUtils;
 
 import androidx.annotation.NonNull;
 
 import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.QuickstepModelDelegate.PredictorState;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.widget.PendingAddWidgetInfo;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
 /** Task to update model as a result of predicted widgets update */
@@ -59,50 +59,43 @@
         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<PackageUserKey, List<WidgetItem>> allWidgets =
                 dataModel.widgetsModel.getAllWidgetsWithoutShortcuts();
 
-        FixedContainerItems fixedContainerItems =
-                new FixedContainerItems(mPredictorState.containerId);
+        List<WidgetItem> servicePredictedItems = new ArrayList<>();
+        List<WidgetItem> localFilteredWidgets = new ArrayList<>();
 
-        if (FeatureFlags.ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER.get()) {
-            for (AppTarget app : mTargets) {
-                PackageUserKey packageUserKey = new PackageUserKey(app.getPackageName(),
-                        app.getUser());
-                if (allWidgets.containsKey(packageUserKey)) {
-                    List<WidgetItem> notAddedWidgets = allWidgets.get(packageUserKey).stream()
-                            .filter(item ->
-                                    !widgetsInWorkspace.contains(
-                                            new ComponentKey(item.componentName, item.user)))
-                            .collect(Collectors.toList());
-                    if (notAddedWidgets.size() > 0) {
-                        // Even an apps have more than one widgets, we only include one widget.
-                        fixedContainerItems.items.add(
-                                new PendingAddWidgetInfo(
-                                        notAddedWidgets.get(0).widgetInfo,
-                                        CONTAINER_WIDGETS_PREDICTION));
-                    }
-                }
+        for (AppTarget app : mTargets) {
+            PackageUserKey packageUserKey = new PackageUserKey(app.getPackageName(), app.getUser());
+            List<WidgetItem> widgets = allWidgets.get(packageUserKey);
+            if (widgets == null || widgets.isEmpty()) {
+                continue;
             }
-        } else {
-            Map<ComponentKey, WidgetItem> widgetItems =
-                    allWidgets.values().stream().flatMap(List::stream).distinct()
-                            .collect(Collectors.toMap(widget -> (ComponentKey) widget,
-                                    widget -> widget));
-            for (AppTarget app : mTargets) {
-                if (TextUtils.isEmpty(app.getClassName())) {
+            String className = app.getClassName();
+            if (!TextUtils.isEmpty(className)) {
+                WidgetItem item = widgets.stream()
+                        .filter(w -> className.equals(w.componentName.getClassName()))
+                        .filter(notOnWorkspace)
+                        .findFirst()
+                        .orElse(null);
+                if (item != null) {
+                    servicePredictedItems.add(item);
                     continue;
                 }
-                ComponentKey targetWidget = new ComponentKey(
-                        new ComponentName(app.getPackageName(), app.getClassName()), app.getUser());
-                if (widgetItems.containsKey(targetWidget)) {
-                    fixedContainerItems.items.add(
-                            new PendingAddWidgetInfo(widgetItems.get(
-                                    targetWidget).widgetInfo,
-                                    CONTAINER_WIDGETS_PREDICTION));
-                }
             }
+            // No widget was added by the service, try local filtering
+            widgets.stream().filter(notOnWorkspace).findFirst()
+                    .ifPresent(localFilteredWidgets::add);
         }
+        if (servicePredictedItems.isEmpty()) {
+            servicePredictedItems.addAll(localFilteredWidgets);
+        }
+        FixedContainerItems fixedContainerItems =
+                new FixedContainerItems(mPredictorState.containerId);
+        servicePredictedItems.forEach(w -> fixedContainerItems.items.add(
+                new PendingAddWidgetInfo(w.widgetInfo, CONTAINER_WIDGETS_PREDICTION)));
+
         dataModel.extraItems.put(mPredictorState.containerId, fixedContainerItems);
         bindExtraContainerItems(fixedContainerItems);
 
diff --git a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
index b903691..83341cb 100644
--- a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
+++ b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -41,7 +41,6 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
@@ -136,21 +135,21 @@
     public void widgetsRecommendationRan_shouldOnlyReturnNotAddedWidgetsInAppPredictionOrder()
             throws Exception {
         // WHEN newPredicationTask is executed with app predication of 5 apps.
-        AppTarget app1 = new AppTarget(new AppTargetId("app1"), "app1", "className",
+        AppTarget app1 = new AppTarget(new AppTargetId("app1"), "app1", "provider1",
                 mUserHandle);
-        AppTarget app2 = new AppTarget(new AppTargetId("app2"), "app2", "className",
+        AppTarget app2 = new AppTarget(new AppTargetId("app2"), "app2", "provider1",
                 mUserHandle);
         AppTarget app3 = new AppTarget(new AppTargetId("app3"), "app3", "className",
                 mUserHandle);
-        AppTarget app4 = new AppTarget(new AppTargetId("app4"), "app4", "className",
+        AppTarget app4 = new AppTarget(new AppTargetId("app4"), "app4", "provider1",
                 mUserHandle);
-        AppTarget app5 = new AppTarget(new AppTargetId("app5"), "app5", "className",
+        AppTarget app5 = new AppTarget(new AppTargetId("app5"), "app5", "provider1",
                 mUserHandle);
         mModelHelper.executeTaskForTest(
                 newWidgetsPredicationTask(List.of(app5, app3, app2, app4, app1)))
                 .forEach(Runnable::run);
 
-        // THEN only 3 widgets are returned because
+        // THEN only 2 widgets are returned because
         // 1. app5/provider1 & app4/provider1 have already been added to workspace. They are
         //    excluded from the result.
         // 2. app3 doesn't have a widget.
@@ -159,45 +158,39 @@
                 .stream()
                 .map(itemInfo -> (PendingAddWidgetInfo) itemInfo)
                 .collect(Collectors.toList());
-        assertThat(recommendedWidgets).hasSize(3);
+        assertThat(recommendedWidgets).hasSize(2);
         assertWidgetInfo(recommendedWidgets.get(0).info, mApp2Provider1);
-        assertWidgetInfo(recommendedWidgets.get(1).info, mApp4Provider2);
-        assertWidgetInfo(recommendedWidgets.get(2).info, mApp1Provider1);
+        assertWidgetInfo(recommendedWidgets.get(1).info, mApp1Provider1);
     }
 
     @Test
-    public void widgetsRecommendationRan_localFilterDisabled_shouldReturnWidgetsInPredicationOrder()
+    public void widgetsRecommendationRan_shouldReturnPackageWidgetsWhenEmpty()
             throws Exception {
-        if (FeatureFlags.ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER.get()) {
-            return;
-        }
 
-        // WHEN newPredicationTask is executed with 5 predicated widgets.
-        AppTarget widget1 = new AppTarget(new AppTargetId("app1"), "app1", "provider1",
-                mUserHandle);
-        AppTarget widget2 = new AppTarget(new AppTargetId("app1"), "app1", "provider2",
+        // Not installed widget
+        AppTarget widget1 = new AppTarget(new AppTargetId("app1"), "app1", "provider3",
                 mUserHandle);
         // Not installed app
         AppTarget widget3 = new AppTarget(new AppTargetId("app2"), "app3", "provider1",
                 mUserHandle);
-        // Not installed widget
-        AppTarget widget4 = new AppTarget(new AppTargetId("app4"), "app4", "provider3",
+        // Workspace added widgets
+        AppTarget widget4 = new AppTarget(new AppTargetId("app4"), "app4", "provider1",
                 mUserHandle);
         AppTarget widget5 = new AppTarget(new AppTargetId("app5"), "app5", "provider1",
                 mUserHandle);
         mModelHelper.executeTaskForTest(
-                newWidgetsPredicationTask(List.of(widget5, widget3, widget2, widget4, widget1)))
+                newWidgetsPredicationTask(List.of(widget5, widget3, widget4, widget1)))
                 .forEach(Runnable::run);
 
-        // THEN only 3 widgets are returned because the launcher only filters out non-exist widgets.
+        // THEN only 2 widgets are returned because the launcher only filters out non-exist widgets.
         List<PendingAddWidgetInfo> recommendedWidgets = mCallback.mRecommendedWidgets.items
                 .stream()
                 .map(itemInfo -> (PendingAddWidgetInfo) itemInfo)
                 .collect(Collectors.toList());
-        assertThat(recommendedWidgets).hasSize(3);
-        assertWidgetInfo(recommendedWidgets.get(0).info, mApp5Provider1);
-        assertWidgetInfo(recommendedWidgets.get(1).info, mApp1Provider2);
-        assertWidgetInfo(recommendedWidgets.get(2).info, mApp1Provider1);
+        assertThat(recommendedWidgets).hasSize(2);
+        // Another widget from the same package
+        assertWidgetInfo(recommendedWidgets.get(0).info, mApp4Provider2);
+        assertWidgetInfo(recommendedWidgets.get(1).info, mApp1Provider1);
     }
 
     private void assertWidgetInfo(
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 694536c..32463a5 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -206,10 +206,6 @@
     public static final BooleanFlag ENABLE_ENFORCED_ROUNDED_CORNERS = new DeviceFlag(
             "ENABLE_ENFORCED_ROUNDED_CORNERS", true, "Enforce rounded corners on all App Widgets");
 
-    public static final BooleanFlag ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER = new DeviceFlag(
-            "ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER", true,
-            "Enables a local filter for recommended widgets.");
-
     public static final BooleanFlag NOTIFY_CRASHES = getDebugFlag("NOTIFY_CRASHES", false,
             "Sends a notification whenever launcher encounters an uncaught exception.");