Non-lazy horizontal pager implementation in Compose

Pager.kt

/*
 * Copyright 2020 The Android Open Source Project
 * Copyright 2023 Albert Chang
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import android.view.ViewConfiguration
import androidx.compose.animation.core.animate
import androidx.compose.foundation.ExperimentalFoundationApi
import androidx.compose.foundation.MutatePriority
import androidx.compose.foundation.clipScrollableContainer
import androidx.compose.foundation.gestures.FlingBehavior
import androidx.compose.foundation.gestures.Orientation
import androidx.compose.foundation.gestures.ScrollScope
import androidx.compose.foundation.gestures.ScrollableDefaults
import androidx.compose.foundation.gestures.ScrollableState
import androidx.compose.foundation.gestures.animateScrollBy
import androidx.compose.foundation.gestures.scrollable
import androidx.compose.foundation.interaction.MutableInteractionSource
import androidx.compose.foundation.layout.Box
import androidx.compose.foundation.layout.BoxScope
import androidx.compose.foundation.overscroll
import androidx.compose.runtime.Composable
import androidx.compose.runtime.Immutable
import androidx.compose.runtime.Stable
import androidx.compose.runtime.getValue
import androidx.compose.runtime.key
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.rememberCoroutineScope
import androidx.compose.runtime.saveable.SaveableStateHolder
import androidx.compose.runtime.saveable.Saver
import androidx.compose.runtime.saveable.rememberSaveable
import androidx.compose.runtime.saveable.rememberSaveableStateHolder
import androidx.compose.runtime.setValue
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.layout.Layout
import androidx.compose.ui.layout.Measurable
import androidx.compose.ui.layout.ParentDataModifier
import androidx.compose.ui.platform.LocalContext
import androidx.compose.ui.platform.LocalLayoutDirection
import androidx.compose.ui.semantics.ScrollAxisRange
import androidx.compose.ui.semantics.horizontalScrollAxisRange
import androidx.compose.ui.semantics.pageLeft
import androidx.compose.ui.semantics.pageRight
import androidx.compose.ui.semantics.scrollBy
import androidx.compose.ui.semantics.scrollToIndex
import androidx.compose.ui.semantics.semantics
import androidx.compose.ui.unit.Density
import androidx.compose.ui.unit.LayoutDirection
import androidx.compose.ui.unit.constrainHeight
import androidx.compose.ui.util.fastForEach
import androidx.compose.ui.util.fastForEachIndexed
import androidx.compose.ui.util.fastMap
import kotlin.math.ceil
import kotlin.math.floor
import kotlin.math.round
import kotlin.math.roundToInt
import kotlinx.coroutines.launch

// Originated from
// https://github.com/android/compose-samples/blob/main/Jetcaster/app/src/main/java/com/example/jetcaster/util/Pager.kt

@Composable
fun rememberPagerState(
    pageCount: Int,
    initialIndex: Int = 0
): PagerState {
    val velocity = ViewConfiguration.get(LocalContext.current).scaledMinimumFlingVelocity
    val stateHolder = rememberSaveableStateHolder()

    val saver = remember(stateHolder) {
        PagerState.Saver(stateHolder)
    }

    val state = rememberSaveable(stateHolder, saver = saver) {
        PagerState(initialIndex, stateHolder)
    }
    state.minimumFlingVelocity = velocity
    state.pageCount = pageCount
    return state
}

@Stable
class PagerState(
    initialIndex: Int = 0,
    internal val stateHolder: SaveableStateHolder
) : ScrollableState, FlingBehavior {

    internal var minimumFlingVelocity = 0

    private var _pageCount: Int by mutableStateOf(0)
    var pageCount: Int
        get() = _pageCount
        set(value) {
            _pageCount = value
            if (pageIndexNonObservable >= pageCount) {
                pageIndexNonObservable = pageCount - 1
            }
            if (pageIndexNonObservable == pageCount - 1 && pageOffsetNonObservable < 0) {
                pageOffsetNonObservable = 0f
            }
        }

    private var pageIndexNonObservable: Int = initialIndex
        set(value) {
            field = value
            pageIndex = value
        }

    var pageIndex: Int by mutableStateOf(pageIndexNonObservable)
        private set

    private var pageOffsetNonObservable: Float = 0f
        set(value) {
            field = value
            pageOffset = value
        }

    var pageOffset: Float by mutableStateOf(0f)
        private set

    var pageWidth: Int by mutableStateOf(0)
        internal set

    val pageOffsetFraction: Float
        get() = if (pageWidth == 0) 0f else pageOffset / pageWidth

    private val state = ScrollableState(::onScroll)

    private val globalOffset: Float
        get() = pageIndexNonObservable * pageWidth - pageOffsetNonObservable

    private fun onScroll(distance: Float): Float {
        if (pageCount == 0 ||
            distance > 0 && pageIndexNonObservable == pageCount - 1 && pageOffsetNonObservable <= 0 ||
            distance < 0 && pageIndexNonObservable == 0 && pageOffsetNonObservable >= 0
        ) {
            return 0f
        }

        val oldGlobalOffset = globalOffset
        val distanceToConsume = distance
            .coerceIn(-oldGlobalOffset, (pageCount - 1) * pageWidth - oldGlobalOffset)
        val newGlobalOffset = oldGlobalOffset + distanceToConsume
        pageIndexNonObservable = (newGlobalOffset / pageWidth).roundToInt()
        pageOffsetNonObservable = pageIndexNonObservable * pageWidth - newGlobalOffset

        return distanceToConsume
    }

    override val isScrollInProgress: Boolean
        get() = state.isScrollInProgress

    override val canScrollForward: Boolean
        get() = pageCount > 1 &&
                (pageIndexNonObservable < pageCount - 1 || pageOffsetNonObservable > 0)

    override val canScrollBackward: Boolean
        get() = pageCount > 1 && (pageIndexNonObservable > 0 || pageOffsetNonObservable < 0)

    override fun dispatchRawDelta(delta: Float): Float = state.dispatchRawDelta(delta)

    override suspend fun scroll(
        scrollPriority: MutatePriority,
        block: suspend ScrollScope.() -> Unit
    ) = state.scroll(scrollPriority, block)

    override suspend fun ScrollScope.performFling(initialVelocity: Float): Float {
        val scrollOffset = calculateNearestPageOffset(initialVelocity)
        if (scrollOffset == 0f) return initialVelocity
        var previousValue = 0f
        animate(
            initialValue = 0f,
            targetValue = scrollOffset,
            initialVelocity = initialVelocity
        ) { value, _ ->
            val delta = value - previousValue
            scrollBy(delta)
            previousValue = value
        }
        // Snap to account for rounding error
        pageOffsetNonObservable = 0f
        return 0f
    }

    private fun calculateNearestPageOffset(velocity: Float): Float {
        val currentOffset = globalOffset
        val times = currentOffset / pageWidth
        val nearestIndex = when {
            velocity < -minimumFlingVelocity -> floor(times)
            velocity > minimumFlingVelocity -> ceil(times)
            else -> round(times)
        }
        val nearestOffset = nearestIndex * pageWidth
        return nearestOffset - currentOffset
    }

    fun snapTo(pageIndex: Int) {
        if (pageIndex == pageIndexNonObservable || pageIndex !in 0 until pageCount) return
        pageIndexNonObservable = pageIndex
        pageOffsetNonObservable = 0f
    }

    suspend fun animateSnapTo(pageIndex: Int) {
        if (pageIndex == pageIndexNonObservable || pageIndex !in 0 until pageCount) return
        val scrollOffset = pageIndex * pageWidth - globalOffset
        animateScrollBy(scrollOffset)
    }

    companion object {
        fun Saver(stateHolder: SaveableStateHolder): Saver<PagerState, *> = Saver(
            save = { it.pageIndexNonObservable },
            restore = {
                PagerState(
                    initialIndex = it,
                    stateHolder = stateHolder
                )
            }
        )
    }

}

@Immutable
private data class PageData(val pageIndex: Int) : ParentDataModifier {
    override fun Density.modifyParentData(parentData: Any?): Any = this@PageData
}

private val Measurable.pageIndex: Int
    get() = (parentData as? PageData)?.pageIndex ?: error("No PageData for measurable $this")

@OptIn(ExperimentalFoundationApi::class)
@Composable
fun Pager(
    state: PagerState,
    modifier: Modifier = Modifier,
    offscreenLimit: Int = 1,
    reverseLayout: Boolean = false,
    pageContent: @Composable BoxScope.(Int) -> Unit
) {
    val isRtl = LocalLayoutDirection.current == LayoutDirection.Rtl
    // reverse scroll by default, to have "natural" gesture that goes reversed to layout
    // if rtl and horizontal, do not reverse to make it right-to-left
    val reverseDirection = if (isRtl) reverseLayout else !reverseLayout
    val contentFactory = wrapWithStateRestoration(state.stateHolder) { index ->
        { pageContent(index) }
    }
    val scope = rememberCoroutineScope()
    val overscrollEffect = ScrollableDefaults.overscrollEffect()

    Layout(
        content = {
            val minExtra = if (offscreenLimit == 0 && state.pageOffset > 0) -1 else 0
            val maxExtra = if (offscreenLimit == 0 && state.pageOffset < 0) 1 else 0
            val minPage = (state.pageIndex - offscreenLimit + minExtra)
                .coerceAtLeast(0)
            val maxPage = (state.pageIndex + offscreenLimit + maxExtra)
                .coerceAtMost(state.pageCount - 1)

            for (page in minPage..maxPage) {
                val pageData = PageData(page)
                key(pageData) {
                    Box(contentAlignment = Alignment.Center, modifier = pageData) {
                        contentFactory(page).invoke()
                    }
                }
            }
        },
        modifier = modifier
            .semantics {
                horizontalScrollAxisRange = ScrollAxisRange(
                    value = { state.pageIndex + state.pageOffsetFraction },
                    maxValue = { state.pageCount.toFloat() }
                )
                scrollBy { x, _ ->
                    scope.launch {
                        state.animateScrollBy(x)
                    }
                    true
                }
                scrollToIndex { index ->
                    scope.launch {
                        state.animateSnapTo(index)
                    }
                    true
                }
                pageLeft {
                    val targetIndex = state.pageIndex - 1
                    if (targetIndex >= 0) {
                        scope.launch {
                            state.animateSnapTo(targetIndex)
                        }
                        true
                    } else {
                        false
                    }
                }
                pageRight {
                    val targetIndex = state.pageIndex + 1
                    if (targetIndex < state.pageCount) {
                        scope.launch {
                            state.animateSnapTo(targetIndex)
                        }
                        true
                    } else {
                        false
                    }
                }
            }
            .clipScrollableContainer(Orientation.Horizontal)
            .overscroll(overscrollEffect)
            .scrollable(
                state = state,
                orientation = Orientation.Horizontal,
                overscrollEffect = overscrollEffect,
                reverseDirection = reverseDirection,
                flingBehavior = state,
                interactionSource = remember { MutableInteractionSource() }
            )
    ) { measurables, constraints ->
        val relaxedConstraints = constraints.copy(minHeight = 0)
        val placeables = measurables.fastMap { it.measure(relaxedConstraints) }
        val indices = measurables.fastMap { it.pageIndex }

        var maxWidth = 0
        var maxHeight = 0
        placeables.fastForEach {
            if (it.width > maxWidth) maxWidth = it.width
            if (it.height > maxHeight) maxHeight = it.height
        }

        state.pageWidth = maxWidth

        layout(maxWidth, constraints.constrainHeight(maxHeight)) {
            placeables.fastForEachIndexed { i, placeable ->
                val x = state.pageOffset + (indices[i] - state.pageIndex) * maxWidth
                placeable.placeRelative(x.roundToInt(), 0)
            }
        }
    }
}

/**
 * Converts page content factory to another one which adds auto state restoration functionality.
 */
@Composable
private fun wrapWithStateRestoration(
    stateHolder: SaveableStateHolder,
    pageContent: BoxScope.(Int) -> @Composable () -> Unit
): BoxScope.(Int) -> @Composable () -> Unit {
    return remember(stateHolder, pageContent) {
        { index ->
            val content = pageContent(index)
            // we just wrap our original lambda with the one which auto restores the state
            // currently we use index in the list as a key for the restoration, but in the future
            // we will use the user provided key
            (@Composable { stateHolder.SaveableStateProvider(index, content) })
        }
    }
}