[20/N] Desks: Implement IDesktopTaskListener add/remove/activation APIs

Forwards desk-level changes seen by the DesktopRepository to listeners
registered through IDesktopMode.

Flag: com.android.window.flags.enable_multiple_desktops_backend
Bug: 393961813
Test: atest WMShellUnitTests
Change-Id: Id551692e0b3dd35fe97ae7a35e32d3fbe1fe4eb7

3 files changed, 278 insertions, 3 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
index 926e0f287869..1a864b0cfc7a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
@@ -35,6 +35,9 @@ import com.android.wm.shell.shared.annotations.ShellMainThread
 import java.io.PrintWriter
 import java.util.concurrent.Executor
 import java.util.function.Consumer
+import kotlin.collections.component1
+import kotlin.collections.component2
+import kotlin.collections.forEach
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
 
@@ -116,6 +119,7 @@ class DesktopRepository(
         }
     }
 
+    private val deskChangeListeners = ArrayMap<DeskChangeListener, Executor>()
     private val activeTasksListeners = ArraySet<ActiveTasksListener>()
     private val visibleTasksListeners = ArrayMap<VisibleTasksListener, Executor>()
 
@@ -144,6 +148,11 @@ class DesktopRepository(
             SingleDesktopData()
         }
 
+    /** Adds a listener to be notified of updates about desk changes. */
+    fun addDeskChangeListener(listener: DeskChangeListener, executor: Executor) {
+        deskChangeListeners[listener] = executor
+    }
+
     /** Adds [activeTasksListener] to be notified of updates to active tasks. */
     fun addActiveTaskListener(activeTasksListener: ActiveTasksListener) {
         activeTasksListeners.add(activeTasksListener)
@@ -196,6 +205,11 @@ class DesktopRepository(
         return desktopExclusionRegion
     }
 
+    /** Removes the previously registered listener. */
+    fun removeDeskChangeListener(listener: DeskChangeListener) {
+        deskChangeListeners.remove(listener)
+    }
+
     /** Remove the previously registered [activeTasksListener] */
     fun removeActiveTasksListener(activeTasksListener: ActiveTasksListener) {
         activeTasksListeners.remove(activeTasksListener)
@@ -210,6 +224,9 @@ class DesktopRepository(
     fun addDesk(displayId: Int, deskId: Int) {
         logD("addDesk for displayId=%d and deskId=%d", displayId, deskId)
         desktopData.createDesk(displayId, deskId)
+        deskChangeListeners.forEach { (listener, executor) ->
+            executor.execute { listener.onDeskAdded(displayId = displayId, deskId = deskId) }
+        }
     }
 
     /** Returns the ids of the existing desks in the given display. */
@@ -229,12 +246,37 @@ class DesktopRepository(
     /** Sets the given desk as the active one in the given display. */
     fun setActiveDesk(displayId: Int, deskId: Int) {
         logD("setActiveDesk for displayId=%d and deskId=%d", displayId, deskId)
+        val oldActiveDeskId = desktopData.getActiveDesk(displayId)?.deskId ?: INVALID_DESK_ID
         desktopData.setActiveDesk(displayId = displayId, deskId = deskId)
+        deskChangeListeners.forEach { (listener, executor) ->
+            executor.execute {
+                listener.onActiveDeskChanged(
+                    displayId = displayId,
+                    newActiveDeskId = deskId,
+                    oldActiveDeskId = oldActiveDeskId,
+                )
+            }
+        }
     }
 
     /** Sets the given desk as inactive if it was active. */
     fun setDeskInactive(deskId: Int) {
+        val displayId = desktopData.getDisplayForDesk(deskId)
+        val activeDeskId = desktopData.getActiveDesk(displayId)?.deskId ?: INVALID_DESK_ID
+        if (activeDeskId == INVALID_DESK_ID || activeDeskId != deskId) {
+            // Desk wasn't active.
+            return
+        }
         desktopData.setDeskInactive(deskId)
+        deskChangeListeners.forEach { (listener, executor) ->
+            executor.execute {
+                listener.onActiveDeskChanged(
+                    displayId = displayId,
+                    newActiveDeskId = INVALID_DESK_ID,
+                    oldActiveDeskId = deskId,
+                )
+            }
+        }
     }
 
     /** Returns the id of the active desk in the given display, if any. */
@@ -898,8 +940,21 @@ class DesktopRepository(
                 ?: return emptySet<Int>().also {
                     logW("Could not find desk to remove: deskId=%d", deskId)
                 }
+        val wasActive = desktopData.getActiveDesk(desk.displayId)?.deskId == desk.deskId
         val activeTasks = ArraySet(desk.activeTasks)
         desktopData.remove(desk.deskId)
+        deskChangeListeners.forEach { (listener, executor) ->
+            executor.execute {
+                if (wasActive) {
+                    listener.onActiveDeskChanged(
+                        displayId = desk.displayId,
+                        newActiveDeskId = INVALID_DESK_ID,
+                        oldActiveDeskId = desk.deskId,
+                    )
+                }
+                listener.onDeskRemoved(displayId = desk.displayId, deskId = desk.deskId)
+            }
+        }
         return activeTasks
     }
 
@@ -1022,6 +1077,18 @@ class DesktopRepository(
             }
     }
 
+    /** Listens to changes of desks state. */
+    interface DeskChangeListener {
+        /** Called when a new desk is added to a display. */
+        fun onDeskAdded(displayId: Int, deskId: Int)
+
+        /** Called when a desk is removed from a display. */
+        fun onDeskRemoved(displayId: Int, deskId: Int)
+
+        /** Called when the active desk in a display has changed. */
+        fun onActiveDeskChanged(displayId: Int, newActiveDeskId: Int, oldActiveDeskId: Int)
+    }
+
     /** Listens to changes for active tasks in desktop mode. */
     interface ActiveTasksListener {
         fun onActiveTasksChanged(displayId: Int) {}
@@ -1282,6 +1349,8 @@ class DesktopRepository(
 
     companion object {
         private const val TAG = "DesktopRepository"
+
+        @VisibleForTesting const val INVALID_DESK_ID = -1
     }
 }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index 87d967427d88..85717f10ed34 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -97,6 +97,7 @@ import com.android.wm.shell.desktopmode.DesktopModeEventLogger.Companion.Unminim
 import com.android.wm.shell.desktopmode.DesktopModeUiEventLogger.DesktopUiEventEnum
 import com.android.wm.shell.desktopmode.DesktopModeVisualIndicator.DragStartState
 import com.android.wm.shell.desktopmode.DesktopModeVisualIndicator.IndicatorType
+import com.android.wm.shell.desktopmode.DesktopRepository.DeskChangeListener
 import com.android.wm.shell.desktopmode.DesktopRepository.VisibleTasksListener
 import com.android.wm.shell.desktopmode.DragToDesktopTransitionHandler.Companion.DRAG_TO_DESKTOP_FINISH_ANIM_DURATION_MS
 import com.android.wm.shell.desktopmode.DragToDesktopTransitionHandler.DragToDesktopStateListener
@@ -3446,7 +3447,47 @@ class DesktopTasksController(
         private lateinit var remoteListener:
             SingleInstanceRemoteListener<DesktopTasksController, IDesktopTaskListener>
 
-        private val listener: VisibleTasksListener =
+        private val deskChangeListener: DeskChangeListener =
+            object : DeskChangeListener {
+                override fun onDeskAdded(displayId: Int, deskId: Int) {
+                    ProtoLog.v(
+                        WM_SHELL_DESKTOP_MODE,
+                        "IDesktopModeImpl: onDeskAdded display=%d deskId=%d",
+                        displayId,
+                        deskId,
+                    )
+                    remoteListener.call { l -> l.onDeskAdded(displayId, deskId) }
+                }
+
+                override fun onDeskRemoved(displayId: Int, deskId: Int) {
+                    ProtoLog.v(
+                        WM_SHELL_DESKTOP_MODE,
+                        "IDesktopModeImpl: onDeskRemoved display=%d deskId=%d",
+                        displayId,
+                        deskId,
+                    )
+                    remoteListener.call { l -> l.onDeskRemoved(displayId, deskId) }
+                }
+
+                override fun onActiveDeskChanged(
+                    displayId: Int,
+                    newActiveDeskId: Int,
+                    oldActiveDeskId: Int,
+                ) {
+                    ProtoLog.v(
+                        WM_SHELL_DESKTOP_MODE,
+                        "IDesktopModeImpl: onActiveDeskChanged display=%d new=%d old=%d",
+                        displayId,
+                        newActiveDeskId,
+                        oldActiveDeskId,
+                    )
+                    remoteListener.call { l ->
+                        l.onActiveDeskChanged(displayId, newActiveDeskId, oldActiveDeskId)
+                    }
+                }
+            }
+
+        private val visibleTasksListener: VisibleTasksListener =
             object : VisibleTasksListener {
                 override fun onTasksVisibilityChanged(displayId: Int, visibleTasksCount: Int) {
                     ProtoLog.v(
@@ -3510,7 +3551,14 @@ class DesktopTasksController(
                     controller,
                     { c ->
                         run {
-                            c.taskRepository.addVisibleTasksListener(listener, c.mainExecutor)
+                            c.taskRepository.addDeskChangeListener(
+                                deskChangeListener,
+                                c.mainExecutor,
+                            )
+                            c.taskRepository.addVisibleTasksListener(
+                                visibleTasksListener,
+                                c.mainExecutor,
+                            )
                             c.taskbarDesktopTaskListener = taskbarDesktopTaskListener
                             c.desktopModeEnterExitTransitionListener =
                                 desktopModeEntryExitTransitionListener
@@ -3518,7 +3566,8 @@ class DesktopTasksController(
                     },
                     { c ->
                         run {
-                            c.taskRepository.removeVisibleTasksListener(listener)
+                            c.taskRepository.removeDeskChangeListener(deskChangeListener)
+                            c.taskRepository.removeVisibleTasksListener(visibleTasksListener)
                             c.taskbarDesktopTaskListener = null
                             c.desktopModeEnterExitTransitionListener = null
                         }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
index 9bff287e314a..a43b4dd00911 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
@@ -31,12 +31,14 @@ import com.android.window.flags.Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND
 import com.android.wm.shell.ShellTestCase
 import com.android.wm.shell.TestShellExecutor
 import com.android.wm.shell.common.ShellExecutor
+import com.android.wm.shell.desktopmode.DesktopRepository.Companion.INVALID_DESK_ID
 import com.android.wm.shell.desktopmode.persistence.Desktop
 import com.android.wm.shell.desktopmode.persistence.DesktopPersistentRepository
 import com.android.wm.shell.sysui.ShellInit
 import com.google.common.truth.Truth.assertThat
 import junit.framework.Assert.fail
 import kotlin.test.assertEquals
+import kotlin.test.assertNotNull
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -1430,6 +1432,161 @@ class DesktopRepositoryTest(flags: FlagsParameterization) : ShellTestCase() {
             )
     }
 
+    @Test
+    @EnableFlags(FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun addDesk_updatesListener() {
+        val listener = TestDeskChangeListener()
+        val executor = TestShellExecutor()
+        repo.addDeskChangeListener(listener, executor)
+
+        repo.addDesk(displayId = 0, deskId = 1)
+        executor.flushAll()
+
+        val lastAddition = assertNotNull(listener.lastAddition)
+        assertThat(lastAddition.displayId).isEqualTo(0)
+        assertThat(lastAddition.deskId).isEqualTo(1)
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun removeDesk_updatesListener() {
+        val listener = TestDeskChangeListener()
+        val executor = TestShellExecutor()
+        repo.addDeskChangeListener(listener, executor)
+        repo.addDesk(displayId = 0, deskId = 1)
+
+        repo.removeDesk(deskId = 1)
+        executor.flushAll()
+
+        val lastRemoval = assertNotNull(listener.lastRemoval)
+        assertThat(lastRemoval.displayId).isEqualTo(0)
+        assertThat(lastRemoval.deskId).isEqualTo(1)
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun removeDesk_didNotExist_doesNotUpdateListener() {
+        val listener = TestDeskChangeListener()
+        val executor = TestShellExecutor()
+        repo.addDeskChangeListener(listener, executor)
+        repo.addDesk(displayId = 0, deskId = 1)
+
+        repo.removeDesk(deskId = 2)
+        executor.flushAll()
+
+        assertThat(listener.lastRemoval).isNull()
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun removeDesk_wasActive_updatesActiveChangeListener() {
+        val listener = TestDeskChangeListener()
+        val executor = TestShellExecutor()
+        repo.addDeskChangeListener(listener, executor)
+        repo.addDesk(displayId = 0, deskId = 1)
+        repo.setActiveDesk(displayId = 0, deskId = 1)
+
+        repo.removeDesk(deskId = 1)
+        executor.flushAll()
+
+        val lastActivationChange = assertNotNull(listener.lastActivationChange)
+        assertThat(lastActivationChange.displayId).isEqualTo(0)
+        assertThat(lastActivationChange.oldActive).isEqualTo(1)
+        assertThat(lastActivationChange.newActive).isEqualTo(INVALID_DESK_ID)
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun setDeskActive_fromNoActive_updatesListener() {
+        val listener = TestDeskChangeListener()
+        val executor = TestShellExecutor()
+        repo.addDeskChangeListener(listener, executor)
+        repo.addDesk(displayId = 1, deskId = 1)
+
+        repo.setActiveDesk(displayId = 1, deskId = 1)
+        executor.flushAll()
+
+        val lastActivationChange = assertNotNull(listener.lastActivationChange)
+        assertThat(lastActivationChange.displayId).isEqualTo(1)
+        assertThat(lastActivationChange.oldActive).isEqualTo(INVALID_DESK_ID)
+        assertThat(lastActivationChange.newActive).isEqualTo(1)
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun setDeskActive_fromOtherActive_updatesListener() {
+        val listener = TestDeskChangeListener()
+        val executor = TestShellExecutor()
+        repo.addDeskChangeListener(listener, executor)
+        repo.addDesk(displayId = 1, deskId = 1)
+        repo.addDesk(displayId = 1, deskId = 2)
+        repo.setActiveDesk(displayId = 1, deskId = 1)
+
+        repo.setActiveDesk(displayId = 1, deskId = 2)
+        executor.flushAll()
+
+        val lastActivationChange = assertNotNull(listener.lastActivationChange)
+        assertThat(lastActivationChange.displayId).isEqualTo(1)
+        assertThat(lastActivationChange.oldActive).isEqualTo(1)
+        assertThat(lastActivationChange.newActive).isEqualTo(2)
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun setDeskInactive_updatesListener() {
+        val listener = TestDeskChangeListener()
+        val executor = TestShellExecutor()
+        repo.addDeskChangeListener(listener, executor)
+        repo.addDesk(displayId = 0, deskId = 1)
+        repo.setActiveDesk(displayId = 0, deskId = 1)
+
+        repo.setDeskInactive(deskId = 1)
+        executor.flushAll()
+
+        val lastActivationChange = assertNotNull(listener.lastActivationChange)
+        assertThat(lastActivationChange.displayId).isEqualTo(0)
+        assertThat(lastActivationChange.oldActive).isEqualTo(1)
+        assertThat(lastActivationChange.newActive).isEqualTo(INVALID_DESK_ID)
+    }
+
+    private class TestDeskChangeListener : DesktopRepository.DeskChangeListener {
+        var lastAddition: LastAddition? = null
+            private set
+
+        var lastRemoval: LastRemoval? = null
+            private set
+
+        var lastActivationChange: LastActivationChange? = null
+            private set
+
+        override fun onDeskAdded(displayId: Int, deskId: Int) {
+            lastAddition = LastAddition(displayId, deskId)
+        }
+
+        override fun onDeskRemoved(displayId: Int, deskId: Int) {
+            lastRemoval = LastRemoval(displayId, deskId)
+        }
+
+        override fun onActiveDeskChanged(
+            displayId: Int,
+            newActiveDeskId: Int,
+            oldActiveDeskId: Int,
+        ) {
+            lastActivationChange =
+                LastActivationChange(
+                    displayId = displayId,
+                    oldActive = oldActiveDeskId,
+                    newActive = newActiveDeskId,
+                )
+        }
+
+        data class LastAddition(val displayId: Int, val deskId: Int)
+
+        data class LastRemoval(val displayId: Int, val deskId: Int)
+
+        data class LastActivationChange(val displayId: Int, val oldActive: Int, val newActive: Int)
+    }
+
     class TestListener : DesktopRepository.ActiveTasksListener {
         var activeChangesOnDefaultDisplay = 0
         var activeChangesOnSecondaryDisplay = 0