Merge "Rework VolumePanelRadioButtons using Layout" into main
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/selector/ui/composable/VolumePanelRadioButtons.kt b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/selector/ui/composable/VolumePanelRadioButtons.kt
index ae267e2..98d1afd 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/selector/ui/composable/VolumePanelRadioButtons.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/selector/ui/composable/VolumePanelRadioButtons.kt
@@ -16,38 +16,38 @@
 
 package com.android.systemui.volume.panel.component.selector.ui.composable
 
-import androidx.compose.animation.core.animateOffsetAsState
-import androidx.compose.foundation.Canvas
+import androidx.compose.animation.core.Animatable
+import androidx.compose.animation.core.VectorConverter
 import androidx.compose.foundation.background
 import androidx.compose.foundation.layout.Arrangement
-import androidx.compose.foundation.layout.Box
-import androidx.compose.foundation.layout.Column
-import androidx.compose.foundation.layout.IntrinsicSize
 import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.RowScope
-import androidx.compose.foundation.layout.fillMaxSize
-import androidx.compose.foundation.layout.height
+import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.layout.offset
 import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.shape.CornerSize
 import androidx.compose.foundation.shape.RoundedCornerShape
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.TextButton
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.mutableIntStateOf
-import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
-import androidx.compose.runtime.setValue
+import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.geometry.CornerRadius
-import androidx.compose.ui.geometry.Offset
-import androidx.compose.ui.geometry.Size
 import androidx.compose.ui.graphics.Color
-import androidx.compose.ui.layout.onGloballyPositioned
+import androidx.compose.ui.layout.Layout
+import androidx.compose.ui.layout.Measurable
+import androidx.compose.ui.layout.MeasurePolicy
+import androidx.compose.ui.layout.MeasureResult
+import androidx.compose.ui.layout.MeasureScope
+import androidx.compose.ui.layout.Placeable
+import androidx.compose.ui.layout.layoutId
 import androidx.compose.ui.platform.LocalDensity
+import androidx.compose.ui.unit.Constraints
 import androidx.compose.ui.unit.Dp
-import androidx.compose.ui.unit.IntSize
+import androidx.compose.ui.unit.IntOffset
 import androidx.compose.ui.unit.dp
+import androidx.compose.ui.util.fastFirst
+import kotlinx.coroutines.launch
 
 /**
  * Radio button group for the Volume Panel. It allows selecting a single item
@@ -65,8 +65,8 @@
     spacing: Dp = VolumePanelRadioButtonBarDefaults.DefaultSpacing,
     labelIndicatorBackgroundSpacing: Dp =
         VolumePanelRadioButtonBarDefaults.DefaultLabelIndicatorBackgroundSpacing,
-    indicatorCornerRadius: CornerRadius =
-        VolumePanelRadioButtonBarDefaults.defaultIndicatorCornerRadius(),
+    indicatorCornerSize: CornerSize =
+        CornerSize(VolumePanelRadioButtonBarDefaults.DefaultIndicatorCornerRadius),
     indicatorBackgroundCornerSize: CornerSize =
         CornerSize(VolumePanelRadioButtonBarDefaults.DefaultIndicatorBackgroundCornerRadius),
     colors: VolumePanelRadioButtonBarColors = VolumePanelRadioButtonBarDefaults.defaultColors(),
@@ -76,60 +76,41 @@
         VolumePanelRadioButtonBarScopeImpl().apply(content).apply {
             require(hasSelectedItem) { "At least one item should be selected" }
         }
-
     val items = scope.items
 
-    var selectedIndex by remember { mutableIntStateOf(items.indexOfFirst { it.isSelected }) }
-
-    var size by remember { mutableStateOf(IntSize(0, 0)) }
-    val spacingPx = with(LocalDensity.current) { spacing.toPx() }
-    val indicatorWidth = size.width / items.size - (spacingPx * (items.size - 1) / items.size)
-    val offset by
-        animateOffsetAsState(
-            targetValue =
-                Offset(
-                    selectedIndex * indicatorWidth + (spacingPx * selectedIndex),
-                    0f,
-                ),
-            label = "VolumePanelRadioButtonOffsetAnimation",
-            finishedListener = {
-                for (itemIndex in items.indices) {
-                    val item = items[itemIndex]
-                    if (itemIndex == selectedIndex) {
-                        item.onItemSelected()
-                        break
-                    }
-                }
-            }
-        )
-
-    Column(modifier = modifier) {
-        Box(modifier = Modifier.height(IntrinsicSize.Max)) {
-            Canvas(
+    val coroutineScope = rememberCoroutineScope()
+    val offsetAnimatable = remember { Animatable(UNSET_OFFSET, Int.VectorConverter) }
+    Layout(
+        modifier = modifier,
+        content = {
+            Spacer(
                 modifier =
-                    Modifier.fillMaxSize()
+                    Modifier.layoutId(RadioButtonBarComponent.ButtonsBackground)
                         .background(
                             colors.indicatorBackgroundColor,
                             RoundedCornerShape(indicatorBackgroundCornerSize),
                         )
+            )
+            Spacer(
+                modifier =
+                    Modifier.layoutId(RadioButtonBarComponent.Indicator)
+                        .offset { IntOffset(offsetAnimatable.value, 0) }
                         .padding(indicatorBackgroundPadding)
-                        .onGloballyPositioned { size = it.size }
-            ) {
-                drawRoundRect(
-                    color = colors.indicatorColor,
-                    topLeft = offset,
-                    size = Size(indicatorWidth, size.height.toFloat()),
-                    cornerRadius = indicatorCornerRadius,
-                )
-            }
+                        .background(
+                            colors.indicatorColor,
+                            RoundedCornerShape(indicatorCornerSize),
+                        )
+            )
             Row(
-                modifier = Modifier.padding(indicatorBackgroundPadding),
+                modifier =
+                    Modifier.layoutId(RadioButtonBarComponent.Buttons)
+                        .padding(indicatorBackgroundPadding),
                 horizontalArrangement = Arrangement.spacedBy(spacing)
             ) {
                 for (itemIndex in items.indices) {
                     TextButton(
                         modifier = Modifier.weight(1f),
-                        onClick = { selectedIndex = itemIndex },
+                        onClick = { items[itemIndex].onItemSelected() },
                     ) {
                         val item = items[itemIndex]
                         if (item.icon !== Empty) {
@@ -138,28 +119,116 @@
                     }
                 }
             }
-        }
-
-        Row(
-            modifier =
-                Modifier.padding(
-                    start = indicatorBackgroundPadding,
-                    top = labelIndicatorBackgroundSpacing,
-                    end = indicatorBackgroundPadding
-                ),
-            horizontalArrangement = Arrangement.spacedBy(spacing),
-        ) {
-            for (itemIndex in items.indices) {
-                TextButton(
-                    modifier = Modifier.weight(1f),
-                    onClick = { selectedIndex = itemIndex },
-                ) {
-                    val item = items[itemIndex]
-                    if (item.icon !== Empty) {
-                        with(items[itemIndex]) { label() }
+            Row(
+                modifier =
+                    Modifier.layoutId(RadioButtonBarComponent.Labels)
+                        .padding(
+                            start = indicatorBackgroundPadding,
+                            top = labelIndicatorBackgroundSpacing,
+                            end = indicatorBackgroundPadding
+                        ),
+                horizontalArrangement = Arrangement.spacedBy(spacing),
+            ) {
+                for (itemIndex in items.indices) {
+                    TextButton(
+                        modifier = Modifier.weight(1f),
+                        onClick = { items[itemIndex].onItemSelected() },
+                    ) {
+                        val item = items[itemIndex]
+                        if (item.icon !== Empty) {
+                            with(items[itemIndex]) { label() }
+                        }
                     }
                 }
             }
+        },
+        measurePolicy =
+            with(LocalDensity.current) {
+                val spacingPx =
+                    (spacing - indicatorBackgroundPadding * 2).roundToPx().coerceAtLeast(0)
+
+                BarMeasurePolicy(
+                    buttonsCount = items.size,
+                    selectedIndex = scope.selectedIndex,
+                    spacingPx = spacingPx,
+                ) {
+                    coroutineScope.launch {
+                        if (offsetAnimatable.value == UNSET_OFFSET) {
+                            offsetAnimatable.snapTo(it)
+                        } else {
+                            offsetAnimatable.animateTo(it)
+                        }
+                    }
+                }
+            },
+    )
+}
+
+private class BarMeasurePolicy(
+    private val buttonsCount: Int,
+    private val selectedIndex: Int,
+    private val spacingPx: Int,
+    private val onTargetIndicatorOffsetMeasured: (Int) -> Unit,
+) : MeasurePolicy {
+
+    override fun MeasureScope.measure(
+        measurables: List<Measurable>,
+        constraints: Constraints
+    ): MeasureResult {
+        val fillWidthConstraints = constraints.copy(minWidth = constraints.maxWidth)
+        val buttonsPlaceable: Placeable =
+            measurables
+                .fastFirst { it.layoutId == RadioButtonBarComponent.Buttons }
+                .measure(fillWidthConstraints)
+        val labelsPlaceable: Placeable =
+            measurables
+                .fastFirst { it.layoutId == RadioButtonBarComponent.Labels }
+                .measure(fillWidthConstraints)
+
+        val buttonsBackgroundPlaceable: Placeable =
+            measurables
+                .fastFirst { it.layoutId == RadioButtonBarComponent.ButtonsBackground }
+                .measure(
+                    Constraints(
+                        minWidth = buttonsPlaceable.width,
+                        maxWidth = buttonsPlaceable.width,
+                        minHeight = buttonsPlaceable.height,
+                        maxHeight = buttonsPlaceable.height,
+                    )
+                )
+
+        val totalSpacing = spacingPx * (buttonsCount - 1)
+        val indicatorWidth = (buttonsBackgroundPlaceable.width - totalSpacing) / buttonsCount
+        val indicatorPlaceable: Placeable =
+            measurables
+                .fastFirst { it.layoutId == RadioButtonBarComponent.Indicator }
+                .measure(
+                    Constraints(
+                        minWidth = indicatorWidth,
+                        maxWidth = indicatorWidth,
+                        minHeight = buttonsBackgroundPlaceable.height,
+                        maxHeight = buttonsBackgroundPlaceable.height,
+                    )
+                )
+
+        onTargetIndicatorOffsetMeasured(
+            selectedIndex * indicatorWidth + (spacingPx * selectedIndex)
+        )
+
+        return layout(constraints.maxWidth, buttonsPlaceable.height + labelsPlaceable.height) {
+            buttonsBackgroundPlaceable.placeRelative(
+                0,
+                0,
+                RadioButtonBarComponent.ButtonsBackground.zIndex,
+            )
+            indicatorPlaceable.placeRelative(0, 0, RadioButtonBarComponent.Indicator.zIndex)
+
+            buttonsPlaceable.placeRelative(0, 0, RadioButtonBarComponent.Buttons.zIndex)
+            labelsPlaceable.placeRelative(
+                0,
+                buttonsBackgroundPlaceable.height,
+                RadioButtonBarComponent.Labels.zIndex,
+            )
         }
     }
 }
@@ -179,12 +248,6 @@
     val DefaultIndicatorCornerRadius = 20.dp
     val DefaultIndicatorBackgroundCornerRadius = 20.dp
 
-    @Composable
-    fun defaultIndicatorCornerRadius(
-        x: Dp = DefaultIndicatorCornerRadius,
-        y: Dp = DefaultIndicatorCornerRadius,
-    ): CornerRadius = with(LocalDensity.current) { CornerRadius(x.toPx(), y.toPx()) }
-
     /**
      * Returns the default VolumePanelRadioButtonBar colors.
      *
@@ -225,9 +288,12 @@
 
 private class VolumePanelRadioButtonBarScopeImpl : VolumePanelRadioButtonBarScope {
 
-    var hasSelectedItem: Boolean = false
+    var selectedIndex: Int = UNSET_INDEX
         private set
 
+    val hasSelectedItem: Boolean
+        get() = selectedIndex != UNSET_INDEX
+
     private val mutableItems: MutableList<Item> = mutableListOf()
     val items: List<Item> = mutableItems
 
@@ -238,21 +304,34 @@
         label: @Composable RowScope.() -> Unit,
     ) {
         require(!isSelected || !hasSelectedItem) { "Only one item should be selected at a time" }
-        hasSelectedItem = hasSelectedItem || isSelected
+        if (isSelected) {
+            selectedIndex = mutableItems.size
+        }
         mutableItems.add(
             Item(
-                isSelected = isSelected,
                 onItemSelected = onItemSelected,
                 icon = icon,
                 label = label,
             )
         )
     }
+
+    private companion object {
+        const val UNSET_INDEX = -1
+    }
 }
 
 private class Item(
-    val isSelected: Boolean,
     val onItemSelected: () -> Unit,
     val icon: @Composable RowScope.() -> Unit,
     val label: @Composable RowScope.() -> Unit,
 )
+
+private const val UNSET_OFFSET = -1
+
+private enum class RadioButtonBarComponent(val zIndex: Float) {
+    ButtonsBackground(0f),
+    Indicator(1f),
+    Buttons(2f),
+    Labels(2f),
+}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/spatialaudio/ui/composable/SpatialAudioPopup.kt b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/spatialaudio/ui/composable/SpatialAudioPopup.kt
index bed0ae8..71b3e8a 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/spatialaudio/ui/composable/SpatialAudioPopup.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/spatialaudio/ui/composable/SpatialAudioPopup.kt
@@ -65,7 +65,7 @@
             return
         }
 
-        val enabledModelStates by viewModel.spatialAudioButtonByEnabled.collectAsState()
+        val enabledModelStates by viewModel.spatialAudioButtons.collectAsState()
         if (enabledModelStates.isEmpty()) {
             return
         }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioViewModel.kt
index 30715d1..f260d61 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioViewModel.kt
@@ -56,7 +56,7 @@
     val isAvailable: StateFlow<Boolean> =
         availabilityCriteria.isAvailable().stateIn(scope, SharingStarted.Eagerly, true)
 
-    val spatialAudioButtonByEnabled: StateFlow<List<SpatialAudioButtonViewModel>> =
+    val spatialAudioButtons: StateFlow<List<SpatialAudioButtonViewModel>> =
         combine(interactor.isEnabled, interactor.isAvailable) { currentIsEnabled, isAvailable ->
                 SpatialAudioEnabledModel.values
                     .filter {