Fix missing space between umo and shelf

StackSizeCalculator needs to account for possible gap
before the shelf, in case the first element in the shelf
starts a new section.

Fixes: 227107113
Test: manual
	=> see gap before umo and shelf
        => no gap between shelf and previous notif if
		previous notif and first notif in shelf
		belong to same section
Test: NotificationStackSizeCalculatorTest
Change-Id: I8724f81327610a49a95b2aff69846bb06119c054
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculator.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculator.kt
index 6c6ed85..d68f371 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculator.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculator.kt
@@ -139,6 +139,8 @@
                 height += spaceNeeded
                 count += 1
             } else {
+                val gapBeforeFirstViewInShelf = current.calculateGapHeight(stack, previous, count)
+                height += gapBeforeFirstViewInShelf
                 height += shelfHeight
                 log { "returning height with shelf -> $height" }
                 return height
@@ -178,7 +180,9 @@
         if (visibleIndex != 0) {
             size += notificationPadding
         }
-        size += calculateGapHeight(stack, previousView, visibleIndex)
+        val gapHeight = calculateGapHeight(stack, previousView, visibleIndex)
+        log { "\ti=$visibleIndex gapHeight=$gapHeight"}
+        size += gapHeight
         return size
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculatorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculatorTest.kt
index 9a4e10c..497a857 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculatorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculatorTest.kt
@@ -22,8 +22,6 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.statusbar.StatusBarState.KEYGUARD
-import com.android.systemui.statusbar.StatusBarState.SHADE
 import com.android.systemui.statusbar.SysuiStatusBarStateController
 import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
@@ -142,11 +140,13 @@
     }
 
     @Test
-    fun computeHeight_returnsLessThanAvailableSpaceUsedToCalculateMaxNotifications() {
+    fun computeHeight_returnsAtMostSpaceAvailable_withGapBeforeShelf() {
         val rowHeight = ROW_HEIGHT
         val shelfHeight = SHELF_HEIGHT
         val totalSpaceForEachRow = GAP_HEIGHT + rowHeight + NOTIFICATION_PADDING
         val availableSpace = totalSpaceForEachRow * 2
+
+        // All rows in separate sections (default setup).
         val rows =
             listOf(
                 createMockRow(rowHeight),
@@ -157,6 +157,28 @@
         assertThat(maxNotifications).isEqualTo(2)
 
         val height = sizeCalculator.computeHeight(stackLayout, maxNotifications, SHELF_HEIGHT)
+        assertThat(height).isAtMost(availableSpace + GAP_HEIGHT + SHELF_HEIGHT)
+    }
+
+    @Test
+    fun computeHeight_returnsAtMostSpaceAvailable_noGapBeforeShelf() {
+        val rowHeight = ROW_HEIGHT
+        val shelfHeight = SHELF_HEIGHT
+        val totalSpaceForEachRow = GAP_HEIGHT + rowHeight + NOTIFICATION_PADDING
+        val availableSpace = totalSpaceForEachRow * 1
+
+        // Both rows are in the same section.
+        whenever(stackLayout.calculateGapHeight(nullable(), nullable(), any()))
+                .thenReturn(0f)
+        val rows =
+                listOf(
+                        createMockRow(rowHeight),
+                        createMockRow(rowHeight))
+
+        val maxNotifications = computeMaxKeyguardNotifications(rows, availableSpace, shelfHeight)
+        assertThat(maxNotifications).isEqualTo(1)
+
+        val height = sizeCalculator.computeHeight(stackLayout, maxNotifications, SHELF_HEIGHT)
         assertThat(height).isAtMost(availableSpace + SHELF_HEIGHT)
     }