Merge pull request #1063 from android/ben/scrollbars_perf

Move Scrollbar state reads out of composition

core/designsystem/src/main/kotlin/com/google/samples/apps/nowinandroid/core/designsystem/component/scrollbar/AppScrollbars.kt

@@ -16,10 +16,10 @@
 
 package com.google.samples.apps.nowinandroid.core.designsystem.component.scrollbar
 
+import android.annotation.SuppressLint
 import androidx.compose.animation.animateColorAsState
 import androidx.compose.animation.core.Spring
 import androidx.compose.animation.core.SpringSpec
-import androidx.compose.foundation.background
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.foundation.gestures.Orientation.Horizontal
 import androidx.compose.foundation.gestures.Orientation.Vertical
@@ -38,12 +38,22 @@ import androidx.compose.foundation.shape.RoundedCornerShape
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.State
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.geometry.Size
 import androidx.compose.ui.graphics.Color
+import androidx.compose.ui.graphics.ColorProducer
+import androidx.compose.ui.graphics.Outline
+import androidx.compose.ui.graphics.drawOutline
+import androidx.compose.ui.graphics.drawscope.ContentDrawScope
+import androidx.compose.ui.node.DrawModifierNode
+import androidx.compose.ui.node.ModifierNodeElement
+import androidx.compose.ui.node.invalidateDraw
+import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.unit.dp
 import com.google.samples.apps.nowinandroid.core.designsystem.component.scrollbar.ThumbState.Active
 import com.google.samples.apps.nowinandroid.core.designsystem.component.scrollbar.ThumbState.Dormant
@@ -130,12 +140,7 @@ private fun ScrollableState.DraggableScrollbarThumb(
                     Horizontal -> height(12.dp).fillMaxWidth()
                 }
             }
-            .background(
-                color = scrollbarThumbColor(
-                    interactionSource = interactionSource,
-                ),
-                shape = RoundedCornerShape(16.dp),
-            ),
+            .scrollThumb(this, interactionSource),
     )
 }
 
@@ -155,31 +160,72 @@ private fun ScrollableState.DecorativeScrollbarThumb(
                     Horizontal -> height(2.dp).fillMaxWidth()
                 }
             }
-            .background(
-                color = scrollbarThumbColor(
-                    interactionSource = interactionSource,
-                ),
-                shape = RoundedCornerShape(16.dp),
-            ),
+            .scrollThumb(this, interactionSource),
     )
 }
 
+// TODO: This lint is removed in 1.6 as the recommendation has changed
+// remove when project is upgraded
+@SuppressLint("ComposableModifierFactory")
+@Composable
+private fun Modifier.scrollThumb(
+    scrollableState: ScrollableState,
+    interactionSource: InteractionSource,
+): Modifier {
+    val colorState = scrollbarThumbColor(scrollableState, interactionSource)
+    return this then ScrollThumbElement { colorState.value }
+}
+
+private data class ScrollThumbElement(val colorProducer: ColorProducer) :
+    ModifierNodeElement<ScrollThumbNode>() {
+    override fun create(): ScrollThumbNode = ScrollThumbNode(colorProducer)
+    override fun update(node: ScrollThumbNode) {
+        node.colorProducer = colorProducer
+        node.invalidateDraw()
+    }
+}
+
+private class ScrollThumbNode(var colorProducer: ColorProducer) : DrawModifierNode, Modifier.Node() {
+    private val shape = RoundedCornerShape(16.dp)
+
+    // naive cache outline calculation if size is the same
+    private var lastSize: Size? = null
+    private var lastLayoutDirection: LayoutDirection? = null
+    private var lastOutline: Outline? = null
+
+    override fun ContentDrawScope.draw() {
+        val color = colorProducer()
+        val outline =
+            if (size == lastSize && layoutDirection == lastLayoutDirection) {
+                lastOutline!!
+            } else {
+                shape.createOutline(size, layoutDirection, this)
+            }
+        if (color != Color.Unspecified) drawOutline(outline, color = color)
+
+        lastOutline = outline
+        lastSize = size
+        lastLayoutDirection = layoutDirection
+    }
+}
+
 /**
  * The color of the scrollbar thumb as a function of its interaction state.
  * @param interactionSource source of interactions in the scrolling container
  */
 @Composable
-private fun ScrollableState.scrollbarThumbColor(
+private fun scrollbarThumbColor(
+    scrollableState: ScrollableState,
     interactionSource: InteractionSource,
-): Color {
+): State<Color> {
     var state by remember { mutableStateOf(Dormant) }
     val pressed by interactionSource.collectIsPressedAsState()
     val hovered by interactionSource.collectIsHoveredAsState()
     val dragged by interactionSource.collectIsDraggedAsState()
-    val active = (canScrollForward || canScrollForward) &&
-        (pressed || hovered || dragged || isScrollInProgress)
+    val active = (scrollableState.canScrollForward || scrollableState.canScrollForward) &&
+        (pressed || hovered || dragged || scrollableState.isScrollInProgress)
 
-    val color by animateColorAsState(
+    val color = animateColorAsState(
         targetValue = when (state) {
             Active -> MaterialTheme.colorScheme.onSurface.copy(0.5f)
             Inactive -> MaterialTheme.colorScheme.onSurface.copy(alpha = 0.2f)

core/designsystem/src/main/kotlin/com/google/samples/apps/nowinandroid/core/designsystem/component/scrollbar/Scrollbar.kt

@@ -16,8 +16,9 @@
 
 package com.google.samples.apps.nowinandroid.core.designsystem.component.scrollbar
 
-import androidx.compose.animation.core.animateDpAsState
 import androidx.compose.foundation.gestures.Orientation
+import androidx.compose.foundation.gestures.Orientation.Horizontal
+import androidx.compose.foundation.gestures.Orientation.Vertical
 import androidx.compose.foundation.gestures.detectHorizontalDragGestures
 import androidx.compose.foundation.gestures.detectTapGestures
 import androidx.compose.foundation.gestures.detectVerticalDragGestures
@@ -28,31 +29,28 @@ import androidx.compose.foundation.interaction.PressInteraction
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.fillMaxHeight
 import androidx.compose.foundation.layout.fillMaxWidth
-import androidx.compose.foundation.layout.height
-import androidx.compose.foundation.layout.offset
-import androidx.compose.foundation.layout.width
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.Immutable
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableFloatStateOf
+import androidx.compose.runtime.mutableLongStateOf
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
-import androidx.compose.runtime.rememberUpdatedState
 import androidx.compose.runtime.setValue
-import androidx.compose.ui.Alignment
+import androidx.compose.runtime.snapshotFlow
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.input.pointer.PointerInputChange
 import androidx.compose.ui.input.pointer.pointerInput
+import androidx.compose.ui.layout.Layout
+import androidx.compose.ui.layout.layout
 import androidx.compose.ui.layout.onGloballyPositioned
 import androidx.compose.ui.layout.positionInRoot
-import androidx.compose.ui.platform.LocalDensity
 import androidx.compose.ui.unit.Dp
 import androidx.compose.ui.unit.IntOffset
 import androidx.compose.ui.unit.IntSize
 import androidx.compose.ui.unit.dp
-import androidx.compose.ui.unit.max
 import androidx.compose.ui.util.packFloats
 import androidx.compose.ui.util.unpackFloat1
 import androidx.compose.ui.util.unpackFloat2
@@ -61,6 +59,7 @@ import kotlinx.coroutines.delay
 import kotlinx.coroutines.withTimeout
 import kotlin.math.max
 import kotlin.math.min
+import kotlin.math.roundToInt
 
 /**
  * The delay between scrolls when a user long presses on the scrollbar track to initiate a scroll
@@ -74,21 +73,53 @@ private const val SCROLLBAR_PRESS_DELAY_MS = 10L
  */
 private const val SCROLLBAR_PRESS_DELTA_PCT = 0.02f
 
+class ScrollbarState {
+    private var packedValue by mutableLongStateOf(0L)
+
+    internal fun onScroll(stateValue: ScrollbarStateValue) {
+        packedValue = stateValue.packedValue
+    }
+
+    /**
+     * Returns the thumb size of the scrollbar as a percentage of the total track size
+     */
+    val thumbSizePercent
+        get() = unpackFloat1(packedValue)
+
+    /**
+     * Returns the distance the thumb has traveled as a percentage of total track size
+     */
+    val thumbMovedPercent
+        get() = unpackFloat2(packedValue)
+}
+
+/**
+ * Returns the size of the scrollbar track in pixels
+ */
+private val ScrollbarTrack.size
+    get() = unpackFloat2(packedValue) - unpackFloat1(packedValue)
+
+/**
+ * Returns the position of the scrollbar thumb on the track as a percentage
+ */
+private fun ScrollbarTrack.thumbPosition(
+    dimension: Float,
+): Float = max(
+    a = min(
+        a = dimension / size,
+        b = 1f,
+    ),
+    b = 0f,
+)
+
 /**
  * Class definition for the core properties of a scroll bar
  */
 @Immutable
 @JvmInline
-value class ScrollbarState internal constructor(
+value class ScrollbarStateValue internal constructor(
     internal val packedValue: Long,
-) {
-    companion object {
-        val FULL = ScrollbarState(
-            thumbSizePercent = 1f,
-            thumbMovedPercent = 0f,
-        )
-    }
-}
+)
 
 /**
  * Class definition for the core properties of a scroll bar track
@@ -105,54 +136,23 @@ private value class ScrollbarTrack(
 }
 
 /**
- * Creates a [ScrollbarState] with the listed properties
+ * Creates a [ScrollbarStateValue] with the listed properties
  * @param thumbSizePercent the thumb size of the scrollbar as a percentage of the total track size.
  *  Refers to either the thumb width (for horizontal scrollbars)
  *  or height (for vertical scrollbars).
  * @param thumbMovedPercent the distance the thumb has traveled as a percentage of total
  * track size.
  */
-fun ScrollbarState(
+fun scrollbarStateValue(
     thumbSizePercent: Float,
     thumbMovedPercent: Float,
-) = ScrollbarState(
+) = ScrollbarStateValue(
     packFloats(
         val1 = thumbSizePercent,
         val2 = thumbMovedPercent,
     ),
 )
 
-/**
- * Returns the thumb size of the scrollbar as a percentage of the total track size
- */
-val ScrollbarState.thumbSizePercent
-    get() = unpackFloat1(packedValue)
-
-/**
- * Returns the distance the thumb has traveled as a percentage of total track size
- */
-val ScrollbarState.thumbMovedPercent
-    get() = unpackFloat2(packedValue)
-
-/**
- * Returns the size of the scrollbar track in pixels
- */
-private val ScrollbarTrack.size
-    get() = unpackFloat2(packedValue) - unpackFloat1(packedValue)
-
-/**
- * Returns the position of the scrollbar thumb on the track as a percentage
- */
-private fun ScrollbarTrack.thumbPosition(
-    dimension: Float,
-): Float = max(
-    a = min(
-        a = dimension / size,
-        b = 1f,
-    ),
-    b = 0f,
-)
-
 /**
  * Returns the value of [offset] along the axis specified by [this]
  */
@@ -197,8 +197,6 @@ fun Scrollbar(
     thumb: @Composable () -> Unit,
     onThumbMoved: ((Float) -> Unit)? = null,
 ) {
-    val localDensity = LocalDensity.current
-
     // Using Offset.Unspecified and Float.NaN instead of null
     // to prevent unnecessary boxing of primitives
     var pressedOffset by remember { mutableStateOf(Offset.Unspecified) }
@@ -210,23 +208,6 @@ fun Scrollbar(
 
     var track by remember { mutableStateOf(ScrollbarTrack(packedValue = 0)) }
 
-    val thumbTravelPercent = when {
-        interactionThumbTravelPercent.isNaN() -> state.thumbMovedPercent
-        else -> interactionThumbTravelPercent
-    }
-    val thumbSizePx = max(
-        a = state.thumbSizePercent * track.size,
-        b = with(localDensity) { minThumbSize.toPx() },
-    )
-    val thumbSizeDp by animateDpAsState(
-        targetValue = with(localDensity) { thumbSizePx.toDp() },
-        label = "scrollbar thumb size",
-    )
-    val thumbMovedPx = min(
-        a = track.size * thumbTravelPercent,
-        b = track.size - thumbSizePx,
-    )
-
     // scrollbar track container
     Box(
         modifier = modifier
@@ -320,84 +301,110 @@ fun Scrollbar(
                 }
             },
     ) {
-        val scrollbarThumbMovedDp = max(
-            a = with(localDensity) { thumbMovedPx.toDp() },
-            b = 0.dp,
-        )
         // scrollbar thumb container
-        Box(
-            modifier = Modifier
-                .align(Alignment.TopStart)
-                .run {
-                    when (orientation) {
-                        Orientation.Horizontal -> width(thumbSizeDp)
-                        Orientation.Vertical -> height(thumbSizeDp)
-                    }
+        Layout(content = { thumb() }) { measurables, constraints ->
+            val measurable = measurables.first()
+
+            val thumbSizePx = max(
+                a = state.thumbSizePercent * track.size,
+                b = minThumbSize.toPx(),
+            )
+
+            val thumbTravelPercent = when {
+                interactionThumbTravelPercent.isNaN() -> state.thumbMovedPercent
+                else -> interactionThumbTravelPercent
+            }
+
+            val thumbMovedPx = min(
+                a = track.size * thumbTravelPercent,
+                b = track.size - thumbSizePx,
+            )
+
+            val scrollbarThumbMovedPx = max(
+                a = thumbMovedPx.roundToInt(),
+                b = 0,
+            )
+
+            val y = when (orientation) {
+                Horizontal -> 0
+                Vertical -> scrollbarThumbMovedPx
+            }
+            val x = when (orientation) {
+                Horizontal -> scrollbarThumbMovedPx
+                Vertical -> 0
+            }
+
+            val updatedConstraints = when (orientation) {
+                Horizontal -> {
+                    constraints.copy(
+                        minWidth = thumbSizePx.roundToInt(),
+                        maxWidth = thumbSizePx.roundToInt(),
+                    )
                 }
-                .offset(
-                    y = when (orientation) {
-                        Orientation.Horizontal -> 0.dp
-                        Orientation.Vertical -> scrollbarThumbMovedDp
-                    },
-                    x = when (orientation) {
-                        Orientation.Horizontal -> scrollbarThumbMovedDp
-                        Orientation.Vertical -> 0.dp
-                    },
-                ),
-        ) {
-            thumb()
+                Vertical -> {
+                    constraints.copy(
+                        minHeight = thumbSizePx.roundToInt(),
+                        maxHeight = thumbSizePx.roundToInt(),
+                    )
+                }
+            }
+
+            val placeable = measurable.measure(updatedConstraints)
+            layout(placeable.width, placeable.height) {
+                placeable.place(x, y)
+            }
         }
     }
 
     if (onThumbMoved == null) return
 
-    // State that will be read inside the effects that follow
-    // but will not cause re-triggering of them
-    val updatedState by rememberUpdatedState(state)
-
     // Process presses
-    LaunchedEffect(pressedOffset) {
-        // Press ended, reset interactionThumbTravelPercent
-        if (pressedOffset == Offset.Unspecified) {
-            interactionThumbTravelPercent = Float.NaN
-            return@LaunchedEffect
-        }
+    LaunchedEffect(Unit) {
+        snapshotFlow { pressedOffset }.collect { pressedOffset ->
+            // Press ended, reset interactionThumbTravelPercent
+            if (pressedOffset == Offset.Unspecified) {
+                interactionThumbTravelPercent = Float.NaN
+                return@collect
+            }
 
-        var currentThumbMovedPercent = updatedState.thumbMovedPercent
-        val destinationThumbMovedPercent = track.thumbPosition(
-            dimension = orientation.valueOf(pressedOffset),
-        )
-        val isPositive = currentThumbMovedPercent < destinationThumbMovedPercent
-        val delta = SCROLLBAR_PRESS_DELTA_PCT * if (isPositive) 1f else -1f
-
-        while (currentThumbMovedPercent != destinationThumbMovedPercent) {
-            currentThumbMovedPercent = when {
-                isPositive -> min(
-                    a = currentThumbMovedPercent + delta,
-                    b = destinationThumbMovedPercent,
-                )
+            var currentThumbMovedPercent = state.thumbMovedPercent
+            val destinationThumbMovedPercent = track.thumbPosition(
+                dimension = orientation.valueOf(pressedOffset),
+            )
+            val isPositive = currentThumbMovedPercent < destinationThumbMovedPercent
+            val delta = SCROLLBAR_PRESS_DELTA_PCT * if (isPositive) 1f else -1f
+
+            while (currentThumbMovedPercent != destinationThumbMovedPercent) {
+                currentThumbMovedPercent = when {
+                    isPositive -> min(
+                        a = currentThumbMovedPercent + delta,
+                        b = destinationThumbMovedPercent,
+                    )
 
-                else -> max(
-                    a = currentThumbMovedPercent + delta,
-                    b = destinationThumbMovedPercent,
-                )
+                    else -> max(
+                        a = currentThumbMovedPercent + delta,
+                        b = destinationThumbMovedPercent,
+                    )
+                }
+                onThumbMoved(currentThumbMovedPercent)
+                interactionThumbTravelPercent = currentThumbMovedPercent
+                delay(SCROLLBAR_PRESS_DELAY_MS)
             }
-            onThumbMoved(currentThumbMovedPercent)
-            interactionThumbTravelPercent = currentThumbMovedPercent
-            delay(SCROLLBAR_PRESS_DELAY_MS)
         }
     }
 
     // Process drags
-    LaunchedEffect(draggedOffset) {
-        if (draggedOffset == Offset.Unspecified) {
-            interactionThumbTravelPercent = Float.NaN
-            return@LaunchedEffect
+    LaunchedEffect(Unit) {
+        snapshotFlow { draggedOffset }.collect { draggedOffset ->
+            if (draggedOffset == Offset.Unspecified) {
+                interactionThumbTravelPercent = Float.NaN
+                return@collect
+            }
+            val currentTravel = track.thumbPosition(
+                dimension = orientation.valueOf(draggedOffset),
+            )
+            onThumbMoved(currentTravel)
+            interactionThumbTravelPercent = currentTravel
         }
-        val currentTravel = track.thumbPosition(
-            dimension = orientation.valueOf(draggedOffset),
-        )
-        onThumbMoved(currentTravel)
-        interactionThumbTravelPercent = currentTravel
     }
 }

core/designsystem/src/main/kotlin/com/google/samples/apps/nowinandroid/core/designsystem/component/scrollbar/ScrollbarExt.kt

@@ -24,7 +24,8 @@ import androidx.compose.foundation.lazy.grid.LazyGridState
 import androidx.compose.foundation.lazy.staggeredgrid.LazyStaggeredGridItemInfo
 import androidx.compose.foundation.lazy.staggeredgrid.LazyStaggeredGridState
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.produceState
+import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.remember
 import androidx.compose.runtime.snapshotFlow
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.filterNotNull
@@ -40,58 +41,58 @@ import kotlin.math.min
 fun LazyListState.scrollbarState(
     itemsAvailable: Int,
     itemIndex: (LazyListItemInfo) -> Int = LazyListItemInfo::index,
-): ScrollbarState = produceState(
-    initialValue = ScrollbarState.FULL,
-    key1 = this,
-    key2 = itemsAvailable,
-) {
-    snapshotFlow {
-        if (itemsAvailable == 0) return@snapshotFlow null
-
-        val visibleItemsInfo = layoutInfo.visibleItemsInfo
-        if (visibleItemsInfo.isEmpty()) return@snapshotFlow null
-
-        val firstIndex = min(
-            a = interpolateFirstItemIndex(
-                visibleItems = visibleItemsInfo,
-                itemSize = { it.size },
-                offset = { it.offset },
-                nextItemOnMainAxis = { first -> visibleItemsInfo.find { it != first } },
-                itemIndex = itemIndex,
-            ),
-            b = itemsAvailable.toFloat(),
-        )
-        if (firstIndex.isNaN()) return@snapshotFlow null
-
-        val itemsVisible = visibleItemsInfo.floatSumOf { itemInfo ->
-            itemVisibilityPercentage(
-                itemSize = itemInfo.size,
-                itemStartOffset = itemInfo.offset,
-                viewportStartOffset = layoutInfo.viewportStartOffset,
-                viewportEndOffset = layoutInfo.viewportEndOffset,
+): ScrollbarState {
+    val state = remember { ScrollbarState() }
+    LaunchedEffect(this, itemsAvailable) {
+        snapshotFlow {
+            if (itemsAvailable == 0) return@snapshotFlow null
+
+            val visibleItemsInfo = layoutInfo.visibleItemsInfo
+            if (visibleItemsInfo.isEmpty()) return@snapshotFlow null
+
+            val firstIndex = min(
+                a = interpolateFirstItemIndex(
+                    visibleItems = visibleItemsInfo,
+                    itemSize = { it.size },
+                    offset = { it.offset },
+                    nextItemOnMainAxis = { first -> visibleItemsInfo.find { it != first } },
+                    itemIndex = itemIndex,
+                ),
+                b = itemsAvailable.toFloat(),
+            )
+            if (firstIndex.isNaN()) return@snapshotFlow null
+
+            val itemsVisible = visibleItemsInfo.floatSumOf { itemInfo ->
+                itemVisibilityPercentage(
+                    itemSize = itemInfo.size,
+                    itemStartOffset = itemInfo.offset,
+                    viewportStartOffset = layoutInfo.viewportStartOffset,
+                    viewportEndOffset = layoutInfo.viewportEndOffset,
+                )
+            }
+
+            val thumbTravelPercent = min(
+                a = firstIndex / itemsAvailable,
+                b = 1f,
+            )
+            val thumbSizePercent = min(
+                a = itemsVisible / itemsAvailable,
+                b = 1f,
+            )
+            scrollbarStateValue(
+                thumbSizePercent = thumbSizePercent,
+                thumbMovedPercent = when {
+                    layoutInfo.reverseLayout -> 1f - thumbTravelPercent
+                    else -> thumbTravelPercent
+                },
             )
         }
-
-        val thumbTravelPercent = min(
-            a = firstIndex / itemsAvailable,
-            b = 1f,
-        )
-        val thumbSizePercent = min(
-            a = itemsVisible / itemsAvailable,
-            b = 1f,
-        )
-        ScrollbarState(
-            thumbSizePercent = thumbSizePercent,
-            thumbMovedPercent = when {
-                layoutInfo.reverseLayout -> 1f - thumbTravelPercent
-                else -> thumbTravelPercent
-            },
-        )
+            .filterNotNull()
+            .distinctUntilChanged()
+            .collect { state.onScroll(it) }
     }
-        .filterNotNull()
-        .distinctUntilChanged()
-        .collect { value = it }
-}.value
+    return state
+}
 
 /**
  * Calculates a [ScrollbarState] driven by the changes in a [LazyGridState]
@@ -103,68 +104,68 @@ fun LazyListState.scrollbarState(
 fun LazyGridState.scrollbarState(
     itemsAvailable: Int,
     itemIndex: (LazyGridItemInfo) -> Int = LazyGridItemInfo::index,
-): ScrollbarState = produceState(
-    initialValue = ScrollbarState.FULL,
-    key1 = this,
-    key2 = itemsAvailable,
-) {
-    snapshotFlow {
-        if (itemsAvailable == 0) return@snapshotFlow null
-
-        val visibleItemsInfo = layoutInfo.visibleItemsInfo
-        if (visibleItemsInfo.isEmpty()) return@snapshotFlow null
-
-        val firstIndex = min(
-            a = interpolateFirstItemIndex(
-                visibleItems = visibleItemsInfo,
-                itemSize = { layoutInfo.orientation.valueOf(it.size) },
-                offset = { layoutInfo.orientation.valueOf(it.offset) },
-                nextItemOnMainAxis = { first ->
-                    when (layoutInfo.orientation) {
-                        Orientation.Vertical -> visibleItemsInfo.find {
-                            it != first && it.row != first.row
-                        }
-
-                        Orientation.Horizontal -> visibleItemsInfo.find {
-                            it != first && it.column != first.column
+): ScrollbarState {
+    val state = remember { ScrollbarState() }
+    LaunchedEffect(this, itemsAvailable) {
+        snapshotFlow {
+            if (itemsAvailable == 0) return@snapshotFlow null
+
+            val visibleItemsInfo = layoutInfo.visibleItemsInfo
+            if (visibleItemsInfo.isEmpty()) return@snapshotFlow null
+
+            val firstIndex = min(
+                a = interpolateFirstItemIndex(
+                    visibleItems = visibleItemsInfo,
+                    itemSize = { layoutInfo.orientation.valueOf(it.size) },
+                    offset = { layoutInfo.orientation.valueOf(it.offset) },
+                    nextItemOnMainAxis = { first ->
+                        when (layoutInfo.orientation) {
+                            Orientation.Vertical -> visibleItemsInfo.find {
+                                it != first && it.row != first.row
+                            }
+
+                            Orientation.Horizontal -> visibleItemsInfo.find {
+                                it != first && it.column != first.column
+                            }
                         }
-                    }
+                    },
+                    itemIndex = itemIndex,
+                ),
+                b = itemsAvailable.toFloat(),
+            )
+            if (firstIndex.isNaN()) return@snapshotFlow null
+
+            val itemsVisible = visibleItemsInfo.floatSumOf { itemInfo ->
+                itemVisibilityPercentage(
+                    itemSize = layoutInfo.orientation.valueOf(itemInfo.size),
+                    itemStartOffset = layoutInfo.orientation.valueOf(itemInfo.offset),
+                    viewportStartOffset = layoutInfo.viewportStartOffset,
+                    viewportEndOffset = layoutInfo.viewportEndOffset,
+                )
+            }
+
+            val thumbTravelPercent = min(
+                a = firstIndex / itemsAvailable,
+                b = 1f,
+            )
+            val thumbSizePercent = min(
+                a = itemsVisible / itemsAvailable,
+                b = 1f,
+            )
+            scrollbarStateValue(
+                thumbSizePercent = thumbSizePercent,
+                thumbMovedPercent = when {
+                    layoutInfo.reverseLayout -> 1f - thumbTravelPercent
+                    else -> thumbTravelPercent
                 },
-                itemIndex = itemIndex,
-            ),
-            b = itemsAvailable.toFloat(),
-        )
-        if (firstIndex.isNaN()) return@snapshotFlow null
-
-        val itemsVisible = visibleItemsInfo.floatSumOf { itemInfo ->
-            itemVisibilityPercentage(
-                itemSize = layoutInfo.orientation.valueOf(itemInfo.size),
-                itemStartOffset = layoutInfo.orientation.valueOf(itemInfo.offset),
-                viewportStartOffset = layoutInfo.viewportStartOffset,
-                viewportEndOffset = layoutInfo.viewportEndOffset,
             )
         }
-
-        val thumbTravelPercent = min(
-            a = firstIndex / itemsAvailable,
-            b = 1f,
-        )
-        val thumbSizePercent = min(
-            a = itemsVisible / itemsAvailable,
-            b = 1f,
-        )
-        ScrollbarState(
-            thumbSizePercent = thumbSizePercent,
-            thumbMovedPercent = when {
-                layoutInfo.reverseLayout -> 1f - thumbTravelPercent
-                else -> thumbTravelPercent
-            },
-        )
+            .filterNotNull()
+            .distinctUntilChanged()
+            .collect { state.onScroll(it) }
     }
-        .filterNotNull()
-        .distinctUntilChanged()
-        .collect { value = it }
-}.value
+    return state
+}
 
 /**
  * Remembers a [ScrollbarState] driven by the changes in a [LazyStaggeredGridState]
@@ -177,57 +178,57 @@ fun LazyGridState.scrollbarState(
 fun LazyStaggeredGridState.scrollbarState(
     itemsAvailable: Int,
     itemIndex: (LazyStaggeredGridItemInfo) -> Int = LazyStaggeredGridItemInfo::index,
-): ScrollbarState = produceState(
-    initialValue = ScrollbarState.FULL,
-    key1 = this,
-    key2 = itemsAvailable,
-) {
-    snapshotFlow {
-        if (itemsAvailable == 0) return@snapshotFlow null
-
-        val visibleItemsInfo = layoutInfo.visibleItemsInfo
-        if (visibleItemsInfo.isEmpty()) return@snapshotFlow null
-
-        val firstIndex = min(
-            a = interpolateFirstItemIndex(
-                visibleItems = visibleItemsInfo,
-                itemSize = { layoutInfo.orientation.valueOf(it.size) },
-                offset = { layoutInfo.orientation.valueOf(it.offset) },
-                nextItemOnMainAxis = { first ->
-                    visibleItemsInfo.find { it != first && it.lane == first.lane }
-                },
-                itemIndex = itemIndex,
-            ),
-            b = itemsAvailable.toFloat(),
-        )
-        if (firstIndex.isNaN()) return@snapshotFlow null
-
-        val itemsVisible = visibleItemsInfo.floatSumOf { itemInfo ->
-            itemVisibilityPercentage(
-                itemSize = layoutInfo.orientation.valueOf(itemInfo.size),
-                itemStartOffset = layoutInfo.orientation.valueOf(itemInfo.offset),
-                viewportStartOffset = layoutInfo.viewportStartOffset,
-                viewportEndOffset = layoutInfo.viewportEndOffset,
+): ScrollbarState {
+    val state = remember { ScrollbarState() }
+    LaunchedEffect(this, itemsAvailable) {
+        snapshotFlow {
+            if (itemsAvailable == 0) return@snapshotFlow null
+
+            val visibleItemsInfo = layoutInfo.visibleItemsInfo
+            if (visibleItemsInfo.isEmpty()) return@snapshotFlow null
+
+            val firstIndex = min(
+                a = interpolateFirstItemIndex(
+                    visibleItems = visibleItemsInfo,
+                    itemSize = { layoutInfo.orientation.valueOf(it.size) },
+                    offset = { layoutInfo.orientation.valueOf(it.offset) },
+                    nextItemOnMainAxis = { first ->
+                        visibleItemsInfo.find { it != first && it.lane == first.lane }
+                    },
+                    itemIndex = itemIndex,
+                ),
+                b = itemsAvailable.toFloat(),
+            )
+            if (firstIndex.isNaN()) return@snapshotFlow null
+
+            val itemsVisible = visibleItemsInfo.floatSumOf { itemInfo ->
+                itemVisibilityPercentage(
+                    itemSize = layoutInfo.orientation.valueOf(itemInfo.size),
+                    itemStartOffset = layoutInfo.orientation.valueOf(itemInfo.offset),
+                    viewportStartOffset = layoutInfo.viewportStartOffset,
+                    viewportEndOffset = layoutInfo.viewportEndOffset,
+                )
+            }
+
+            val thumbTravelPercent = min(
+                a = firstIndex / itemsAvailable,
+                b = 1f,
+            )
+            val thumbSizePercent = min(
+                a = itemsVisible / itemsAvailable,
+                b = 1f,
+            )
+            scrollbarStateValue(
+                thumbSizePercent = thumbSizePercent,
+                thumbMovedPercent = thumbTravelPercent,
             )
         }
-
-        val thumbTravelPercent = min(
-            a = firstIndex / itemsAvailable,
-            b = 1f,
-        )
-        val thumbSizePercent = min(
-            a = itemsVisible / itemsAvailable,
-            b = 1f,
-        )
-        ScrollbarState(
-            thumbSizePercent = thumbSizePercent,
-            thumbMovedPercent = thumbTravelPercent,
-        )
+            .filterNotNull()
+            .distinctUntilChanged()
+            .collect { state.onScroll(it) }
     }
-        .filterNotNull()
-        .distinctUntilChanged()
-        .collect { value = it }
-}.value
+    return state
+}
 
 private inline fun <T> List<T>.floatSumOf(selector: (T) -> Float): Float {
     var sum = 0f