IndexScrollAnimator.kt

import androidx.compose.animation.core.AnimationSpec
import androidx.compose.animation.core.animate
import androidx.compose.animation.core.tween
import androidx.compose.foundation.gestures.ScrollableState
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.ColumnScope
import androidx.compose.foundation.lazy.LazyListState
import androidx.compose.runtime.Composable
import androidx.compose.runtime.Stable
import androidx.compose.runtime.saveable.Saver
import androidx.compose.runtime.saveable.listSaver
import androidx.compose.runtime.saveable.rememberSaveable
import androidx.compose.ui.Modifier
import androidx.compose.ui.geometry.Offset
import androidx.compose.ui.input.pointer.util.VelocityTracker
import androidx.compose.ui.layout.onSizeChanged

/** Remember an instance of [IndexScrollAnimator].
 * @see IndexScrollAnimator*/
@Composable
fun rememberIndexScrollAnimator(
    firstVisibleItemIndex: Int = 0,
    firstVisibleItemScrollOffset: Int = 0
) = rememberSaveable(saver = IndexScrollAnimator.Saver) {
    IndexScrollAnimator(firstVisibleItemIndex, firstVisibleItemScrollOffset)
}

/** This class is used to animate scroll to a particular index with an [AnimationSpec].
 *
 * There are some assumptions/limitations that this implementation is supposed to work with:
 * 1. Heights for all indices have to remain fixed. Its not that you would have to supply constant heights
 *    but if an item, once rendered, should not be re-rendered with a different height. You can change height of
 *    an item if its inside viewport but before or after animation, NOT while animation is in progress (can improve this
 *    part by some logic where additional height is automatically added to ongoing scroll animation).
 *
 * 2. Trying to animate to a certain index which has not been rendered or its scroll facing neighbours are not rendered even
 *    once will cause this implementation to use `LazyListState.animateScrollToItem` instead. This is because the height of
 *    those items are not currently recorded inorder to calculate the amount to scroll.
 *
 * 3. Activity recreation will cause this implementation to loose all heights causing [animateToIndex] to fall back to
 *    `LazyListState.animateScrollToItem` due to reasons discussed above.
 *
 * Prefer using this implementation to scroll smoothly to nearby indices or indices which were rendered in past.
 *
 * Usage:
 * ```
 * val animator = rememberIndexScrollAnimator()
 *
 * Box(
 *     modifier = Modifier.scrollable(
 *         state = animator.scrollableState,
 *         orientation = Orientation.Vertical,
 *         reverseDirection = true
 *     )
 * ) {
 *     LazyColumn(
 *         state = animator.listState,
 *         userScrollEnabled = false,
 *     ) {
 *         itemsIndexed(items) { index, item ->
 *             // Has to be top-level.
 *             animator.RecordItemSize(index) {
 *                 content()
 *             }
 *         }
 *     }
 * }
 *
 * ```
 */
@Stable
class IndexScrollAnimator(
    firstVisibleItemIndex: Int = 0,
    firstVisibleItemScrollOffset: Int = 0,
) {
    val listState = LazyListState(firstVisibleItemIndex, firstVisibleItemScrollOffset)

    private val velocityTracker = VelocityTracker()
    private var scrolledOffset = firstVisibleItemScrollOffset.toFloat()
    val scrollableState = ScrollableState { delta ->
        updateVelocity(delta)
        listState.dispatchRawDelta(delta)
        delta
    }

    private val heightList = mutableListOf<Int?>()

    private fun updateVelocity(delta: Float) {
        scrolledOffset += delta
        velocityTracker.addPosition(System.currentTimeMillis(), Offset(0f, scrolledOffset))
    }

    private fun updateHeightMap(index: Int, height: Int) {
        if (heightList.lastIndex >= index) {
            heightList[index] = height
        } else {
            repeat(index - heightList.lastIndex) {
                heightList.add(null)
            }
            heightList[index] = height
        }
    }

    @Composable
    fun RecordItemSize(
        index: Int,
        modifier: Modifier = Modifier.Companion,
        content: @Composable ColumnScope.() -> Unit,
    ) {
        Column(
            modifier = modifier.onSizeChanged {
                updateHeightMap(index, it.height)
            },
            content = content,
        )
    }

    /** Animates scroll with given animation spec only if it is rendered or was previously rendered (with heights fixed),
     * else fallbacks to [LazyListState.animateScrollToItem].*/
    suspend fun animateToIndex(
        index: Int,
        animationSpec: AnimationSpec<Float> = tween()
    ) {
        val offsetToScroll = run {
            val firstVisibleItemIndex = listState.firstVisibleItemIndex
            val firstVisibleItemScrollOffset = listState.firstVisibleItemScrollOffset

            if (index > firstVisibleItemIndex) {
                val firstVisibleItemHeight = heightList.getOrNull(firstVisibleItemIndex) ?: return@run null
                val offsetToSecondVisibleItem = firstVisibleItemHeight - firstVisibleItemScrollOffset

                var totalOffset = offsetToSecondVisibleItem
                // Add offset of all items after next visible item and before `index`.
                for (i in (firstVisibleItemIndex + 1)..(index - 1)) {
                    totalOffset += heightList.getOrNull(i) ?: return@run null
                }

                val overscrollAdjustment = run {
                    var scrollingBuffer = 0
                    for (i in index..heightList.lastIndex) {
                        scrollingBuffer += heightList[i]!!
                    }

                    // If negative, overscroll shouldn't occur
                    (listState.layoutInfo.viewportSize.height - scrollingBuffer).coerceAtLeast(0)
                }

                (totalOffset - overscrollAdjustment).coerceAtLeast(0)
            } else {
                val previousVisibleItemIndex = firstVisibleItemIndex - 1
                val offsetToPreviousVisibleItem = firstVisibleItemScrollOffset

                var totalOffset = -offsetToPreviousVisibleItem
                // Add offset of all items after next visible item and before `index`.
                for (i in index..previousVisibleItemIndex) {
                    totalOffset -= heightList.getOrNull(i) ?: return@run null
                }

                totalOffset
            }.toFloat()
        }

        if (offsetToScroll != null) {
            var previousValue = 0f
            scrollableState.scroll {
                animate(
                    0f,
                    offsetToScroll,
                    velocityTracker.calculateVelocity().y,
                    animationSpec = animationSpec
                ) { currentValue, velocity ->
                    val delta = currentValue - previousValue
                    updateVelocity(delta)
                    previousValue += scrollBy(delta)
                }
            }
        } else {
            // Happens in case we do not have heights available for some indices.
            listState.animateScrollToItem(index)
        }
    }

    companion object {
        val Saver: Saver<IndexScrollAnimator, *> = listSaver(
            save = {
                with(LazyListState.Companion.Saver) {
                    save(it.listState) as List<Int>
                }
            },
            restore = {
                IndexScrollAnimator(
                    firstVisibleItemIndex = it[0],
                    firstVisibleItemScrollOffset = it[1]
                )
            }
        )
    }
}