Merge "Send widget added and removed events to AiAi." into sc-v2-dev
diff --git a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionModel.java b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionModel.java
index 080633a..56945ba 100644
--- a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionModel.java
+++ b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionModel.java
@@ -16,33 +16,25 @@
 package com.android.launcher3.hybridhotseat;
 
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION;
+import static com.android.launcher3.model.PredictionHelper.getAppTargetFromItemInfo;
+import static com.android.launcher3.model.PredictionHelper.isTrackedForHotseatPrediction;
+import static com.android.launcher3.model.PredictionHelper.wrapAppTargetWithItemLocation;
 
 import android.app.prediction.AppTarget;
 import android.app.prediction.AppTargetEvent;
-import android.app.prediction.AppTargetId;
-import android.content.ComponentName;
 import android.content.Context;
 import android.os.Bundle;
 
-import com.android.launcher3.LauncherSettings;
-import com.android.launcher3.Workspace;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.model.data.LauncherAppWidgetInfo;
-import com.android.launcher3.model.data.WorkspaceItemInfo;
-import com.android.launcher3.shortcuts.ShortcutKey;
 
 import java.util.ArrayList;
-import java.util.Locale;
 
 /**
  * Model helper for app predictions in workspace
  */
 public class HotseatPredictionModel {
-    private static final String APP_LOCATION_HOTSEAT = "hotseat";
-    private static final String APP_LOCATION_WORKSPACE = "workspace";
-
     private static final String BUNDLE_KEY_PIN_EVENTS = "pin_events";
     private static final String BUNDLE_KEY_CURRENT_ITEMS = "current_items";
 
@@ -54,15 +46,15 @@
         ArrayList<AppTargetEvent> events = new ArrayList<>();
         ArrayList<ItemInfo> workspaceItems = dataModel.getAllWorkspaceItems();
         for (ItemInfo item : workspaceItems) {
-            AppTarget target = getAppTargetFromInfo(context, item);
-            if (target != null && !isTrackedForPrediction(item)) continue;
-            events.add(wrapAppTargetWithLocation(target, AppTargetEvent.ACTION_PIN, item));
+            AppTarget target = getAppTargetFromItemInfo(context, item);
+            if (target != null && !isTrackedForHotseatPrediction(item)) continue;
+            events.add(wrapAppTargetWithItemLocation(target, AppTargetEvent.ACTION_PIN, item));
         }
         ArrayList<AppTarget> currentTargets = new ArrayList<>();
         FixedContainerItems hotseatItems = dataModel.extraItems.get(CONTAINER_HOTSEAT_PREDICTION);
         if (hotseatItems != null) {
             for (ItemInfo itemInfo : hotseatItems.items) {
-                AppTarget target = getAppTargetFromInfo(context, itemInfo);
+                AppTarget target = getAppTargetFromItemInfo(context, itemInfo);
                 if (target != null) currentTargets.add(target);
             }
         }
@@ -70,56 +62,4 @@
         bundle.putParcelableArrayList(BUNDLE_KEY_CURRENT_ITEMS, currentTargets);
         return bundle;
     }
-
-    /**
-     * Creates and returns for {@link AppTarget} object given an {@link ItemInfo}. Returns null
-     * if item is not supported prediction
-     */
-    public static AppTarget getAppTargetFromInfo(Context context, ItemInfo info) {
-        if (info == null) return null;
-        if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET
-                && info instanceof LauncherAppWidgetInfo
-                && ((LauncherAppWidgetInfo) info).providerName != null) {
-            ComponentName cn = ((LauncherAppWidgetInfo) info).providerName;
-            return new AppTarget.Builder(new AppTargetId("widget:" + cn.getPackageName()),
-                    cn.getPackageName(), info.user).setClassName(cn.getClassName()).build();
-        } else if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_APPLICATION
-                && info.getTargetComponent() != null) {
-            ComponentName cn = info.getTargetComponent();
-            return new AppTarget.Builder(new AppTargetId("app:" + cn.getPackageName()),
-                    cn.getPackageName(), info.user).setClassName(cn.getClassName()).build();
-        } else if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT
-                && info instanceof WorkspaceItemInfo) {
-            ShortcutKey shortcutKey = ShortcutKey.fromItemInfo(info);
-            //TODO: switch to using full shortcut info
-            return new AppTarget.Builder(new AppTargetId("shortcut:" + shortcutKey.getId()),
-                    shortcutKey.componentName.getPackageName(), shortcutKey.user).build();
-        } else if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_FOLDER) {
-            return new AppTarget.Builder(new AppTargetId("folder:" + info.id),
-                    context.getPackageName(), info.user).build();
-        }
-        return null;
-    }
-
-    /**
-     * Creates and returns {@link AppTargetEvent} from an {@link AppTarget}, action, and item
-     * location using {@link ItemInfo}
-     */
-    public static AppTargetEvent wrapAppTargetWithLocation(
-            AppTarget target, int action, ItemInfo info) {
-        String location = String.format(Locale.ENGLISH, "%s/%d/[%d,%d]/[%d,%d]",
-                info.container == LauncherSettings.Favorites.CONTAINER_HOTSEAT
-                        ? APP_LOCATION_HOTSEAT : APP_LOCATION_WORKSPACE,
-                info.screenId, info.cellX, info.cellY, info.spanX, info.spanY);
-        return new AppTargetEvent.Builder(target, action).setLaunchLocation(location).build();
-    }
-
-    /**
-     * Helper method to determine if {@link ItemInfo} should be tracked and reported to predictors
-     */
-    public static boolean isTrackedForPrediction(ItemInfo info) {
-        return info.container == LauncherSettings.Favorites.CONTAINER_HOTSEAT || (
-                info.container == LauncherSettings.Favorites.CONTAINER_DESKTOP
-                        && info.screenId == Workspace.FIRST_SCREEN_ID);
-    }
 }
diff --git a/quickstep/src/com/android/launcher3/model/AppEventProducer.java b/quickstep/src/com/android/launcher3/model/AppEventProducer.java
index b665db6..09233a4 100644
--- a/quickstep/src/com/android/launcher3/model/AppEventProducer.java
+++ b/quickstep/src/com/android/launcher3/model/AppEventProducer.java
@@ -22,6 +22,7 @@
 
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_PREDICTION;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_APP_LAUNCH_TAP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_FOLDER_CONVERTED_TO_ICON;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_HOTSEAT_PREDICTION_PINNED;
@@ -35,6 +36,8 @@
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_QUICKSWITCH_RIGHT;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASK_LAUNCH_SWIPE_DOWN;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASK_LAUNCH_TAP;
+import static com.android.launcher3.model.PredictionHelper.isTrackedForHotseatPrediction;
+import static com.android.launcher3.model.PredictionHelper.isTrackedForWidgetPrediction;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
 import android.annotation.TargetApi;
@@ -62,7 +65,6 @@
 import com.android.launcher3.logger.LauncherAtom.HotseatContainer;
 import com.android.launcher3.logger.LauncherAtom.WorkspaceContainer;
 import com.android.launcher3.logging.StatsLogManager.EventEnum;
-import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.shortcuts.ShortcutRequest;
 import com.android.quickstep.logging.StatsLogCompatManager.StatsLogConsumer;
@@ -141,6 +143,9 @@
             if (isTrackedForHotseatPrediction(mLastDragItem)) {
                 sendEvent(mLastDragItem, ACTION_UNPIN, CONTAINER_HOTSEAT_PREDICTION);
             }
+            if (isTrackedForWidgetPrediction(atomInfo)) {
+                sendEvent(atomInfo, ACTION_PIN, CONTAINER_WIDGETS_PREDICTION);
+            }
             mLastDragItem = null;
         } else if (event == LAUNCHER_ITEM_DROP_FOLDER_CREATED) {
             if (isTrackedForHotseatPrediction(atomInfo)) {
@@ -158,6 +163,9 @@
             if (mLastDragItem != null && isTrackedForHotseatPrediction(mLastDragItem)) {
                 sendEvent(mLastDragItem, ACTION_UNPIN, CONTAINER_HOTSEAT_PREDICTION);
             }
+            if (mLastDragItem != null && isTrackedForWidgetPrediction(mLastDragItem)) {
+                sendEvent(mLastDragItem, ACTION_UNPIN, CONTAINER_WIDGETS_PREDICTION);
+            }
         } else if (event == LAUNCHER_HOTSEAT_PREDICTION_PINNED) {
             if (isTrackedForHotseatPrediction(atomInfo)) {
                 sendEvent(atomInfo, ACTION_PIN, CONTAINER_HOTSEAT_PREDICTION);
@@ -302,19 +310,4 @@
         return TextUtils.isEmpty(componentNameString)
                 ? null : ComponentName.unflattenFromString(componentNameString);
     }
-
-    /**
-     * Helper method to determine if {@link ItemInfo} should be tracked and reported to predictors
-     */
-    private static boolean isTrackedForHotseatPrediction(LauncherAtom.ItemInfo info) {
-        ContainerInfo ci = info.getContainerInfo();
-        switch (ci.getContainerCase()) {
-            case HOTSEAT:
-                return true;
-            case WORKSPACE:
-                return ci.getWorkspace().getPageIndex() == 0;
-            default:
-                return false;
-        }
-    }
 }
diff --git a/quickstep/src/com/android/launcher3/model/PredictionHelper.java b/quickstep/src/com/android/launcher3/model/PredictionHelper.java
new file mode 100644
index 0000000..738dd83
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/model/PredictionHelper.java
@@ -0,0 +1,130 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.model;
+
+import static com.android.launcher3.logger.LauncherAtom.ContainerInfo.ContainerCase.WORKSPACE;
+
+import android.app.prediction.AppTarget;
+import android.app.prediction.AppTargetEvent;
+import android.app.prediction.AppTargetId;
+import android.content.ComponentName;
+import android.content.Context;
+
+import androidx.annotation.Nullable;
+
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.Workspace;
+import com.android.launcher3.logger.LauncherAtom;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.LauncherAppWidgetInfo;
+import com.android.launcher3.model.data.WorkspaceItemInfo;
+import com.android.launcher3.shortcuts.ShortcutKey;
+
+import java.util.Locale;
+
+/** Helper class with methods for converting launcher items to form usable by predictors */
+public final class PredictionHelper {
+    private static final String APP_LOCATION_HOTSEAT = "hotseat";
+    private static final String APP_LOCATION_WORKSPACE = "workspace";
+
+    /**
+     * Creates and returns an {@link AppTarget} object for an {@link ItemInfo}. Returns null
+     * if item type is not supported in predictions
+     */
+    @Nullable
+    public static AppTarget getAppTargetFromItemInfo(Context context, ItemInfo info) {
+        if (info == null) return null;
+        if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET
+                && info instanceof LauncherAppWidgetInfo
+                && ((LauncherAppWidgetInfo) info).providerName != null) {
+            ComponentName cn = ((LauncherAppWidgetInfo) info).providerName;
+            return new AppTarget.Builder(new AppTargetId("widget:" + cn.getPackageName()),
+                    cn.getPackageName(), info.user).setClassName(cn.getClassName()).build();
+        } else if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_APPLICATION
+                && info.getTargetComponent() != null) {
+            ComponentName cn = info.getTargetComponent();
+            return new AppTarget.Builder(new AppTargetId("app:" + cn.getPackageName()),
+                    cn.getPackageName(), info.user).setClassName(cn.getClassName()).build();
+        } else if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT
+                && info instanceof WorkspaceItemInfo) {
+            ShortcutKey shortcutKey = ShortcutKey.fromItemInfo(info);
+            //TODO: switch to using full shortcut info
+            return new AppTarget.Builder(new AppTargetId("shortcut:" + shortcutKey.getId()),
+                    shortcutKey.componentName.getPackageName(), shortcutKey.user).build();
+        } else if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_FOLDER) {
+            return new AppTarget.Builder(new AppTargetId("folder:" + info.id),
+                    context.getPackageName(), info.user).build();
+        }
+        return null;
+    }
+
+    /**
+     * Creates and returns {@link AppTargetEvent} from an {@link AppTarget}, action, and item
+     * location using {@link ItemInfo}
+     */
+    public static AppTargetEvent wrapAppTargetWithItemLocation(
+            AppTarget target, int action, ItemInfo info) {
+        String location = String.format(Locale.ENGLISH, "%s/%d/[%d,%d]/[%d,%d]",
+                info.container == LauncherSettings.Favorites.CONTAINER_HOTSEAT
+                        ? APP_LOCATION_HOTSEAT : APP_LOCATION_WORKSPACE,
+                info.screenId, info.cellX, info.cellY, info.spanX, info.spanY);
+        return new AppTargetEvent.Builder(target, action).setLaunchLocation(location).build();
+    }
+
+    /**
+     * Helper method to determine if {@link ItemInfo} should be tracked and reported to hotseat
+     * predictors
+     */
+    public static boolean isTrackedForHotseatPrediction(ItemInfo info) {
+        return info.container == LauncherSettings.Favorites.CONTAINER_HOTSEAT || (
+                info.container == LauncherSettings.Favorites.CONTAINER_DESKTOP
+                        && info.screenId == Workspace.FIRST_SCREEN_ID);
+    }
+
+    /**
+     * Helper method to determine if {@link LauncherAtom.ItemInfo} should be tracked and reported to
+     * hotseat predictors
+     */
+    public static boolean isTrackedForHotseatPrediction(LauncherAtom.ItemInfo info) {
+        LauncherAtom.ContainerInfo ci = info.getContainerInfo();
+        switch (ci.getContainerCase()) {
+            case HOTSEAT:
+                return true;
+            case WORKSPACE:
+                return ci.getWorkspace().getPageIndex() == 0;
+            default:
+                return false;
+        }
+    }
+
+    /**
+     * Helper method to determine if {@link ItemInfo} should be tracked and reported to widget
+     * predictors
+     */
+    public static boolean isTrackedForWidgetPrediction(ItemInfo info) {
+        return info.itemType == LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET
+                && info.container == LauncherSettings.Favorites.CONTAINER_DESKTOP;
+    }
+
+    /**
+     * Helper method to determine if {@link LauncherAtom.ItemInfo} should be tracked and reported
+     * to widget predictors
+     */
+    public static boolean isTrackedForWidgetPrediction(LauncherAtom.ItemInfo info) {
+        return info.getItemCase() == LauncherAtom.ItemInfo.ItemCase.WIDGET
+                && info.getContainerInfo().getContainerCase() == WORKSPACE;
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index 55a140d..7794d27 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -25,8 +25,12 @@
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
 import static com.android.launcher3.Utilities.getDevicePrefs;
 import static com.android.launcher3.hybridhotseat.HotseatPredictionModel.convertDataModelToAppTargetBundle;
+import static com.android.launcher3.model.PredictionHelper.getAppTargetFromItemInfo;
+import static com.android.launcher3.model.PredictionHelper.wrapAppTargetWithItemLocation;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
+import static java.util.stream.Collectors.toCollection;
+
 import android.app.StatsManager;
 import android.app.prediction.AppPredictionContext;
 import android.app.prediction.AppPredictionManager;
@@ -39,6 +43,7 @@
 import android.content.pm.LauncherActivityInfo;
 import android.content.pm.LauncherApps;
 import android.content.pm.ShortcutInfo;
+import android.os.Bundle;
 import android.os.UserHandle;
 import android.util.Log;
 import android.util.StatsEvent;
@@ -62,6 +67,7 @@
 import com.android.quickstep.logging.StatsLogCompatManager;
 import com.android.systemui.shared.system.SysUiStatsLog;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -75,6 +81,7 @@
 
     public static final String LAST_PREDICTION_ENABLED_STATE = "last_prediction_enabled_state";
     private static final String LAST_SNAPSHOT_TIME_MILLIS = "LAST_SNAPSHOT_TIME_MILLIS";
+    private static final String BUNDLE_KEY_ADDED_APP_WIDGETS = "added_app_widgets";
     private static final int NUM_OF_RECOMMENDED_WIDGETS_PREDICATION = 20;
 
     private static final boolean IS_DEBUG = false;
@@ -272,6 +279,7 @@
         registerWidgetsPredictor(apm.createAppPredictionSession(
                 new AppPredictionContext.Builder(context)
                         .setUiSurface("widgets")
+                        .setExtras(getBundleForWidgetsOnWorkspace(context, mDataModel))
                         .setPredictedTargetCount(NUM_OF_RECOMMENDED_WIDGETS_PREDICATION)
                         .build()));
     }
@@ -306,12 +314,41 @@
     }
 
     private void onAppTargetEvent(AppTargetEvent event, int client) {
-        PredictorState state = client == CONTAINER_PREDICTION ? mAllAppsState : mHotseatState;
+        PredictorState state;
+        switch(client) {
+            case CONTAINER_PREDICTION:
+                state = mAllAppsState;
+                break;
+            case CONTAINER_WIDGETS_PREDICTION:
+                state = mWidgetsRecommendationState;
+                break;
+            case CONTAINER_HOTSEAT_PREDICTION:
+            default:
+                state = mHotseatState;
+                break;
+        }
         if (state.predictor != null) {
             state.predictor.notifyAppTargetEvent(event);
         }
     }
 
+    private Bundle getBundleForWidgetsOnWorkspace(Context context, BgDataModel dataModel) {
+        Bundle bundle = new Bundle();
+        ArrayList<AppTargetEvent> widgetEvents =
+                dataModel.getAllWorkspaceItems().stream()
+                        .filter(PredictionHelper::isTrackedForWidgetPrediction)
+                        .map(item -> {
+                            AppTarget target = getAppTargetFromItemInfo(context, item);
+                            if (target == null) return null;
+                            return wrapAppTargetWithItemLocation(
+                                    target, AppTargetEvent.ACTION_PIN, item);
+                        })
+                        .filter(Objects::nonNull)
+                        .collect(toCollection(ArrayList::new));
+        bundle.putParcelableArrayList(BUNDLE_KEY_ADDED_APP_WIDGETS, widgetEvents);
+        return bundle;
+    }
+
     static class PredictorState {
 
         public final FixedContainerItems items;