Recalculate max notifications on lockscreen

NSSL has a listener that will be invoked when a notification changes,
and can be used to recalculate the number of notifications to show on
lockscreen. Use this signal to emit an event for the view model to
use.

Fixes: 296606746
Test: atest SharedNotificationContainerViewModelTest
Test: use Notify apk to add delayed notifications, ensure it stays to
the bounds
Flag: LEGACY MIGRATE_NSSL DISABLED

Change-Id: I7b59b56d167d2dec12761bda80fbea27f74bf239
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultNotificationStackScrollLayoutSection.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultNotificationStackScrollLayoutSection.kt
index 078feff..c2aedca 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultNotificationStackScrollLayoutSection.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultNotificationStackScrollLayoutSection.kt
@@ -30,6 +30,7 @@
 import com.android.systemui.res.R
 import com.android.systemui.shade.NotificationPanelView
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayoutController
+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 javax.inject.Inject
@@ -44,6 +45,7 @@
     sharedNotificationContainer: SharedNotificationContainer,
     sharedNotificationContainerViewModel: SharedNotificationContainerViewModel,
     controller: NotificationStackScrollLayoutController,
+    notificationStackSizeCalculator: NotificationStackSizeCalculator,
     private val smartspaceViewModel: KeyguardSmartspaceViewModel,
 ) :
     NotificationStackScrollLayoutSection(
@@ -53,6 +55,7 @@
         sharedNotificationContainer,
         sharedNotificationContainerViewModel,
         controller,
+        notificationStackSizeCalculator,
     ) {
     override fun applyConstraints(constraintSet: ConstraintSet) {
         if (!featureFlags.isEnabled(Flags.MIGRATE_NSSL)) {
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 00966f2..ea2bdf7 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
@@ -27,6 +27,7 @@
 import com.android.systemui.res.R
 import com.android.systemui.shade.NotificationPanelView
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayoutController
+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.viewbinder.SharedNotificationContainerBinder
 import com.android.systemui.statusbar.notification.stack.ui.viewmodel.SharedNotificationContainerViewModel
@@ -40,6 +41,7 @@
     private val sharedNotificationContainer: SharedNotificationContainer,
     private val sharedNotificationContainerViewModel: SharedNotificationContainerViewModel,
     private val controller: NotificationStackScrollLayoutController,
+    private val notificationStackSizeCalculator: NotificationStackSizeCalculator,
 ) : KeyguardSection() {
     private val placeHolderId = R.id.nssl_placeholder
     private var disposableHandle: DisposableHandle? = null
@@ -69,6 +71,7 @@
                 sharedNotificationContainer,
                 sharedNotificationContainerViewModel,
                 controller,
+                notificationStackSizeCalculator,
             )
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/SplitShadeNotificationStackScrollLayoutSection.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/SplitShadeNotificationStackScrollLayoutSection.kt
index bf95c77..dc2ad8d 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/SplitShadeNotificationStackScrollLayoutSection.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/SplitShadeNotificationStackScrollLayoutSection.kt
@@ -30,6 +30,7 @@
 import com.android.systemui.res.R
 import com.android.systemui.shade.NotificationPanelView
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayoutController
+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 javax.inject.Inject
@@ -44,6 +45,7 @@
     sharedNotificationContainer: SharedNotificationContainer,
     sharedNotificationContainerViewModel: SharedNotificationContainerViewModel,
     controller: NotificationStackScrollLayoutController,
+    notificationStackSizeCalculator: NotificationStackSizeCalculator,
     private val smartspaceViewModel: KeyguardSmartspaceViewModel,
 ) :
     NotificationStackScrollLayoutSection(
@@ -53,6 +55,7 @@
         sharedNotificationContainer,
         sharedNotificationContainerViewModel,
         controller,
+        notificationStackSizeCalculator,
     ) {
     override fun applyConstraints(constraintSet: ConstraintSet) {
         if (!featureFlags.isEnabled(Flags.MIGRATE_NSSL)) {
diff --git a/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java b/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
index 285cb5a..e9c930a 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
@@ -1483,16 +1483,16 @@
     }
 
     private void updateMaxDisplayedNotifications(boolean recompute) {
+        if (mFeatureFlags.isEnabled(Flags.MIGRATE_NSSL)) {
+            return;
+        }
+
         if (recompute) {
             setMaxDisplayedNotifications(Math.max(computeMaxKeyguardNotifications(), 1));
         } else {
             if (SPEW_LOGCAT) Log.d(TAG, "Skipping computeMaxKeyguardNotifications() by request");
         }
 
-        if (mFeatureFlags.isEnabled(Flags.MIGRATE_NSSL)) {
-            return;
-        }
-
         if (isKeyguardShowing() && !mKeyguardBypassController.getBypassEnabled()) {
             mNotificationStackScrollLayoutController.setMaxDisplayedNotifications(
                     mMaxAllowedKeyguardNotifications);
diff --git a/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt b/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt
index e2e4556..8bab669 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt
@@ -65,6 +65,10 @@
      */
     @Deprecated("Use ShadeInteractor instead") val legacyShadeTracking: StateFlow<Boolean>
 
+    /** Specifically tracks the user expanding the shade on the lockscreen only */
+    @Deprecated("Use ShadeInteractor.isUserInteractingWithShade instead")
+    val legacyLockscreenShadeTracking: MutableStateFlow<Boolean>
+
     /**
      * QuickSettingsController.mTracking as a flow. "Tracking" means that the user is moving quick
      * settings up or down with a pointer. Going forward, this concept will be replaced by checks
@@ -106,6 +110,9 @@
     /** Sets whether the user is moving the shade with a pointer */
     fun setLegacyShadeTracking(tracking: Boolean)
 
+    /** Sets whether the user is moving the shade with a pointer, on lockscreen only */
+    fun setLegacyLockscreenShadeTracking(tracking: Boolean)
+
     /** Amount shade has expanded with regard to the UDFPS location */
     val udfpsTransitionToFullShadeProgress: StateFlow<Float>
 
@@ -177,6 +184,8 @@
     @Deprecated("Use ShadeInteractor instead")
     override val legacyShadeTracking: StateFlow<Boolean> = _legacyShadeTracking.asStateFlow()
 
+    override val legacyLockscreenShadeTracking = MutableStateFlow(false)
+
     private val _legacyQsTracking = MutableStateFlow(false)
     @Deprecated("Use ShadeInteractor instead")
     override val legacyQsTracking: StateFlow<Boolean> = _legacyQsTracking.asStateFlow()
@@ -212,6 +221,11 @@
         _legacyShadeTracking.value = tracking
     }
 
+    @Deprecated("Should only be called by NPVC and tests")
+    override fun setLegacyLockscreenShadeTracking(tracking: Boolean) {
+        legacyLockscreenShadeTracking.value = tracking
+    }
+
     override fun setQsExpansion(qsExpansion: Float) {
         _qsExpansion.value = qsExpansion
     }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractor.kt b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractor.kt
index b2ffeb3..d687ef6 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractor.kt
@@ -72,7 +72,7 @@
     userSetupRepository: UserSetupRepository,
     userSwitcherInteractor: UserSwitcherInteractor,
     sharedNotificationContainerInteractor: SharedNotificationContainerInteractor,
-    repository: ShadeRepository,
+    private val repository: ShadeRepository,
 ) {
     /** Emits true if the shade is currently allowed and false otherwise. */
     val isShadeEnabled: StateFlow<Boolean> =
@@ -185,7 +185,15 @@
         if (sceneContainerFlags.isEnabled()) {
             sceneBasedInteracting(sceneInteractorProvider.get(), SceneKey.Shade)
         } else {
-            userInteractingFlow(repository.legacyShadeTracking, repository.legacyShadeExpansion)
+            combine(
+                userInteractingFlow(
+                    repository.legacyShadeTracking,
+                    repository.legacyShadeExpansion
+                ),
+                repository.legacyLockscreenShadeTracking
+            ) { legacyShadeTracking, legacyLockscreenShadeTracking ->
+                legacyShadeTracking || legacyLockscreenShadeTracking
+            }
         }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/LockscreenShadeTransitionController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/LockscreenShadeTransitionController.kt
index bf722af..2e3f3f8 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/LockscreenShadeTransitionController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/LockscreenShadeTransitionController.kt
@@ -340,6 +340,7 @@
         )
         nsslController.resetScrollPosition()
         nsslController.resetCheckSnoozeLeavebehind()
+        shadeRepository.setLegacyLockscreenShadeTracking(false)
         setDragDownAmountAnimated(0f)
     }
 
@@ -366,6 +367,7 @@
                 cancel()
             }
         }
+        shadeRepository.setLegacyLockscreenShadeTracking(true)
     }
 
     /** Do we need a falsing check currently? */
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
index a0ffba3..14ec08f35 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
@@ -253,6 +253,7 @@
     private NotificationLogger.OnChildLocationsChangedListener mListener;
     private OnOverscrollTopChangedListener mOverscrollTopChangedListener;
     private ExpandableView.OnHeightChangedListener mOnHeightChangedListener;
+    private Runnable mOnHeightChangedRunnable;
     private OnEmptySpaceClickListener mOnEmptySpaceClickListener;
     private boolean mNeedsAnimation;
     private boolean mTopPaddingNeedsAnimation;
@@ -1121,6 +1122,10 @@
         if (mOnHeightChangedListener != null) {
             mOnHeightChangedListener.onHeightChanged(view, needsAnimation);
         }
+
+        if (mOnHeightChangedRunnable != null) {
+            mOnHeightChangedRunnable.run();
+        }
     }
 
     public boolean isPulseExpanding() {
@@ -4252,6 +4257,10 @@
         this.mOnHeightChangedListener = onHeightChangedListener;
     }
 
+    void setOnHeightChangedRunnable(Runnable r) {
+        this.mOnHeightChangedRunnable = r;
+    }
+
     void onChildAnimationFinished() {
         setAnimationRunning(false);
         requestChildrenUpdate();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
index 99b3a00..2cf0c26 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
@@ -20,6 +20,7 @@
 import static android.service.notification.NotificationStats.DISMISS_SENTIMENT_NEUTRAL;
 
 import static com.android.app.animation.Interpolators.STANDARD;
+
 import static com.android.internal.jank.InteractionJankMonitor.CUJ_NOTIFICATION_SHADE_SCROLL_FLING;
 import static com.android.systemui.Dependency.ALLOW_NOTIFICATION_LONG_PRESS_NAME;
 import static com.android.systemui.statusbar.StatusBarState.KEYGUARD;
@@ -34,7 +35,6 @@
 
 import android.animation.ObjectAnimator;
 import android.content.res.Configuration;
-import android.content.res.Resources;
 import android.graphics.Point;
 import android.os.Trace;
 import android.os.UserHandle;
@@ -65,9 +65,7 @@
 import com.android.systemui.bouncer.domain.interactor.PrimaryBouncerInteractor;
 import com.android.systemui.classifier.Classifier;
 import com.android.systemui.classifier.FalsingCollector;
-import com.android.systemui.common.ui.ConfigurationState;
 import com.android.systemui.dagger.SysUISingleton;
-import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.flags.FeatureFlagsClassic;
 import com.android.systemui.flags.Flags;
@@ -94,7 +92,6 @@
 import com.android.systemui.statusbar.SysuiStatusBarStateController;
 import com.android.systemui.statusbar.notification.DynamicPrivacyController;
 import com.android.systemui.statusbar.notification.LaunchAnimationParameters;
-import com.android.systemui.statusbar.notification.NotifPipelineFlags;
 import com.android.systemui.statusbar.notification.NotificationActivityStarter;
 import com.android.systemui.statusbar.notification.collection.NotifCollection;
 import com.android.systemui.statusbar.notification.collection.NotifPipeline;
@@ -112,7 +109,6 @@
 import com.android.systemui.statusbar.notification.collection.render.SectionHeaderController;
 import com.android.systemui.statusbar.notification.dagger.SilentHeader;
 import com.android.systemui.statusbar.notification.domain.interactor.SeenNotificationsInteractor;
-import com.android.systemui.statusbar.notification.icon.ui.viewbinder.ShelfNotificationIconViewStore;
 import com.android.systemui.statusbar.notification.init.NotificationsController;
 import com.android.systemui.statusbar.notification.logging.NotificationLogger;
 import com.android.systemui.statusbar.notification.row.ActivatableNotificationView;
@@ -122,11 +118,9 @@
 import com.android.systemui.statusbar.notification.row.NotificationGutsManager;
 import com.android.systemui.statusbar.notification.row.NotificationSnooze;
 import com.android.systemui.statusbar.notification.stack.ui.viewbinder.NotificationListViewBinder;
-import com.android.systemui.statusbar.notification.stack.ui.viewmodel.NotificationListViewModel;
 import com.android.systemui.statusbar.phone.HeadsUpAppearanceController;
 import com.android.systemui.statusbar.phone.HeadsUpTouchHelper;
 import com.android.systemui.statusbar.phone.KeyguardBypassController;
-import com.android.systemui.statusbar.phone.NotificationIconAreaController;
 import com.android.systemui.statusbar.phone.ScrimController;
 import com.android.systemui.statusbar.policy.ConfigurationController;
 import com.android.systemui.statusbar.policy.ConfigurationController.ConfigurationListener;
@@ -960,6 +954,13 @@
         mView.setOnHeightChangedListener(listener);
     }
 
+    /**
+     * Invoked in addition to {@see #setOnHeightChangedListener}
+     */
+    public void setOnHeightChangedRunnable(Runnable r) {
+        mView.setOnHeightChangedRunnable(r);
+    }
+
     public void setOverscrollTopChangedListener(
             OnOverscrollTopChangedListener listener) {
         mView.setOverscrollTopChangedListener(listener);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/domain/interactor/SharedNotificationContainerInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/domain/interactor/SharedNotificationContainerInteractor.kt
index 57cea5d..eb1c17a 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/domain/interactor/SharedNotificationContainerInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/domain/interactor/SharedNotificationContainerInteractor.kt
@@ -18,13 +18,15 @@
 package com.android.systemui.statusbar.notification.stack.domain.interactor
 
 import android.content.Context
-import com.android.systemui.res.R
 import com.android.systemui.common.ui.data.repository.ConfigurationRepository
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.res.R
 import com.android.systemui.statusbar.policy.SplitShadeStateController
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.map
@@ -43,6 +45,10 @@
     private val _topPosition = MutableStateFlow(0f)
     val topPosition = _topPosition.asStateFlow()
 
+    private val _notificationStackChanged = MutableSharedFlow<Unit>(extraBufferCapacity = 1)
+    /** An internal modification was made to notifications */
+    val notificationStackChanged = _notificationStackChanged.asSharedFlow()
+
     val configurationBasedDimensions: Flow<ConfigurationBasedDimensions> =
         configurationRepository.onAnyConfigurationChange
             .onStart { emit(Unit) }
@@ -72,6 +78,11 @@
         _topPosition.value = top
     }
 
+    /** An internal modification was made to notifications */
+    fun notificationStackChanged() {
+        _notificationStackChanged.tryEmit(Unit)
+    }
+
     data class ConfigurationBasedDimensions(
         val useSplitShade: Boolean,
         val useLargeScreenHeader: Boolean,
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 a1a0cca..0ff1bec 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
@@ -20,6 +20,7 @@
 import androidx.lifecycle.repeatOnLifecycle
 import com.android.systemui.lifecycle.repeatWhenAttached
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayoutController
+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 kotlinx.coroutines.DisposableHandle
@@ -33,6 +34,7 @@
         view: SharedNotificationContainer,
         viewModel: SharedNotificationContainerViewModel,
         controller: NotificationStackScrollLayoutController,
+        notificationStackSizeCalculator: NotificationStackSizeCalculator,
     ): DisposableHandle {
         val disposableHandle =
             view.repeatWhenAttached {
@@ -54,9 +56,16 @@
                     }
 
                     launch {
-                        viewModel.maxNotifications.collect {
-                            controller.setMaxDisplayedNotifications(it)
-                        }
+                        viewModel
+                            .getMaxNotifications { space ->
+                                notificationStackSizeCalculator.computeMaxKeyguardNotifications(
+                                    controller.getView(),
+                                    space,
+                                    0f, // Vertical space for shelf is already accounted for
+                                    controller.getShelfHeight().toFloat(),
+                                )
+                            }
+                            .collect { controller.setMaxDisplayedNotifications(it) }
                     }
 
                     launch {
@@ -70,9 +79,12 @@
                 }
             }
 
+        controller.setOnHeightChangedRunnable(Runnable { viewModel.notificationStackChanged() })
+
         return object : DisposableHandle {
             override fun dispose() {
                 disposableHandle.dispose()
+                controller.setOnHeightChangedRunnable(null)
             }
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
index b86b5dc..d6b6f75 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
@@ -22,28 +22,26 @@
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
-import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayoutController
-import com.android.systemui.statusbar.notification.stack.NotificationStackSizeCalculator
 import com.android.systemui.statusbar.notification.stack.domain.interactor.SharedNotificationContainerInteractor
 import com.android.systemui.util.kotlin.sample
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.combineTransform
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.merge
 import kotlinx.coroutines.flow.onStart
 
 /** View-model for the shared notification container, used by both the shade and keyguard spaces */
 class SharedNotificationContainerViewModel
 @Inject
 constructor(
-    interactor: SharedNotificationContainerInteractor,
+    private val interactor: SharedNotificationContainerInteractor,
     keyguardInteractor: KeyguardInteractor,
     keyguardTransitionInteractor: KeyguardTransitionInteractor,
-    notificationStackSizeCalculator: NotificationStackSizeCalculator,
-    controller: NotificationStackScrollLayoutController,
-    shadeInteractor: ShadeInteractor,
+    private val shadeInteractor: ShadeInteractor,
 ) {
     private val statesForConstrainedNotifications =
         setOf(
@@ -151,24 +149,46 @@
      * When on keyguard, there is limited space to display notifications so calculate how many could
      * be shown. Otherwise, there is no limit since the vertical space will be scrollable.
      *
-     * TODO: b/296606746 - Need to rerun logic when notifs change
+     * When expanding or when the user is interacting with the shade, keep the count stable; do not
+     * emit a value.
      */
-    val maxNotifications: Flow<Int> =
-        combine(isOnLockscreen, shadeInteractor.shadeExpansion, position) {
-            onLockscreen,
-            shadeExpansion,
-            positionInfo ->
-            if (onLockscreen && shadeExpansion < 1f) {
-                notificationStackSizeCalculator.computeMaxKeyguardNotifications(
-                    controller.getView(),
-                    positionInfo.bottom - positionInfo.top,
-                    0f, // Vertical space for shelf is already accounted for
-                    controller.getShelfHeight().toFloat(),
-                )
-            } else {
-                -1 // No limit
+    fun getMaxNotifications(calculateSpace: (Float) -> Int): Flow<Int> {
+        // When to limit notifications: on lockscreen with an unexpanded shade. Also, recalculate
+        // when the notification stack has changed internally
+        val limitedNotifications =
+            combineTransform(
+                isOnLockscreen,
+                position,
+                shadeInteractor.shadeExpansion,
+                interactor.notificationStackChanged.onStart { emit(Unit) },
+            ) { isOnLockscreen, position, shadeExpansion, _ ->
+                if (isOnLockscreen && shadeExpansion == 0f) {
+                    emit(calculateSpace(position.bottom - position.top))
+                }
             }
-        }
+
+        // When to show unlimited notifications: When the shade is fully expanded and the user is
+        // not actively dragging the shade
+        val unlimitedNotifications =
+            combineTransform(
+                shadeInteractor.shadeExpansion,
+                shadeInteractor.isUserInteracting,
+            ) { shadeExpansion, isUserInteracting ->
+                if (shadeExpansion == 1f && !isUserInteracting) {
+                    emit(-1)
+                }
+            }
+
+        return merge(
+                limitedNotifications,
+                unlimitedNotifications,
+            )
+            .distinctUntilChanged()
+    }
+
+    fun notificationStackChanged() {
+        interactor.notificationStackChanged()
+    }
 
     data class ConfigurationBasedDimensions(
         val marginStart: Int,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/data/repository/ShadeRepositoryImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/shade/data/repository/ShadeRepositoryImplTest.kt
index e920687..20b19fd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/data/repository/ShadeRepositoryImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/data/repository/ShadeRepositoryImplTest.kt
@@ -158,6 +158,15 @@
         }
 
     @Test
+    fun updateLegacyLockscreenShadeTracking() =
+        testScope.runTest {
+            assertThat(underTest.legacyLockscreenShadeTracking.value).isEqualTo(false)
+
+            underTest.setLegacyLockscreenShadeTracking(true)
+            assertThat(underTest.legacyLockscreenShadeTracking.value).isEqualTo(true)
+        }
+
+    @Test
     fun updateLegacyQsTracking() =
         testScope.runTest {
             assertThat(underTest.legacyQsTracking.value).isEqualTo(false)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
index 6203531..e91d6d7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
@@ -83,6 +83,7 @@
 import com.android.systemui.statusbar.notification.footer.ui.view.FooterView;
 import com.android.systemui.statusbar.notification.init.NotificationsController;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
+import com.android.systemui.statusbar.notification.row.ExpandableView;
 import com.android.systemui.statusbar.phone.KeyguardBypassController;
 import com.android.systemui.statusbar.phone.ScreenOffAnimationController;
 import com.android.systemui.statusbar.phone.StatusBarKeyguardViewManager;
@@ -906,6 +907,20 @@
         assertEquals(bottomImeInset, mStackScrollerInternal.mBottomInset);
     }
 
+    @Test
+    public void testSetMaxDisplayedNotifications_notifiesListeners() {
+        ExpandableView.OnHeightChangedListener listener =
+                mock(ExpandableView.OnHeightChangedListener.class);
+        Runnable runnable = mock(Runnable.class);
+        mStackScroller.setOnHeightChangedListener(listener);
+        mStackScroller.setOnHeightChangedRunnable(runnable);
+
+        mStackScroller.setMaxDisplayedNotifications(50);
+
+        verify(listener).onHeightChanged(mNotificationShelf, false);
+        verify(runnable).run();
+    }
+
     private void setBarStateForTest(int state) {
         // Can't inject this through the listener or we end up on the actual implementation
         // rather than the mock because the spy just coppied the anonymous inner /shruggie.
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
index 22553df..db8f217 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
@@ -40,13 +40,8 @@
 import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.res.R
 import com.android.systemui.shade.data.repository.FakeShadeRepository
-import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayoutController
-import com.android.systemui.statusbar.notification.stack.NotificationStackSizeCalculator
 import com.android.systemui.statusbar.notification.stack.domain.interactor.SharedNotificationContainerInteractor
 import com.android.systemui.user.domain.UserDomainLayerModule
-import com.android.systemui.util.mockito.any
-import com.android.systemui.util.mockito.mock
-import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import dagger.BindsInstance
 import dagger.Component
@@ -84,25 +79,13 @@
         }
     }
 
-    private val notificationStackSizeCalculator: NotificationStackSizeCalculator = mock()
-    private val notificationStackScrollLayoutController: NotificationStackScrollLayoutController =
-        mock {
-            whenever(view).thenReturn(mock())
-            whenever(shelfHeight).thenReturn(0)
-        }
-
     private val testComponent: TestComponent =
         DaggerSharedNotificationContainerViewModelTest_TestComponent.factory()
             .create(
                 test = this,
                 featureFlags =
                     FakeFeatureFlagsClassicModule { set(Flags.FULL_SCREEN_USER_SWITCHER, true) },
-                mocks =
-                    TestMocksModule(
-                        notificationStackSizeCalculator = notificationStackSizeCalculator,
-                        notificationStackScrollLayoutController =
-                            notificationStackScrollLayoutController,
-                    )
+                mocks = TestMocksModule(),
             )
 
     @Test
@@ -336,17 +319,9 @@
     @Test
     fun maxNotificationsOnLockscreen() =
         testComponent.runTest {
-            whenever(
-                    notificationStackSizeCalculator.computeMaxKeyguardNotifications(
-                        any(),
-                        any(),
-                        any(),
-                        any()
-                    )
-                )
-                .thenReturn(10)
-
-            val maxNotifications by collectLastValue(underTest.maxNotifications)
+            var notificationCount = 10
+            val maxNotifications by
+                collectLastValue(underTest.getMaxNotifications { notificationCount })
 
             showLockscreen()
 
@@ -356,21 +331,52 @@
                 SharedNotificationContainerPosition(top = 1f, bottom = 2f)
 
             assertThat(maxNotifications).isEqualTo(10)
+
+            // Also updates when directly requested (as it would from NotificationStackScrollLayout)
+            notificationCount = 25
+            sharedNotificationContainerInteractor.notificationStackChanged()
+            assertThat(maxNotifications).isEqualTo(25)
+        }
+
+    @Test
+    fun maxNotificationsOnLockscreen_DoesNotUpdateWhenUserInteracting() =
+        testComponent.runTest {
+            var notificationCount = 10
+            val maxNotifications by
+                collectLastValue(underTest.getMaxNotifications { notificationCount })
+
+            showLockscreen()
+
+            overrideResource(R.bool.config_use_split_notification_shade, false)
+            configurationRepository.onAnyConfigurationChange()
+            keyguardInteractor.sharedNotificationContainerPosition.value =
+                SharedNotificationContainerPosition(top = 1f, bottom = 2f)
+
+            assertThat(maxNotifications).isEqualTo(10)
+
+            // Shade expanding... still 10
+            shadeRepository.setLockscreenShadeExpansion(0.5f)
+            assertThat(maxNotifications).isEqualTo(10)
+
+            notificationCount = 25
+
+            // When shade is expanding by user interaction
+            shadeRepository.setLegacyLockscreenShadeTracking(true)
+
+            // Should still be 10, since the user is interacting
+            assertThat(maxNotifications).isEqualTo(10)
+
+            shadeRepository.setLegacyLockscreenShadeTracking(false)
+            shadeRepository.setLockscreenShadeExpansion(0f)
+
+            // Stopped tracking, show 25
+            assertThat(maxNotifications).isEqualTo(25)
         }
 
     @Test
     fun maxNotificationsOnShade() =
         testComponent.runTest {
-            whenever(
-                    notificationStackSizeCalculator.computeMaxKeyguardNotifications(
-                        any(),
-                        any(),
-                        any(),
-                        any()
-                    )
-                )
-                .thenReturn(10)
-            val maxNotifications by collectLastValue(underTest.maxNotifications)
+            val maxNotifications by collectLastValue(underTest.getMaxNotifications { 10 })
 
             // Show lockscreen with shade expanded
             showLockscreenWithShadeExpanded()
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/FakeShadeRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/FakeShadeRepository.kt
index 800593f..02318ab 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/FakeShadeRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/FakeShadeRepository.kt
@@ -59,6 +59,8 @@
     private val _legacyIsQsExpanded = MutableStateFlow(false)
     @Deprecated("Use ShadeInteractor instead") override val legacyIsQsExpanded = _legacyIsQsExpanded
 
+    override val legacyLockscreenShadeTracking = MutableStateFlow(false)
+
     @Deprecated("Use ShadeInteractor instead")
     override fun setLegacyIsQsExpanded(legacyIsQsExpanded: Boolean) {
         _legacyIsQsExpanded.value = legacyIsQsExpanded
@@ -81,6 +83,11 @@
         _legacyShadeTracking.value = tracking
     }
 
+    @Deprecated("Should only be called by NPVC and tests")
+    override fun setLegacyLockscreenShadeTracking(tracking: Boolean) {
+        legacyLockscreenShadeTracking.value = tracking
+    }
+
     fun setShadeModel(model: ShadeModel) {
         _shadeModel.value = model
     }