Merge "Add DisposableHandles and use it in a few places." into main
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/util/kotlin/DisposableHandlesTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/util/kotlin/DisposableHandlesTest.kt
new file mode 100644
index 0000000..a5ad3c3
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/util/kotlin/DisposableHandlesTest.kt
@@ -0,0 +1,68 @@
+/*
+ * Copyright (C) 2024 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.systemui.util.kotlin
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.DisposableHandle
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class DisposableHandlesTest : SysuiTestCase() {
+    @Test
+    fun disposeWorksOnce() {
+        var handleDisposeCount = 0
+        val underTest = DisposableHandles()
+
+        // Add a handle
+        underTest += DisposableHandle { handleDisposeCount++ }
+
+        // dispose() calls dispose() on children
+        underTest.dispose()
+        assertThat(handleDisposeCount).isEqualTo(1)
+
+        // Once disposed, children are not disposed again
+        underTest.dispose()
+        assertThat(handleDisposeCount).isEqualTo(1)
+    }
+
+    @Test
+    fun replaceCallsDispose() {
+        var handleDisposeCount1 = 0
+        var handleDisposeCount2 = 0
+        val underTest = DisposableHandles()
+        val handle1 = DisposableHandle { handleDisposeCount1++ }
+        val handle2 = DisposableHandle { handleDisposeCount2++ }
+
+        // First add handle1
+        underTest += handle1
+
+        // replace() calls dispose() on existing children
+        underTest.replaceAll(handle2)
+        assertThat(handleDisposeCount1).isEqualTo(1)
+        assertThat(handleDisposeCount2).isEqualTo(0)
+
+        // Once disposed, replaced children are not disposed again
+        underTest.dispose()
+        assertThat(handleDisposeCount1).isEqualTo(1)
+        assertThat(handleDisposeCount2).isEqualTo(1)
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt
index 7c76e6a..f60da0e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt
@@ -90,6 +90,7 @@
 import com.android.systemui.statusbar.phone.KeyguardBottomAreaView
 import com.android.systemui.statusbar.phone.ScreenOffAnimationController
 import com.android.systemui.temporarydisplay.chipbar.ChipbarCoordinator
+import com.android.systemui.util.kotlin.DisposableHandles
 import com.android.systemui.util.settings.SecureSettings
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedInject
@@ -173,7 +174,7 @@
     private lateinit var smallClockHostView: FrameLayout
     private var smartSpaceView: View? = null
 
-    private val disposables = mutableSetOf<DisposableHandle>()
+    private val disposables = DisposableHandles()
     private var isDestroyed = false
 
     private val shortcutsBindings = mutableSetOf<KeyguardQuickAffordanceViewBinder.Binding>()
@@ -183,7 +184,7 @@
 
     init {
         coroutineScope = CoroutineScope(applicationScope.coroutineContext + Job())
-        disposables.add(DisposableHandle { coroutineScope.cancel() })
+        disposables += DisposableHandle { coroutineScope.cancel() }
 
         if (keyguardBottomAreaRefactor()) {
             quickAffordancesCombinedViewModel.enablePreviewMode(
@@ -214,7 +215,7 @@
                     if (hostToken == null) null else InputTransferToken(hostToken),
                     "KeyguardPreviewRenderer"
                 )
-            disposables.add(DisposableHandle { host.release() })
+            disposables += DisposableHandle { host.release() }
         }
     }
 
@@ -284,7 +285,7 @@
     fun destroy() {
         isDestroyed = true
         lockscreenSmartspaceController.disconnect()
-        disposables.forEach { it.dispose() }
+        disposables.dispose()
         if (keyguardBottomAreaRefactor()) {
             shortcutsBindings.forEach { it.destroy() }
         }
@@ -372,7 +373,7 @@
     private fun setupKeyguardRootView(previewContext: Context, rootView: FrameLayout) {
         val keyguardRootView = KeyguardRootView(previewContext, null)
         if (!keyguardBottomAreaRefactor()) {
-            disposables.add(
+            disposables +=
                 KeyguardRootViewBinder.bind(
                     keyguardRootView,
                     keyguardRootViewModel,
@@ -387,7 +388,6 @@
                     null, // device entry haptics not required for preview mode
                     null, // falsing manager not required for preview mode
                 )
-            )
         }
         rootView.addView(
             keyguardRootView,
@@ -555,14 +555,12 @@
                     }
                 }
             clockRegistry.registerClockChangeListener(clockChangeListener)
-            disposables.add(
-                DisposableHandle {
-                    clockRegistry.unregisterClockChangeListener(clockChangeListener)
-                }
-            )
+            disposables += DisposableHandle {
+                clockRegistry.unregisterClockChangeListener(clockChangeListener)
+            }
 
             clockController.registerListeners(parentView)
-            disposables.add(DisposableHandle { clockController.unregisterListeners() })
+            disposables += DisposableHandle { clockController.unregisterListeners() }
         }
 
         val receiver =
@@ -581,7 +579,7 @@
                 addAction(Intent.ACTION_TIME_CHANGED)
             },
         )
-        disposables.add(DisposableHandle { broadcastDispatcher.unregisterReceiver(receiver) })
+        disposables += DisposableHandle { broadcastDispatcher.unregisterReceiver(receiver) }
 
         if (!migrateClocksToBlueprint()) {
             val layoutChangeListener =
@@ -602,9 +600,9 @@
                     }
                 }
             parentView.addOnLayoutChangeListener(layoutChangeListener)
-            disposables.add(
-                DisposableHandle { parentView.removeOnLayoutChangeListener(layoutChangeListener) }
-            )
+            disposables += DisposableHandle {
+                parentView.removeOnLayoutChangeListener(layoutChangeListener)
+            }
         }
 
         onClockChanged()
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/NotificationStackScrollLayoutSection.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/NotificationStackScrollLayoutSection.kt
index e73fa72..8323502 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/NotificationStackScrollLayoutSection.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/NotificationStackScrollLayoutSection.kt
@@ -35,7 +35,7 @@
 import com.android.systemui.statusbar.notification.stack.ui.viewbinder.NotificationStackViewBinder
 import com.android.systemui.statusbar.notification.stack.ui.viewbinder.SharedNotificationContainerBinder
 import com.android.systemui.statusbar.notification.stack.ui.viewmodel.SharedNotificationContainerViewModel
-import kotlinx.coroutines.DisposableHandle
+import com.android.systemui.util.kotlin.DisposableHandles
 
 abstract class NotificationStackScrollLayoutSection
 constructor(
@@ -48,7 +48,7 @@
     private val notificationStackViewBinder: NotificationStackViewBinder,
 ) : KeyguardSection() {
     private val placeHolderId = R.id.nssl_placeholder
-    private val disposableHandles: MutableList<DisposableHandle> = mutableListOf()
+    private val disposableHandles = DisposableHandles()
 
     /**
      * Align the notification placeholder bottom to the top of either the lock icon or the ambient
@@ -94,26 +94,20 @@
             return
         }
 
-        disposeHandles()
-        disposableHandles.add(
+        disposableHandles.dispose()
+        disposableHandles +=
             sharedNotificationContainerBinder.bind(
                 sharedNotificationContainer,
                 sharedNotificationContainerViewModel,
             )
-        )
 
         if (sceneContainerFlags.isEnabled()) {
-            disposableHandles.add(notificationStackViewBinder.bindWhileAttached())
+            disposableHandles += notificationStackViewBinder.bindWhileAttached()
         }
     }
 
     override fun removeViews(constraintLayout: ConstraintLayout) {
-        disposeHandles()
+        disposableHandles.dispose()
         constraintLayout.removeView(placeHolderId)
     }
-
-    private fun disposeHandles() {
-        disposableHandles.forEach { it.dispose() }
-        disposableHandles.clear()
-    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/SharedNotificationContainerBinder.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/SharedNotificationContainerBinder.kt
index 3352079..6db6719 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/SharedNotificationContainerBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/SharedNotificationContainerBinder.kt
@@ -31,6 +31,7 @@
 import com.android.systemui.statusbar.notification.stack.NotificationStackSizeCalculator
 import com.android.systemui.statusbar.notification.stack.ui.view.SharedNotificationContainer
 import com.android.systemui.statusbar.notification.stack.ui.viewmodel.SharedNotificationContainerViewModel
+import com.android.systemui.util.kotlin.DisposableHandles
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.DisposableHandle
@@ -54,7 +55,9 @@
         view: SharedNotificationContainer,
         viewModel: SharedNotificationContainerViewModel,
     ): DisposableHandle {
-        val disposableHandle =
+        val disposables = DisposableHandles()
+
+        disposables +=
             view.repeatWhenAttached {
                 repeatOnLifecycle(Lifecycle.State.CREATED) {
                     launch {
@@ -77,24 +80,6 @@
                 }
             }
 
-        // Required to capture keyguard media changes and ensure the notification count is correct
-        val layoutChangeListener =
-            object : View.OnLayoutChangeListener {
-                override fun onLayoutChange(
-                    view: View,
-                    left: Int,
-                    top: Int,
-                    right: Int,
-                    bottom: Int,
-                    oldLeft: Int,
-                    oldTop: Int,
-                    oldRight: Int,
-                    oldBottom: Int
-                ) {
-                    viewModel.notificationStackChanged()
-                }
-            }
-
         val burnInParams = MutableStateFlow(BurnInParameters())
         val viewState =
             ViewStateAccessor(
@@ -105,7 +90,7 @@
          * For animation sensitive coroutines, immediately run just like applicationScope does
          * instead of doing a post() to the main thread. This extra delay can cause visible jitter.
          */
-        val disposableHandleMainImmediate =
+        disposables +=
             view.repeatWhenAttached(mainImmediateDispatcher) {
                 repeatOnLifecycle(Lifecycle.State.CREATED) {
                     launch {
@@ -172,7 +157,8 @@
                 }
             }
 
-        controller.setOnHeightChangedRunnable(Runnable { viewModel.notificationStackChanged() })
+        controller.setOnHeightChangedRunnable { viewModel.notificationStackChanged() }
+        disposables += DisposableHandle { controller.setOnHeightChangedRunnable(null) }
 
         view.setOnApplyWindowInsetsListener { v: View, insets: WindowInsets ->
             val insetTypes = WindowInsets.Type.systemBars() or WindowInsets.Type.displayCutout()
@@ -181,16 +167,16 @@
             }
             insets
         }
-        view.addOnLayoutChangeListener(layoutChangeListener)
+        disposables += DisposableHandle { view.setOnApplyWindowInsetsListener(null) }
 
-        return object : DisposableHandle {
-            override fun dispose() {
-                disposableHandle.dispose()
-                disposableHandleMainImmediate.dispose()
-                controller.setOnHeightChangedRunnable(null)
-                view.setOnApplyWindowInsetsListener(null)
-                view.removeOnLayoutChangeListener(layoutChangeListener)
+        // Required to capture keyguard media changes and ensure the notification count is correct
+        val layoutChangeListener =
+            View.OnLayoutChangeListener { _, _, _, _, _, _, _, _, _ ->
+                viewModel.notificationStackChanged()
             }
-        }
+        view.addOnLayoutChangeListener(layoutChangeListener)
+        disposables += DisposableHandle { view.removeOnLayoutChangeListener(layoutChangeListener) }
+
+        return disposables
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/util/kotlin/DisposableHandles.kt b/packages/SystemUI/src/com/android/systemui/util/kotlin/DisposableHandles.kt
new file mode 100644
index 0000000..de036ea
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/util/kotlin/DisposableHandles.kt
@@ -0,0 +1,51 @@
+/*
+ * Copyright (C) 2024 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.systemui.util.kotlin
+
+import kotlinx.coroutines.DisposableHandle
+
+/** A mutable collection of [DisposableHandle] objects that is itself a [DisposableHandle] */
+class DisposableHandles : DisposableHandle {
+    private val handles = mutableListOf<DisposableHandle>()
+
+    /** Add the provided handles to this collection. */
+    fun add(vararg handles: DisposableHandle) {
+        this.handles.addAll(handles)
+    }
+
+    /** Same as [add] */
+    operator fun plusAssign(handle: DisposableHandle) {
+        this.handles.add(handle)
+    }
+
+    /** Same as [add] */
+    operator fun plusAssign(handles: Iterable<DisposableHandle>) {
+        this.handles.addAll(handles)
+    }
+
+    /** [dispose] the current contents, then [add] the provided [handles] */
+    fun replaceAll(vararg handles: DisposableHandle) {
+        dispose()
+        add(*handles)
+    }
+
+    /** Dispose of all added handles and empty this collection. */
+    override fun dispose() {
+        handles.forEach { it.dispose() }
+        handles.clear()
+    }
+}