Add View.viewModel utility

This allows for finer-grained structured concurrency for view-models
within a view-binder.

Flag: com.android.systemui.scene_container
Bug: 354269846
Test: atest SysUiViewModelTest
Change-Id: I9ec83f9efa85bf1baf930fc0e7b572fbd50cd4e7

5 files changed, 155 insertions, 44 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/lifecycle/SysUiViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/lifecycle/SysUiViewModelTest.kt
index d1f908dfc795..46b370fedf37 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/lifecycle/SysUiViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/lifecycle/SysUiViewModelTest.kt
@@ -16,16 +16,27 @@
 
 package com.android.systemui.lifecycle
 
+import android.view.View
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.util.Assert
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.stub
+import org.mockito.kotlin.verify
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
@@ -110,4 +121,45 @@ class SysUiViewModelTest : SysuiTestCase() {
 
         assertThat(isActive).isFalse()
     }
+
+    @Test
+    fun viewModel_viewBinder() = runTest {
+        Assert.setTestThread(Thread.currentThread())
+
+        val view: View = mock { on { isAttachedToWindow } doReturn false }
+        val viewModel = FakeViewModel()
+        backgroundScope.launch {
+            view.viewModel(
+                minWindowLifecycleState = WindowLifecycleState.ATTACHED,
+                factory = { viewModel },
+            ) {
+                awaitCancellation()
+            }
+        }
+        runCurrent()
+
+        assertThat(viewModel.isActivated).isFalse()
+
+        view.stub { on { isAttachedToWindow } doReturn true }
+        argumentCaptor<View.OnAttachStateChangeListener>()
+            .apply { verify(view).addOnAttachStateChangeListener(capture()) }
+            .allValues
+            .forEach { it.onViewAttachedToWindow(view) }
+        runCurrent()
+
+        assertThat(viewModel.isActivated).isTrue()
+    }
+}
+
+private class FakeViewModel : SysUiViewModel() {
+    var isActivated = false
+
+    override suspend fun onActivated() {
+        isActivated = true
+        try {
+            awaitCancellation()
+        } finally {
+            isActivated = false
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/lifecycle/RepeatWhenAttached.kt b/packages/SystemUI/src/com/android/systemui/lifecycle/RepeatWhenAttached.kt
index 661da6d2af13..c2b5d98699b4 100644
--- a/packages/SystemUI/src/com/android/systemui/lifecycle/RepeatWhenAttached.kt
+++ b/packages/SystemUI/src/com/android/systemui/lifecycle/RepeatWhenAttached.kt
@@ -227,13 +227,33 @@ private fun inferTraceSectionName(): String {
 }
 
 /**
+ * Runs the given [block] in a new coroutine when `this` [View]'s Window's [WindowLifecycleState] is
+ * at least at [state] (or immediately after calling this function if the window is already at least
+ * at [state]), automatically canceling the work when the window is no longer at least at that
+ * state.
+ *
+ * [block] may be run multiple times, running once per every time this` [View]'s Window's
+ * [WindowLifecycleState] becomes at least at [state].
+ */
+suspend fun View.repeatOnWindowLifecycle(
+    state: WindowLifecycleState,
+    block: suspend CoroutineScope.() -> Unit,
+): Nothing {
+    when (state) {
+        WindowLifecycleState.ATTACHED -> repeatWhenAttachedToWindow(block)
+        WindowLifecycleState.VISIBLE -> repeatWhenWindowIsVisible(block)
+        WindowLifecycleState.FOCUSED -> repeatWhenWindowHasFocus(block)
+    }
+}
+
+/**
  * Runs the given [block] every time the [View] becomes attached (or immediately after calling this
  * function, if the view was already attached), automatically canceling the work when the view
  * becomes detached.
  *
  * Only use from the main thread.
  *
- * The [block] may be run multiple times, running once per every time the view is attached.
+ * [block] may be run multiple times, running once per every time the view is attached.
  */
 @MainThread
 suspend fun View.repeatWhenAttachedToWindow(block: suspend CoroutineScope.() -> Unit): Nothing {
@@ -249,7 +269,7 @@ suspend fun View.repeatWhenAttachedToWindow(block: suspend CoroutineScope.() ->
  *
  * Only use from the main thread.
  *
- * The [block] may be run multiple times, running once per every time the window becomes visible.
+ * [block] may be run multiple times, running once per every time the window becomes visible.
  */
 @MainThread
 suspend fun View.repeatWhenWindowIsVisible(block: suspend CoroutineScope.() -> Unit): Nothing {
@@ -265,7 +285,7 @@ suspend fun View.repeatWhenWindowIsVisible(block: suspend CoroutineScope.() -> U
  *
  * Only use from the main thread.
  *
- * The [block] may be run multiple times, running once per every time the window is focused.
+ * [block] may be run multiple times, running once per every time the window is focused.
  */
 @MainThread
 suspend fun View.repeatWhenWindowHasFocus(block: suspend CoroutineScope.() -> Unit): Nothing {
@@ -274,6 +294,21 @@ suspend fun View.repeatWhenWindowHasFocus(block: suspend CoroutineScope.() -> Un
     awaitCancellation() // satisfies return type of Nothing
 }
 
+/** Lifecycle states for a [View]'s interaction with a [android.view.Window]. */
+enum class WindowLifecycleState {
+    /** Indicates that the [View] is attached to a [android.view.Window]. */
+    ATTACHED,
+    /**
+     * Indicates that the [View] is attached to a [android.view.Window], and the window is visible.
+     */
+    VISIBLE,
+    /**
+     * Indicates that the [View] is attached to a [android.view.Window], and the window is visible
+     * and focused.
+     */
+    FOCUSED
+}
+
 private val View.isAttached
     get() = conflatedCallbackFlow {
         val onAttachListener =
diff --git a/packages/SystemUI/src/com/android/systemui/lifecycle/SysUiViewModel.kt b/packages/SystemUI/src/com/android/systemui/lifecycle/SysUiViewModel.kt
index 0af5feaff3b2..77314813c34a 100644
--- a/packages/SystemUI/src/com/android/systemui/lifecycle/SysUiViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/lifecycle/SysUiViewModel.kt
@@ -16,9 +16,10 @@
 
 package com.android.systemui.lifecycle
 
+import android.view.View
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.LaunchedEffect
-import androidx.compose.runtime.remember
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
 
 /** Base class for all System UI view-models. */
 abstract class SysUiViewModel : SafeActivatable() {
@@ -37,8 +38,20 @@ abstract class SysUiViewModel : SafeActivatable() {
 fun <T : SysUiViewModel> rememberViewModel(
     key: Any = Unit,
     factory: () -> T,
-): T {
-    val instance = remember(key) { factory() }
-    LaunchedEffect(instance) { instance.activate() }
-    return instance
-}
+): T = rememberActivated(key, factory)
+
+/**
+ * Invokes [block] in a new coroutine with a new [SysUiViewModel] that is automatically activated
+ * whenever `this` [View]'s Window's [WindowLifecycleState] is at least at
+ * [minWindowLifecycleState], and is automatically canceled once that is no longer the case.
+ */
+suspend fun <T : SysUiViewModel> View.viewModel(
+    minWindowLifecycleState: WindowLifecycleState,
+    factory: () -> T,
+    block: suspend CoroutineScope.(T) -> Unit,
+): Nothing =
+    repeatOnWindowLifecycle(minWindowLifecycleState) {
+        val instance = factory()
+        launch { instance.activate() }
+        block(instance)
+    }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationScrollViewBinder.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationScrollViewBinder.kt
index fd08e898fce3..a30b8772c3d1 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationScrollViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationScrollViewBinder.kt
@@ -17,14 +17,14 @@
 package com.android.systemui.statusbar.notification.stack.ui.viewbinder
 
 import android.util.Log
-import androidx.lifecycle.Lifecycle
-import androidx.lifecycle.repeatOnLifecycle
 import com.android.systemui.common.ui.ConfigurationState
 import com.android.systemui.common.ui.view.onLayoutChanged
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.dump.DumpManager
+import com.android.systemui.lifecycle.WindowLifecycleState
 import com.android.systemui.lifecycle.repeatWhenAttached
+import com.android.systemui.lifecycle.viewModel
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.notification.stack.ui.view.NotificationScrollView
 import com.android.systemui.statusbar.notification.stack.ui.viewmodel.NotificationScrollViewModel
@@ -33,7 +33,6 @@ import com.android.systemui.util.kotlin.launchAndDispose
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.DisposableHandle
-import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.launch
@@ -46,7 +45,7 @@ constructor(
     dumpManager: DumpManager,
     @Main private val mainImmediateDispatcher: CoroutineDispatcher,
     private val view: NotificationScrollView,
-    private val viewModel: NotificationScrollViewModel,
+    private val viewModelFactory: NotificationScrollViewModel.Factory,
     private val configuration: ConfigurationState,
 ) : FlowDumperImpl(dumpManager) {
 
@@ -61,38 +60,42 @@ constructor(
     }
 
     fun bindWhileAttached(): DisposableHandle {
-        return view.asView().repeatWhenAttached(mainImmediateDispatcher) {
-            repeatOnLifecycle(Lifecycle.State.CREATED) { bind() }
-        }
+        return view.asView().repeatWhenAttached(mainImmediateDispatcher) { bind() }
     }
 
-    suspend fun bind() = coroutineScope {
-        launchAndDispose {
-            updateViewPosition()
-            view.asView().onLayoutChanged { updateViewPosition() }
-        }
+    suspend fun bind(): Nothing =
+        view.asView().viewModel(
+            minWindowLifecycleState = WindowLifecycleState.ATTACHED,
+            factory = viewModelFactory::create,
+        ) { viewModel ->
+            launchAndDispose {
+                updateViewPosition()
+                view.asView().onLayoutChanged { updateViewPosition() }
+            }
 
-        launch {
-            viewModel
-                .shadeScrimShape(cornerRadius = scrimRadius, viewLeftOffset = viewLeftOffset)
-                .collect { view.setScrimClippingShape(it) }
-        }
+            launch {
+                viewModel
+                    .shadeScrimShape(cornerRadius = scrimRadius, viewLeftOffset = viewLeftOffset)
+                    .collect { view.setScrimClippingShape(it) }
+            }
 
-        launch { viewModel.maxAlpha.collect { view.setMaxAlpha(it) } }
-        launch { viewModel.scrolledToTop.collect { view.setScrolledToTop(it) } }
-        launch { viewModel.expandFraction.collect { view.setExpandFraction(it.coerceIn(0f, 1f)) } }
-        launch { viewModel.isScrollable.collect { view.setScrollingEnabled(it) } }
-        launch { viewModel.isDozing.collect { isDozing -> view.setDozing(isDozing) } }
+            launch { viewModel.maxAlpha.collect { view.setMaxAlpha(it) } }
+            launch { viewModel.scrolledToTop.collect { view.setScrolledToTop(it) } }
+            launch {
+                viewModel.expandFraction.collect { view.setExpandFraction(it.coerceIn(0f, 1f)) }
+            }
+            launch { viewModel.isScrollable.collect { view.setScrollingEnabled(it) } }
+            launch { viewModel.isDozing.collect { isDozing -> view.setDozing(isDozing) } }
 
-        launchAndDispose {
-            view.setSyntheticScrollConsumer(viewModel.syntheticScrollConsumer)
-            view.setCurrentGestureOverscrollConsumer(viewModel.currentGestureOverscrollConsumer)
-            DisposableHandle {
-                view.setSyntheticScrollConsumer(null)
-                view.setCurrentGestureOverscrollConsumer(null)
+            launchAndDispose {
+                view.setSyntheticScrollConsumer(viewModel.syntheticScrollConsumer)
+                view.setCurrentGestureOverscrollConsumer(viewModel.currentGestureOverscrollConsumer)
+                DisposableHandle {
+                    view.setSyntheticScrollConsumer(null)
+                    view.setCurrentGestureOverscrollConsumer(null)
+                }
             }
         }
-    }
 
     /** flow of the scrim clipping radius */
     private val scrimRadius: Flow<Int>
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
index 2ba79a8612bb..428102530334 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
@@ -19,9 +19,9 @@ package com.android.systemui.statusbar.notification.stack.ui.viewmodel
 
 import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.compose.animation.scene.SceneKey
-import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
+import com.android.systemui.lifecycle.SysUiViewModel
 import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.SceneFamilies
@@ -33,9 +33,11 @@ import com.android.systemui.statusbar.notification.stack.shared.model.ShadeScrim
 import com.android.systemui.statusbar.notification.stack.shared.model.ShadeScrimShape
 import com.android.systemui.statusbar.notification.stack.ui.viewmodel.NotificationTransitionThresholds.EXPANSION_FOR_DELAYED_STACK_FADE_IN
 import com.android.systemui.statusbar.notification.stack.ui.viewmodel.NotificationTransitionThresholds.EXPANSION_FOR_MAX_SCRIM_ALPHA
+import com.android.systemui.util.kotlin.FlowDumper
 import com.android.systemui.util.kotlin.FlowDumperImpl
 import dagger.Lazy
-import javax.inject.Inject
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
@@ -43,9 +45,8 @@ import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
 
 /** ViewModel which represents the state of the NSSL/Controller in the world of flexiglass */
-@SysUISingleton
 class NotificationScrollViewModel
-@Inject
+@AssistedInject
 constructor(
     dumpManager: DumpManager,
     stackAppearanceInteractor: NotificationStackAppearanceInteractor,
@@ -54,7 +55,9 @@ constructor(
     // TODO(b/336364825) Remove Lazy when SceneContainerFlag is released -
     // while the flag is off, creating this object too early results in a crash
     keyguardInteractor: Lazy<KeyguardInteractor>,
-) : FlowDumperImpl(dumpManager) {
+) : FlowDumper by FlowDumperImpl(dumpManager, "NotificationScrollViewModel"),
+    SysUiViewModel() {
+
     /**
      * The expansion fraction of the notification stack. It should go from 0 to 1 when transitioning
      * from Gone to Shade scenes, and remain at 1 when in Lockscreen or Shade scenes and while
@@ -186,4 +189,9 @@ constructor(
             keyguardInteractor.get().isDozing.dumpWhileCollecting("isDozing")
         }
     }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(): NotificationScrollViewModel
+    }
 }