Merge "Modifying the mechanism of A11y window cache"
diff --git a/core/java/android/view/accessibility/AccessibilityCache.java b/core/java/android/view/accessibility/AccessibilityCache.java
index 484d9a1..dc8bf9b 100644
--- a/core/java/android/view/accessibility/AccessibilityCache.java
+++ b/core/java/android/view/accessibility/AccessibilityCache.java
@@ -71,7 +71,9 @@
 
     private boolean mIsAllWindowsCached;
 
-    private final SparseArray<AccessibilityWindowInfo> mWindowCache =
+    // The SparseArray of all {@link AccessibilityWindowInfo}s on all displays.
+    // The key of outer SparseArray is display ID and the key of inner SparseArray is window ID.
+    private final SparseArray<SparseArray<AccessibilityWindowInfo>> mWindowCacheByDisplay =
             new SparseArray<>();
 
     private final SparseArray<LongSparseArray<AccessibilityNodeInfo>> mNodeCache =
@@ -84,34 +86,66 @@
         mAccessibilityNodeRefresher = nodeRefresher;
     }
 
-    public void setWindows(List<AccessibilityWindowInfo> windows) {
+    /**
+     * Sets all {@link AccessibilityWindowInfo}s of all displays into the cache.
+     * The key of SparseArray is display ID.
+     *
+     * @param windowsOnAllDisplays The accessibility windows of all displays.
+     */
+    public void setWindowsOnAllDisplays(
+            SparseArray<List<AccessibilityWindowInfo>> windowsOnAllDisplays) {
         synchronized (mLock) {
             if (DEBUG) {
                 Log.i(LOG_TAG, "Set windows");
             }
-            clearWindowCache();
-            if (windows == null) {
+            clearWindowCacheLocked();
+            if (windowsOnAllDisplays == null) {
                 return;
             }
-            final int windowCount = windows.size();
-            for (int i = 0; i < windowCount; i++) {
-                final AccessibilityWindowInfo window = windows.get(i);
-                addWindow(window);
+
+            final int displayCounts = windowsOnAllDisplays.size();
+            for (int i = 0; i < displayCounts; i++) {
+                final List<AccessibilityWindowInfo> windowsOfDisplay =
+                        windowsOnAllDisplays.valueAt(i);
+
+                if (windowsOfDisplay == null) {
+                    continue;
+                }
+
+                final int displayId = windowsOnAllDisplays.keyAt(i);
+                final int windowCount = windowsOfDisplay.size();
+                for (int j = 0; j < windowCount; j++) {
+                    addWindowByDisplayLocked(displayId, windowsOfDisplay.get(j));
+                }
             }
             mIsAllWindowsCached = true;
         }
     }
 
+    /**
+     * Sets an {@link AccessibilityWindowInfo} into the cache.
+     *
+     * @param window The accessibility window.
+     */
     public void addWindow(AccessibilityWindowInfo window) {
         synchronized (mLock) {
             if (DEBUG) {
-                Log.i(LOG_TAG, "Caching window: " + window.getId());
+                Log.i(LOG_TAG, "Caching window: " + window.getId() + " at display Id [ "
+                        + window.getDisplayId() + " ]");
             }
-            final int windowId = window.getId();
-            mWindowCache.put(windowId, new AccessibilityWindowInfo(window));
+            addWindowByDisplayLocked(window.getDisplayId(), window);
         }
     }
 
+    private void addWindowByDisplayLocked(int displayId, AccessibilityWindowInfo window) {
+        SparseArray<AccessibilityWindowInfo> windows = mWindowCacheByDisplay.get(displayId);
+        if (windows == null) {
+            windows = new SparseArray<>();
+            mWindowCacheByDisplay.put(displayId, windows);
+        }
+        final int windowId = window.getId();
+        windows.put(windowId, new AccessibilityWindowInfo(window));
+    }
     /**
      * Notifies the cache that the something in the UI changed. As a result
      * the cache will either refresh some nodes or evict some nodes.
@@ -236,44 +270,82 @@
         }
     }
 
-    public List<AccessibilityWindowInfo> getWindows() {
+    /**
+     * Gets all {@link AccessibilityWindowInfo}s of all displays from the cache.
+     *
+     * @return All cached {@link AccessibilityWindowInfo}s of all displays
+     *         or null if such not found. The key of SparseArray is display ID.
+     */
+    public SparseArray<List<AccessibilityWindowInfo>> getWindowsOnAllDisplays() {
         synchronized (mLock) {
             if (!mIsAllWindowsCached) {
                 return null;
             }
+            final SparseArray<List<AccessibilityWindowInfo>> returnWindows = new SparseArray<>();
+            final int displayCounts = mWindowCacheByDisplay.size();
 
-            final int windowCount = mWindowCache.size();
-            if (windowCount > 0) {
-                // Careful to return the windows in a decreasing layer order.
-                SparseArray<AccessibilityWindowInfo> sortedWindows = mTempWindowArray;
-                sortedWindows.clear();
+            if (displayCounts > 0) {
+                for (int i = 0; i < displayCounts; i++) {
+                    final int displayId = mWindowCacheByDisplay.keyAt(i);
+                    final SparseArray<AccessibilityWindowInfo> windowsOfDisplay =
+                            mWindowCacheByDisplay.valueAt(i);
 
-                for (int i = 0; i < windowCount; i++) {
-                    AccessibilityWindowInfo window = mWindowCache.valueAt(i);
-                    sortedWindows.put(window.getLayer(), window);
+                    if (windowsOfDisplay == null) {
+                        continue;
+                    }
+
+                    final int windowCount = windowsOfDisplay.size();
+                    if (windowCount > 0) {
+                        // Careful to return the windows in a decreasing layer order.
+                        SparseArray<AccessibilityWindowInfo> sortedWindows = mTempWindowArray;
+                        sortedWindows.clear();
+
+                        for (int j = 0; j < windowCount; j++) {
+                            AccessibilityWindowInfo window = windowsOfDisplay.valueAt(j);
+                            sortedWindows.put(window.getLayer(), window);
+                        }
+
+                        // It's possible in transient conditions for two windows to share the same
+                        // layer, which results in sortedWindows being smaller than
+                        // mWindowCacheByDisplay
+                        final int sortedWindowCount = sortedWindows.size();
+                        List<AccessibilityWindowInfo> windows =
+                                new ArrayList<>(sortedWindowCount);
+                        for (int j = sortedWindowCount - 1; j >= 0; j--) {
+                            AccessibilityWindowInfo window = sortedWindows.valueAt(j);
+                            windows.add(new AccessibilityWindowInfo(window));
+                            sortedWindows.removeAt(j);
+                        }
+                        returnWindows.put(displayId, windows);
+                    }
                 }
-
-                // It's possible in transient conditions for two windows to share the same
-                // layer, which results in sortedWindows being smaller than mWindowCache
-                final int sortedWindowCount = sortedWindows.size();
-                List<AccessibilityWindowInfo> windows = new ArrayList<>(sortedWindowCount);
-                for (int i = sortedWindowCount - 1; i >= 0; i--) {
-                    AccessibilityWindowInfo window = sortedWindows.valueAt(i);
-                    windows.add(new AccessibilityWindowInfo(window));
-                    sortedWindows.removeAt(i);
-                }
-
-                return windows;
+                return returnWindows;
             }
             return null;
         }
     }
 
+    /**
+     * Gets an {@link AccessibilityWindowInfo} by windowId.
+     *
+     * @param windowId The id of the window.
+     *
+     * @return The {@link AccessibilityWindowInfo} or null if such not found.
+     */
     public AccessibilityWindowInfo getWindow(int windowId) {
         synchronized (mLock) {
-            AccessibilityWindowInfo window = mWindowCache.get(windowId);
-            if (window != null) {
-                return new AccessibilityWindowInfo(window);
+            final int displayCounts = mWindowCacheByDisplay.size();
+            for (int i = 0; i < displayCounts; i++) {
+                final SparseArray<AccessibilityWindowInfo> windowsOfDisplay =
+                        mWindowCacheByDisplay.valueAt(i);
+                if (windowsOfDisplay == null) {
+                    continue;
+                }
+
+                AccessibilityWindowInfo window = windowsOfDisplay.get(windowId);
+                if (window != null) {
+                    return new AccessibilityWindowInfo(window);
+                }
             }
             return null;
         }
@@ -358,7 +430,7 @@
             if (DEBUG) {
                 Log.i(LOG_TAG, "clear()");
             }
-            clearWindowCache();
+            clearWindowCacheLocked();
             final int nodesForWindowCount = mNodeCache.size();
             for (int i = nodesForWindowCount - 1; i >= 0; i--) {
                 final int windowId = mNodeCache.keyAt(i);
@@ -370,8 +442,23 @@
         }
     }
 
-    private void clearWindowCache() {
-        mWindowCache.clear();
+    private void clearWindowCacheLocked() {
+        if (DEBUG) {
+            Log.i(LOG_TAG, "clearWindowCacheLocked");
+        }
+        final int displayCounts = mWindowCacheByDisplay.size();
+
+        if (displayCounts > 0) {
+            for (int i = displayCounts - 1; i >= 0; i--) {
+                final int displayId = mWindowCacheByDisplay.keyAt(i);
+                final SparseArray<AccessibilityWindowInfo> windows =
+                        mWindowCacheByDisplay.get(displayId);
+                if (windows != null) {
+                    windows.clear();
+                }
+                mWindowCacheByDisplay.remove(displayId);
+            }
+        }
         mIsAllWindowsCached = false;
     }
 
@@ -444,32 +531,41 @@
     public void checkIntegrity() {
         synchronized (mLock) {
             // Get the root.
-            if (mWindowCache.size() <= 0 && mNodeCache.size() == 0) {
+            if (mWindowCacheByDisplay.size() <= 0 && mNodeCache.size() == 0) {
                 return;
             }
 
             AccessibilityWindowInfo focusedWindow = null;
             AccessibilityWindowInfo activeWindow = null;
 
-            final int windowCount = mWindowCache.size();
-            for (int i = 0; i < windowCount; i++) {
-                AccessibilityWindowInfo window = mWindowCache.valueAt(i);
+            final int displayCounts = mWindowCacheByDisplay.size();
+            for (int i = 0; i < displayCounts; i++) {
+                final SparseArray<AccessibilityWindowInfo> windowsOfDisplay =
+                        mWindowCacheByDisplay.valueAt(i);
 
-                // Check for one active window.
-                if (window.isActive()) {
-                    if (activeWindow != null) {
-                        Log.e(LOG_TAG, "Duplicate active window:" + window);
-                    } else {
-                        activeWindow = window;
-                    }
+                if (windowsOfDisplay == null) {
+                    continue;
                 }
 
-                // Check for one focused window.
-                if (window.isFocused()) {
-                    if (focusedWindow != null) {
-                        Log.e(LOG_TAG, "Duplicate focused window:" + window);
-                    } else {
-                        focusedWindow = window;
+                final int windowCount = windowsOfDisplay.size();
+                for (int j = 0; j < windowCount; j++) {
+                    final AccessibilityWindowInfo window = windowsOfDisplay.valueAt(j);
+
+                    // Check for one active window.
+                    if (window.isActive()) {
+                        if (activeWindow != null) {
+                            Log.e(LOG_TAG, "Duplicate active window:" + window);
+                        } else {
+                            activeWindow = window;
+                        }
+                    }
+                    // Check for one focused window.
+                    if (window.isFocused()) {
+                        if (focusedWindow != null) {
+                            Log.e(LOG_TAG, "Duplicate focused window:" + window);
+                        } else {
+                            focusedWindow = window;
+                        }
                     }
                 }
             }
diff --git a/core/java/android/view/accessibility/AccessibilityInteractionClient.java b/core/java/android/view/accessibility/AccessibilityInteractionClient.java
index 4db6f4f..d9fa9f2 100644
--- a/core/java/android/view/accessibility/AccessibilityInteractionClient.java
+++ b/core/java/android/view/accessibility/AccessibilityInteractionClient.java
@@ -28,6 +28,7 @@
 import android.util.Log;
 import android.util.LongSparseArray;
 import android.util.SparseArray;
+import android.view.Display;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.ArrayUtils;
@@ -267,12 +268,14 @@
         try {
             IAccessibilityServiceConnection connection = getConnection(connectionId);
             if (connection != null) {
-                List<AccessibilityWindowInfo> windows = sAccessibilityCache.getWindows();
-                if (windows != null) {
+                SparseArray<List<AccessibilityWindowInfo>> allWindows =
+                        sAccessibilityCache.getWindowsOnAllDisplays();
+                List<AccessibilityWindowInfo> windows;
+                if (allWindows != null) {
                     if (DEBUG) {
                         Log.i(LOG_TAG, "Windows cache hit");
                     }
-                    return windows;
+                    return allWindows.valueAt(Display.DEFAULT_DISPLAY);
                 }
                 if (DEBUG) {
                     Log.i(LOG_TAG, "Windows cache miss");
@@ -284,7 +287,9 @@
                     Binder.restoreCallingIdentity(identityToken);
                 }
                 if (windows != null) {
-                    sAccessibilityCache.setWindows(windows);
+                    allWindows = new SparseArray<>();
+                    allWindows.put(Display.DEFAULT_DISPLAY, windows);
+                    sAccessibilityCache.setWindowsOnAllDisplays(allWindows);
                     return windows;
                 }
             } else {
diff --git a/core/tests/coretests/src/android/view/accessibility/AccessibilityCacheTest.java b/core/tests/coretests/src/android/view/accessibility/AccessibilityCacheTest.java
index 1a22a70..6bce651 100644
--- a/core/tests/coretests/src/android/view/accessibility/AccessibilityCacheTest.java
+++ b/core/tests/coretests/src/android/view/accessibility/AccessibilityCacheTest.java
@@ -29,6 +29,8 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.util.SparseArray;
+import android.view.Display;
 import android.view.View;
 
 import androidx.test.filters.LargeTest;
@@ -51,12 +53,17 @@
 public class AccessibilityCacheTest {
     private static final int WINDOW_ID_1 = 0xBEEF;
     private static final int WINDOW_ID_2 = 0xFACE;
+    private static final int WINDOW_ID_3 = 0xABCD;
+    private static final int WINDOW_ID_4 = 0xDCBA;
     private static final int SINGLE_VIEW_ID = 0xCAFE;
     private static final int OTHER_VIEW_ID = 0xCAB2;
     private static final int PARENT_VIEW_ID = 0xFED4;
     private static final int CHILD_VIEW_ID = 0xFEED;
     private static final int OTHER_CHILD_VIEW_ID = 0xACE2;
     private static final int MOCK_CONNECTION_ID = 1;
+    private static final int SECONDARY_DISPLAY_ID = Display.DEFAULT_DISPLAY + 1;
+    private static final int DEFAULT_WINDOW_LAYER = 0;
+    private static final int SPECIFIC_WINDOW_LAYER = 5;
 
     AccessibilityCache mAccessibilityCache;
     AccessibilityCache.AccessibilityNodeRefresher mAccessibilityNodeRefresher;
@@ -70,7 +77,7 @@
 
     @After
     public void tearDown() {
-        // Make sure we're recycling all of our window and node infos
+        // Make sure we're recycling all of our window and node infos.
         mAccessibilityCache.clear();
         AccessibilityInteractionClient.getInstance().clearCache();
     }
@@ -78,7 +85,7 @@
     @Test
     public void testEmptyCache_returnsNull() {
         assertNull(mAccessibilityCache.getNode(0, 0));
-        assertNull(mAccessibilityCache.getWindows());
+        assertNull(mAccessibilityCache.getWindowsOnAllDisplays());
         assertNull(mAccessibilityCache.getWindow(0));
     }
 
@@ -114,10 +121,11 @@
         try {
             windowInfo = AccessibilityWindowInfo.obtain();
             windowInfo.setId(WINDOW_ID_1);
+            windowInfo.setDisplayId(Display.DEFAULT_DISPLAY);
             mAccessibilityCache.addWindow(windowInfo);
             // Make a copy
             copyOfInfo = AccessibilityWindowInfo.obtain(windowInfo);
-            windowInfo.setId(WINDOW_ID_2); // Simulate recycling and reusing the original info
+            windowInfo.setId(WINDOW_ID_2); // Simulate recycling and reusing the original info.
             windowFromCache = mAccessibilityCache.getWindow(WINDOW_ID_1);
             assertEquals(copyOfInfo, windowFromCache);
         } finally {
@@ -129,39 +137,40 @@
 
     @Test
     public void addWindowThenClear_noLongerInCache() {
-        putWindowWithIdInCache(WINDOW_ID_1);
+        putWindowWithWindowIdAndDisplayIdInCache(WINDOW_ID_1, Display.DEFAULT_DISPLAY,
+                DEFAULT_WINDOW_LAYER);
         mAccessibilityCache.clear();
         assertNull(mAccessibilityCache.getWindow(WINDOW_ID_1));
     }
 
     @Test
     public void addWindowGetOtherId_returnsNull() {
-        putWindowWithIdInCache(WINDOW_ID_1);
+        putWindowWithWindowIdAndDisplayIdInCache(WINDOW_ID_1, Display.DEFAULT_DISPLAY,
+                DEFAULT_WINDOW_LAYER);
         assertNull(mAccessibilityCache.getWindow(WINDOW_ID_1 + 1));
     }
 
     @Test
     public void addWindowThenGetWindows_returnsNull() {
-        putWindowWithIdInCache(WINDOW_ID_1);
-        assertNull(mAccessibilityCache.getWindows());
+        putWindowWithWindowIdAndDisplayIdInCache(WINDOW_ID_1, Display.DEFAULT_DISPLAY,
+                DEFAULT_WINDOW_LAYER);
+        assertNull(mAccessibilityCache.getWindowsOnAllDisplays());
     }
 
     @Test
     public void setWindowsThenGetWindows_returnsInDecreasingLayerOrder() {
-        AccessibilityWindowInfo windowInfo1 = null, windowInfo2 = null;
-        AccessibilityWindowInfo window1Out = null, window2Out = null;
+        AccessibilityWindowInfo windowInfo1 = null;
+        AccessibilityWindowInfo windowInfo2 = null;
+        AccessibilityWindowInfo window1Out = null;
+        AccessibilityWindowInfo window2Out = null;
         List<AccessibilityWindowInfo> windowsOut = null;
         try {
-            windowInfo1 = AccessibilityWindowInfo.obtain();
-            windowInfo1.setId(WINDOW_ID_1);
-            windowInfo1.setLayer(5);
-            windowInfo2 = AccessibilityWindowInfo.obtain();
-            windowInfo2.setId(WINDOW_ID_2);
-            windowInfo2.setLayer(windowInfo1.getLayer() + 1);
+            windowInfo1 = obtainAccessibilityWindowInfo(WINDOW_ID_1, SPECIFIC_WINDOW_LAYER);
+            windowInfo2 = obtainAccessibilityWindowInfo(WINDOW_ID_2, windowInfo1.getLayer() + 1);
             List<AccessibilityWindowInfo> windowsIn = Arrays.asList(windowInfo1, windowInfo2);
-            mAccessibilityCache.setWindows(windowsIn);
+            setWindowsByDisplay(Display.DEFAULT_DISPLAY, windowsIn);
 
-            windowsOut = mAccessibilityCache.getWindows();
+            windowsOut = getWindowsByDisplay(Display.DEFAULT_DISPLAY);
             window1Out = mAccessibilityCache.getWindow(WINDOW_ID_1);
             window2Out = mAccessibilityCache.getWindow(WINDOW_ID_2);
 
@@ -182,8 +191,151 @@
     }
 
     @Test
+    public void setWindowsAndAddWindow_thenGetWindows_returnsInDecreasingLayerOrder() {
+        AccessibilityWindowInfo windowInfo1 = null;
+        AccessibilityWindowInfo windowInfo2 = null;
+        AccessibilityWindowInfo window1Out = null;
+        AccessibilityWindowInfo window2Out = null;
+        AccessibilityWindowInfo window3Out = null;
+        List<AccessibilityWindowInfo> windowsOut = null;
+        try {
+            windowInfo1 = obtainAccessibilityWindowInfo(WINDOW_ID_1, SPECIFIC_WINDOW_LAYER);
+            windowInfo2 = obtainAccessibilityWindowInfo(WINDOW_ID_2, windowInfo1.getLayer() + 2);
+            List<AccessibilityWindowInfo> windowsIn = Arrays.asList(windowInfo1, windowInfo2);
+            setWindowsByDisplay(Display.DEFAULT_DISPLAY, windowsIn);
+
+            putWindowWithWindowIdAndDisplayIdInCache(WINDOW_ID_3, Display.DEFAULT_DISPLAY,
+                    windowInfo1.getLayer() + 1);
+
+            windowsOut = getWindowsByDisplay(Display.DEFAULT_DISPLAY);
+            window1Out = mAccessibilityCache.getWindow(WINDOW_ID_1);
+            window2Out = mAccessibilityCache.getWindow(WINDOW_ID_2);
+            window3Out = mAccessibilityCache.getWindow(WINDOW_ID_3);
+
+            assertEquals(3, windowsOut.size());
+            assertEquals(windowInfo2, windowsOut.get(0));
+            assertEquals(windowInfo1, windowsOut.get(2));
+            assertEquals(windowInfo1, window1Out);
+            assertEquals(windowInfo2, window2Out);
+            assertEquals(window3Out, windowsOut.get(1));
+        } finally {
+            window1Out.recycle();
+            window2Out.recycle();
+            window3Out.recycle();
+            windowInfo1.recycle();
+            windowInfo2.recycle();
+            for (AccessibilityWindowInfo windowInfo : windowsOut) {
+                windowInfo.recycle();
+            }
+        }
+    }
+
+    @Test
+    public void
+            setWindowsAtFirstDisplay_thenAddWindowAtSecondDisplay_returnWindowLayerOrderUnchange() {
+        AccessibilityWindowInfo windowInfo1 = null;
+        AccessibilityWindowInfo windowInfo2 = null;
+        AccessibilityWindowInfo window1Out = null;
+        AccessibilityWindowInfo window2Out = null;
+        List<AccessibilityWindowInfo> windowsOut = null;
+        try {
+            // Sets windows to default display.
+            windowInfo1 = obtainAccessibilityWindowInfo(WINDOW_ID_1, SPECIFIC_WINDOW_LAYER);
+            windowInfo2 = obtainAccessibilityWindowInfo(WINDOW_ID_2, windowInfo1.getLayer() + 2);
+            List<AccessibilityWindowInfo> windowsIn = Arrays.asList(windowInfo1, windowInfo2);
+            setWindowsByDisplay(Display.DEFAULT_DISPLAY, windowsIn);
+            // Adds one window to second display.
+            putWindowWithWindowIdAndDisplayIdInCache(WINDOW_ID_3, SECONDARY_DISPLAY_ID,
+                    windowInfo1.getLayer() + 1);
+
+            windowsOut = getWindowsByDisplay(Display.DEFAULT_DISPLAY);
+            window1Out = mAccessibilityCache.getWindow(WINDOW_ID_1);
+            window2Out = mAccessibilityCache.getWindow(WINDOW_ID_2);
+
+            assertEquals(2, windowsOut.size());
+            assertEquals(windowInfo2, windowsOut.get(0));
+            assertEquals(windowInfo1, windowsOut.get(1));
+            assertEquals(windowInfo1, window1Out);
+            assertEquals(windowInfo2, window2Out);
+        } finally {
+            window1Out.recycle();
+            window2Out.recycle();
+            windowInfo1.recycle();
+            windowInfo2.recycle();
+            for (AccessibilityWindowInfo windowInfo : windowsOut) {
+                windowInfo.recycle();
+            }
+        }
+    }
+
+    @Test
+    public void setWindowsAtTwoDisplays_thenGetWindows_returnsInDecreasingLayerOrder() {
+        AccessibilityWindowInfo windowInfo1 = null;
+        AccessibilityWindowInfo windowInfo2 = null;
+        AccessibilityWindowInfo window1Out = null;
+        AccessibilityWindowInfo window2Out = null;
+        AccessibilityWindowInfo windowInfo3 = null;
+        AccessibilityWindowInfo windowInfo4 = null;
+        AccessibilityWindowInfo window3Out = null;
+        AccessibilityWindowInfo window4Out = null;
+        List<AccessibilityWindowInfo> windowsOut1 = null;
+        List<AccessibilityWindowInfo> windowsOut2 = null;
+        try {
+            // Prepares all windows for default display.
+            windowInfo1 = obtainAccessibilityWindowInfo(WINDOW_ID_1, SPECIFIC_WINDOW_LAYER);
+            windowInfo2 = obtainAccessibilityWindowInfo(WINDOW_ID_2, windowInfo1.getLayer() + 1);
+            List<AccessibilityWindowInfo> windowsIn1 = Arrays.asList(windowInfo1, windowInfo2);
+            // Prepares all windows for second display.
+            windowInfo3 = obtainAccessibilityWindowInfo(WINDOW_ID_3, windowInfo1.getLayer() + 2);
+            windowInfo4 = obtainAccessibilityWindowInfo(WINDOW_ID_4, windowInfo1.getLayer() + 3);
+            List<AccessibilityWindowInfo> windowsIn2 = Arrays.asList(windowInfo3, windowInfo4);
+            // Sets all windows of all displays into A11y cache.
+            SparseArray<List<AccessibilityWindowInfo>> allWindows = new SparseArray<>();
+            allWindows.put(Display.DEFAULT_DISPLAY, windowsIn1);
+            allWindows.put(SECONDARY_DISPLAY_ID, windowsIn2);
+            mAccessibilityCache.setWindowsOnAllDisplays(allWindows);
+            // Gets windows at default display.
+            windowsOut1 = getWindowsByDisplay(Display.DEFAULT_DISPLAY);
+            window1Out = mAccessibilityCache.getWindow(WINDOW_ID_1);
+            window2Out = mAccessibilityCache.getWindow(WINDOW_ID_2);
+
+            assertEquals(2, windowsOut1.size());
+            assertEquals(windowInfo2, windowsOut1.get(0));
+            assertEquals(windowInfo1, windowsOut1.get(1));
+            assertEquals(windowInfo1, window1Out);
+            assertEquals(windowInfo2, window2Out);
+            // Gets windows at seocnd display.
+            windowsOut2 = getWindowsByDisplay(SECONDARY_DISPLAY_ID);
+            window3Out = mAccessibilityCache.getWindow(WINDOW_ID_3);
+            window4Out = mAccessibilityCache.getWindow(WINDOW_ID_4);
+
+            assertEquals(2, windowsOut2.size());
+            assertEquals(windowInfo4, windowsOut2.get(0));
+            assertEquals(windowInfo3, windowsOut2.get(1));
+            assertEquals(windowInfo3, window3Out);
+            assertEquals(windowInfo4, window4Out);
+        } finally {
+            window1Out.recycle();
+            window2Out.recycle();
+            windowInfo1.recycle();
+            windowInfo2.recycle();
+            window3Out.recycle();
+            window4Out.recycle();
+            windowInfo3.recycle();
+            windowInfo4.recycle();
+            for (AccessibilityWindowInfo windowInfo : windowsOut1) {
+                windowInfo.recycle();
+            }
+            for (AccessibilityWindowInfo windowInfo : windowsOut2) {
+                windowInfo.recycle();
+            }
+        }
+    }
+
+    @Test
     public void addWindowThenStateChangedEvent_noLongerInCache() {
-        putWindowWithIdInCache(WINDOW_ID_1);
+        putWindowWithWindowIdAndDisplayIdInCache(WINDOW_ID_1, Display.DEFAULT_DISPLAY,
+                DEFAULT_WINDOW_LAYER);
         mAccessibilityCache.onAccessibilityEvent(
                 AccessibilityEvent.obtain(AccessibilityEvent.TYPE_WINDOW_STATE_CHANGED));
         assertNull(mAccessibilityCache.getWindow(WINDOW_ID_1));
@@ -191,7 +343,8 @@
 
     @Test
     public void addWindowThenWindowsChangedEvent_noLongerInCache() {
-        putWindowWithIdInCache(WINDOW_ID_1);
+        putWindowWithWindowIdAndDisplayIdInCache(WINDOW_ID_1, Display.DEFAULT_DISPLAY,
+                DEFAULT_WINDOW_LAYER);
         mAccessibilityCache.onAccessibilityEvent(
                 AccessibilityEvent.obtain(AccessibilityEvent.TYPE_WINDOWS_CHANGED));
         assertNull(mAccessibilityCache.getWindow(WINDOW_ID_1));
@@ -622,9 +775,16 @@
         }
     }
 
-    private void putWindowWithIdInCache(int id) {
+    private AccessibilityWindowInfo obtainAccessibilityWindowInfo(int windowId, int layer) {
         AccessibilityWindowInfo windowInfo = AccessibilityWindowInfo.obtain();
-        windowInfo.setId(id);
+        windowInfo.setId(windowId);
+        windowInfo.setLayer(layer);
+        return windowInfo;
+    }
+
+    private void putWindowWithWindowIdAndDisplayIdInCache(int windowId, int displayId, int layer) {
+        AccessibilityWindowInfo windowInfo = obtainAccessibilityWindowInfo(windowId, layer);
+        windowInfo.setDisplayId(displayId);
         mAccessibilityCache.addWindow(windowInfo);
         windowInfo.recycle();
     }
@@ -713,4 +873,20 @@
             }
         }
     }
+
+    private void setWindowsByDisplay(int displayId, List<AccessibilityWindowInfo> windows) {
+        SparseArray<List<AccessibilityWindowInfo>> allWindows = new SparseArray<>();
+        allWindows.put(displayId, windows);
+        mAccessibilityCache.setWindowsOnAllDisplays(allWindows);
+    }
+
+    private List<AccessibilityWindowInfo> getWindowsByDisplay(int displayId) {
+        final SparseArray<List<AccessibilityWindowInfo>> allWindows =
+                mAccessibilityCache.getWindowsOnAllDisplays();
+
+        if (allWindows != null && allWindows.size() > 0) {
+            return allWindows.get(displayId);
+        }
+        return null;
+    }
 }