[AA+] Log app launches from DeviceSearchResultContainer with instance ID.

This InstanceId is used to recreate the AA+ session on the server side.

Bug: 178562918
Change-Id: I1bba94417d3a142351e2470bb6153707d2cadb11
diff --git a/quickstep/src/com/android/launcher3/search/SearchResultIcon.java b/quickstep/src/com/android/launcher3/search/SearchResultIcon.java
index 521174b..f7d5f45 100644
--- a/quickstep/src/com/android/launcher3/search/SearchResultIcon.java
+++ b/quickstep/src/com/android/launcher3/search/SearchResultIcon.java
@@ -149,7 +149,7 @@
             // Workaround to log ItemInfo with DeviceSearchResultContainer without
             // updating ItemInfo.container field.
             @Override
-            protected ContainerInfo getContainerInfo() {
+            public ContainerInfo getContainerInfo() {
                 return buildDeviceSearchResultContainer();
             }
         };
@@ -257,7 +257,7 @@
             // Workaround to log ItemInfo with DeviceSearchResultContainer without
             // updating ItemInfo.container field.
             @Override
-            protected ContainerInfo getContainerInfo() {
+            public ContainerInfo getContainerInfo() {
                 return buildDeviceSearchResultContainer();
             }
         };
@@ -275,7 +275,7 @@
             // Workaround to log ItemInfo with DeviceSearchResultContainer without
             // updating ItemInfo.container field.
             @Override
-            protected ContainerInfo getContainerInfo() {
+            public ContainerInfo getContainerInfo() {
                 return buildDeviceSearchResultContainer();
             }
         };
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index a00ce56..ee00ed2 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -24,6 +24,8 @@
 import static com.android.launcher3.LauncherState.OVERVIEW;
 import static com.android.launcher3.LauncherState.OVERVIEW_MODAL_TASK;
 import static com.android.launcher3.compat.AccessibilityManagerCompat.sendCustomAccessibilityEvent;
+import static com.android.launcher3.logger.LauncherAtom.ContainerInfo.ContainerCase.EXTENDED_CONTAINERS;
+import static com.android.launcher3.logger.LauncherAtomExtensions.ExtendedContainers.ContainerCase.DEVICE_SEARCH_RESULT_CONTAINER;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_APP_LAUNCH_TAP;
 import static com.android.launcher3.testing.TestProtocol.HINT_STATE_ORDINAL;
 import static com.android.launcher3.testing.TestProtocol.OVERVIEW_STATE_ORDINAL;
@@ -81,6 +83,7 @@
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.stream.Stream;
 
 public class QuickstepLauncher extends BaseQuickstepLauncher {
@@ -105,6 +108,15 @@
 
     @Override
     protected void logAppLaunch(ItemInfo info, InstanceId instanceId) {
+        // If the app launch is from DeviceSearchResultContainer then add the InstanceId from
+        // LiveSearchManager to recreate the AllApps search session on the server side.
+        Optional<InstanceId> logInstanceId = this.getLiveSearchManager().getLogInstanceId();
+        if (info.getContainerInfo().getContainerCase() == EXTENDED_CONTAINERS
+                && info.getContainerInfo().getExtendedContainers().getContainerCase()
+                == DEVICE_SEARCH_RESULT_CONTAINER && logInstanceId.isPresent()) {
+            instanceId = logInstanceId.get();
+        }
+
         StatsLogger logger = getStatsLogManager()
                 .logger().withItemInfo(info).withInstanceId(instanceId);
 
diff --git a/src/com/android/launcher3/allapps/search/LiveSearchManager.java b/src/com/android/launcher3/allapps/search/LiveSearchManager.java
index ec33908..c2f0b96 100644
--- a/src/com/android/launcher3/allapps/search/LiveSearchManager.java
+++ b/src/com/android/launcher3/allapps/search/LiveSearchManager.java
@@ -33,10 +33,13 @@
 
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.logging.InstanceId;
+import com.android.launcher3.logging.InstanceIdSequence;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.widget.PendingAddWidgetInfo;
 
 import java.util.HashMap;
+import java.util.Optional;
 
 /**
  * Manages Lifecycle for Live search results
@@ -51,6 +54,7 @@
             new HashMap<>();
     private final HashMap<Uri, LiveData<Slice>> mUriSliceMap = new HashMap<>();
     private SearchWidgetHost mSearchWidgetHost;
+    private InstanceId mLogInstanceId;
 
     public LiveSearchManager(Launcher launcher) {
         mLauncher = launcher;
@@ -113,6 +117,7 @@
      */
     public void start() {
         stop();
+        mLogInstanceId = new InstanceIdSequence().newInstanceId();
         mSearchWidgetHost = new SearchWidgetHost(mLauncher);
         mSearchWidgetHost.startListening();
     }
@@ -136,6 +141,14 @@
         mUriSliceMap.clear();
     }
 
+    /**
+     * Returns {@link InstanceId} that should be used for logging events within search session, if
+     * available.
+     */
+    public Optional<InstanceId> getLogInstanceId() {
+        return Optional.ofNullable(mLogInstanceId);
+    }
+
     static class SearchWidgetHost extends AppWidgetHost {
         SearchWidgetHost(Context context) {
             super(context, SEARCH_APPWIDGET_HOST_ID);
diff --git a/src/com/android/launcher3/model/data/ItemInfo.java b/src/com/android/launcher3/model/data/ItemInfo.java
index 2fc9023..3851ab0 100644
--- a/src/com/android/launcher3/model/data/ItemInfo.java
+++ b/src/com/android/launcher3/model/data/ItemInfo.java
@@ -352,7 +352,10 @@
         return itemBuilder;
     }
 
-    protected ContainerInfo getContainerInfo() {
+    /**
+     * Returns {@link ContainerInfo} used when logging this item.
+     */
+    public ContainerInfo getContainerInfo() {
         switch (container) {
             case CONTAINER_HOTSEAT:
                 return ContainerInfo.newBuilder()