Merge "Make stack queries multi-display aware"
diff --git a/core/java/android/app/IActivityTaskManager.aidl b/core/java/android/app/IActivityTaskManager.aidl
index 770e5f7..37136a1 100644
--- a/core/java/android/app/IActivityTaskManager.aidl
+++ b/core/java/android/app/IActivityTaskManager.aidl
@@ -270,6 +270,9 @@
 
     List<ActivityManager.StackInfo> getAllStackInfos();
     ActivityManager.StackInfo getStackInfo(int windowingMode, int activityType);
+    List<ActivityManager.StackInfo> getAllStackInfosOnDisplay(int displayId);
+    ActivityManager.StackInfo getStackInfoOnDisplay(int windowingMode, int activityType,
+            int displayId);
 
     /**
      * Informs ActivityTaskManagerService that the keyguard is showing.
diff --git a/core/java/android/app/TaskEmbedder.java b/core/java/android/app/TaskEmbedder.java
index a1389bd..a8dc7bc 100644
--- a/core/java/android/app/TaskEmbedder.java
+++ b/core/java/android/app/TaskEmbedder.java
@@ -655,20 +655,11 @@
             // color of the surface view during resizing.
             final int displayId = getDisplayId();
             final List<ActivityManager.StackInfo> stackInfoList =
-                    mActivityTaskManager.getAllStackInfos();
-
-            // Iterate through stacks from top to bottom.
-            final int stackCount = stackInfoList.size();
-            for (int i = 0; i < stackCount; i++) {
-                final ActivityManager.StackInfo stackInfo = stackInfoList.get(i);
-                // Only look for stacks on our virtual display.
-                if (stackInfo.displayId != displayId) {
-                    continue;
-                }
-                // Found the topmost stack on target display.
-                return stackInfo;
+                    mActivityTaskManager.getAllStackInfosOnDisplay(displayId);
+            if (stackInfoList.isEmpty()) {
+                return null;
             }
-            return null;
+            return stackInfoList.get(0);
         }
     }
 }
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index 6977a45..bef4f5a 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -2773,7 +2773,7 @@
         long ident = Binder.clearCallingIdentity();
         try {
             synchronized (mGlobalLock) {
-                return mRootActivityContainer.getAllStackInfos();
+                return mRootActivityContainer.getAllStackInfos(INVALID_DISPLAY);
             }
         } finally {
             Binder.restoreCallingIdentity(ident);
@@ -2794,6 +2794,33 @@
     }
 
     @Override
+    public List<ActivityManager.StackInfo> getAllStackInfosOnDisplay(int displayId) {
+        enforceCallerIsRecentsOrHasPermission(MANAGE_ACTIVITY_STACKS, "getAllStackInfos()");
+        long ident = Binder.clearCallingIdentity();
+        try {
+            synchronized (mGlobalLock) {
+                return mRootActivityContainer.getAllStackInfos(displayId);
+            }
+        } finally {
+            Binder.restoreCallingIdentity(ident);
+        }
+    }
+
+    @Override
+    public ActivityManager.StackInfo getStackInfoOnDisplay(int windowingMode, int activityType,
+            int displayId) {
+        enforceCallerIsRecentsOrHasPermission(MANAGE_ACTIVITY_STACKS, "getStackInfo()");
+        long ident = Binder.clearCallingIdentity();
+        try {
+            synchronized (mGlobalLock) {
+                return mRootActivityContainer.getStackInfo(windowingMode, activityType, displayId);
+            }
+        } finally {
+            Binder.restoreCallingIdentity(ident);
+        }
+    }
+
+    @Override
     public void cancelRecentsAnimation(boolean restoreHomeStackPosition) {
         enforceCallerIsRecentsOrHasPermission(MANAGE_ACTIVITY_STACKS, "cancelRecentsAnimation()");
         final long callingUid = Binder.getCallingUid();
diff --git a/services/core/java/com/android/server/wm/RootActivityContainer.java b/services/core/java/com/android/server/wm/RootActivityContainer.java
index 0995df8..8f1d6ee 100644
--- a/services/core/java/com/android/server/wm/RootActivityContainer.java
+++ b/services/core/java/com/android/server/wm/RootActivityContainer.java
@@ -1243,6 +1243,15 @@
         return null;
     }
 
+    private <T extends ActivityStack> T getStack(int windowingMode, int activityType,
+            int displayId) {
+        ActivityDisplay display = getActivityDisplay(displayId);
+        if (display == null) {
+            return null;
+        }
+        return display.getStack(windowingMode, activityType);
+    }
+
     private ActivityManager.StackInfo getStackInfo(ActivityStack stack) {
         final int displayId = stack.mDisplayId;
         final ActivityDisplay display = getActivityDisplay(displayId);
@@ -1296,14 +1305,31 @@
         return (stack != null) ? getStackInfo(stack) : null;
     }
 
-    ArrayList<ActivityManager.StackInfo> getAllStackInfos() {
+    ActivityManager.StackInfo getStackInfo(int windowingMode, int activityType, int displayId) {
+        final ActivityStack stack = getStack(windowingMode, activityType, displayId);
+        return (stack != null) ? getStackInfo(stack) : null;
+    }
+
+    /** If displayId == INVALID_DISPLAY, this will get stack infos on all displays */
+    ArrayList<ActivityManager.StackInfo> getAllStackInfos(int displayId) {
         ArrayList<ActivityManager.StackInfo> list = new ArrayList<>();
-        for (int displayNdx = 0; displayNdx < mActivityDisplays.size(); ++displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
-            for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
-                final ActivityStack stack = display.getChildAt(stackNdx);
-                list.add(getStackInfo(stack));
+        if (displayId == INVALID_DISPLAY) {
+            for (int displayNdx = 0; displayNdx < mActivityDisplays.size(); ++displayNdx) {
+                final ActivityDisplay display = mActivityDisplays.get(displayNdx);
+                for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
+                    final ActivityStack stack = display.getChildAt(stackNdx);
+                    list.add(getStackInfo(stack));
+                }
             }
+            return list;
+        }
+        final ActivityDisplay display = getActivityDisplay(displayId);
+        if (display == null) {
+            return list;
+        }
+        for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
+            final ActivityStack stack = display.getChildAt(stackNdx);
+            list.add(getStackInfo(stack));
         }
         return list;
     }