commonMain.androidx.compose.foundation.lazy.layout.LazyLayoutPrefetchState.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of foundation Show documentation
Higher level abstractions of the Compose UI primitives. This library is design system agnostic, providing the high-level building blocks for both application and design-system developers
The newest version!
/*
 * Copyright 2021 The Android Open Source Project
 *
 * 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
 *
 *      http://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.
 */

package androidx.compose.foundation.lazy.layout

import androidx.collection.mutableObjectLongMapOf
import androidx.compose.foundation.ExperimentalFoundationApi
import androidx.compose.foundation.internal.checkPrecondition
import androidx.compose.foundation.internal.requirePrecondition
import androidx.compose.foundation.internal.requirePreconditionNotNull
import androidx.compose.foundation.lazy.layout.LazyLayoutPrefetchState.PrefetchHandle
import androidx.compose.runtime.Stable
import androidx.compose.ui.Modifier
import androidx.compose.ui.layout.SubcomposeLayoutState
import androidx.compose.ui.node.ModifierNodeElement
import androidx.compose.ui.node.TraversableNode
import androidx.compose.ui.node.TraversableNode.Companion.TraverseDescendantsAction
import androidx.compose.ui.platform.InspectorInfo
import androidx.compose.ui.unit.Constraints
import androidx.compose.ui.util.trace
import kotlin.time.measureTime

/**
 * State for lazy items prefetching, used by lazy layouts to instruct the prefetcher.
 *
 * Note: this class is a part of [LazyLayout] harness that allows for building custom lazy layouts.
 * LazyLayout and all corresponding APIs are still under development and are subject to change.
 *
 * @param prefetchScheduler the [PrefetchScheduler] implementation to use to execute prefetch
 *   requests. If null is provided, the default [PrefetchScheduler] for the platform will be used.
 * @param onNestedPrefetch a callback which will be invoked when this LazyLayout is prefetched in
 *   context of a parent LazyLayout, giving a chance to recursively prefetch its own children. See
 *   [NestedPrefetchScope].
 */
@ExperimentalFoundationApi
@Stable
class LazyLayoutPrefetchState(
    internal val prefetchScheduler: PrefetchScheduler? = null,
    private val onNestedPrefetch: (NestedPrefetchScope.() -> Unit)? = null
) {

    private val prefetchMetrics: PrefetchMetrics = PrefetchMetrics()
    internal var prefetchHandleProvider: PrefetchHandleProvider? = null

    /**
     * Schedules precomposition for the new item. If you also want to premeasure the item please use
     * a second overload accepting a [Constraints] param.
     *
     * @param index item index to prefetch.
     */
    fun schedulePrefetch(index: Int): PrefetchHandle = schedulePrefetch(index, ZeroConstraints)

    /**
     * Schedules precomposition and premeasure for the new item.
     *
     * @param index item index to prefetch.
     * @param constraints [Constraints] to use for premeasuring.
     */
    fun schedulePrefetch(index: Int, constraints: Constraints): PrefetchHandle {
        return prefetchHandleProvider?.schedulePrefetch(index, constraints, prefetchMetrics)
            ?: DummyHandle
    }

    internal fun collectNestedPrefetchRequests(): List {
        val onNestedPrefetch = onNestedPrefetch ?: return emptyList()

        return NestedPrefetchScopeImpl().run {
            onNestedPrefetch()
            requests
        }
    }

    sealed interface PrefetchHandle {
        /**
         * Notifies the prefetcher that previously scheduled item is no longer needed. If the item
         * was precomposed already it will be disposed.
         */
        fun cancel()

        /**
         * Marks this prefetch request as urgent, which is a way to communicate that the requested
         * item is expected to be needed during the next frame.
         *
         * For urgent requests we can proceed with doing the prefetch even if the available time in
         * the frame is less than we spend on similar prefetch requests on average.
         */
        fun markAsUrgent()
    }

    private inner class NestedPrefetchScopeImpl : NestedPrefetchScope {

        val requests: List
            get() = _requests

        private val _requests: MutableList = mutableListOf()

        override fun schedulePrefetch(index: Int) {
            schedulePrefetch(index, ZeroConstraints)
        }

        override fun schedulePrefetch(index: Int, constraints: Constraints) {
            val prefetchHandleProvider = prefetchHandleProvider ?: return
            _requests.add(
                prefetchHandleProvider.createNestedPrefetchRequest(
                    index,
                    constraints,
                    prefetchMetrics
                )
            )
        }
    }
}

/**
 * A scope which allows nested prefetches to be requested for the precomposition of a LazyLayout.
 */
@ExperimentalFoundationApi
sealed interface NestedPrefetchScope {

    /**
     * Requests a child index to be prefetched as part of the prefetch of a parent LazyLayout.
     *
     * The prefetch will only do the precomposition for the new item. If you also want to premeasure
     * please use a second overload accepting a [Constraints] param.
     *
     * @param index item index to prefetch.
     */
    fun schedulePrefetch(index: Int)

    /**
     * Requests a child index to be prefetched as part of the prefetch of a parent LazyLayout.
     *
     * @param index the index of the child to prefetch.
     * @param constraints [Constraints] to use for premeasuring. If null, the child will not be
     *   premeasured.
     */
    fun schedulePrefetch(index: Int, constraints: Constraints)
}

/**
 * [PrefetchMetrics] tracks composition and measure timings for subcompositions so that they can be
 * used to estimate whether we can fit prefetch work into idle time without delaying the start of
 * the next frame.
 */
@ExperimentalFoundationApi
internal class PrefetchMetrics {

    val averageCompositionTimeNanosByContentType = mutableObjectLongMapOf()
    val averageMeasureTimeNanosByContentType = mutableObjectLongMapOf()

    /** The current average time composition has taken during prefetches of this LazyLayout. */
    var averageCompositionTimeNanos: Long = 0L
        private set

    /** The current average time measure has taken during prefetches of this LazyLayout. */
    var averageMeasureTimeNanos: Long = 0L
        private set

    /**
     * Executes the [doComposition] block and updates [averageCompositionTimeNanos] with the new
     * average.
     */
    internal inline fun recordCompositionTiming(contentType: Any?, doComposition: () -> Unit) {
        val executionTime = measureTime(doComposition)
        contentType?.let {
            val currentAvgCompositionTimeNanos =
                averageCompositionTimeNanosByContentType.getOrDefault(contentType, 0L)
            val newAvgCompositionTimeNanos =
                calculateAverageTime(executionTime.inWholeNanoseconds, currentAvgCompositionTimeNanos)
            averageCompositionTimeNanosByContentType[contentType] = newAvgCompositionTimeNanos
        }
        averageCompositionTimeNanos =
            calculateAverageTime(executionTime.inWholeNanoseconds, averageCompositionTimeNanos)
    }

    /**
     * Executes the [doMeasure] block and updates [averageMeasureTimeNanos] with the new average.
     */
    internal inline fun recordMeasureTiming(contentType: Any?, doMeasure: () -> Unit) {
        val executionTime = measureTime(doMeasure)
        contentType?.let {
            val currentAvgMeasureTimeNanos =
                averageMeasureTimeNanosByContentType.getOrDefault(contentType, 0L)
            val newAvgMeasureTimeNanos =
                calculateAverageTime(executionTime.inWholeNanoseconds, currentAvgMeasureTimeNanos)
            averageMeasureTimeNanosByContentType[contentType] = newAvgMeasureTimeNanos
        }
        averageMeasureTimeNanos = calculateAverageTime(executionTime.inWholeNanoseconds, averageMeasureTimeNanos)
    }

    private fun calculateAverageTime(new: Long, current: Long): Long {
        // Calculate a weighted moving average of time taken to compose an item. We use weighted
        // moving average to bias toward more recent measurements, and to minimize storage /
        // computation cost. (the idea is taken from RecycledViewPool)
        return if (current == 0L) {
            new
        } else {
            // dividing first to avoid a potential overflow
            current / 4 * 3 + new / 4
        }
    }
}

@ExperimentalFoundationApi
private object DummyHandle : PrefetchHandle {
    override fun cancel() {}

    override fun markAsUrgent() {}
}

/**
 * PrefetchHandleProvider is used to connect the [LazyLayoutPrefetchState], which provides the API
 * to schedule prefetches, to a [LazyLayoutItemContentFactory] which resolves key and content from
 * an index, [SubcomposeLayoutState] which knows how to precompose/premeasure, and a specific
 * [PrefetchScheduler] used to execute a request.
 */
@ExperimentalFoundationApi
internal class PrefetchHandleProvider(
    private val itemContentFactory: LazyLayoutItemContentFactory,
    private val subcomposeLayoutState: SubcomposeLayoutState,
    private val executor: PrefetchScheduler
) {
    fun schedulePrefetch(
        index: Int,
        constraints: Constraints,
        prefetchMetrics: PrefetchMetrics
    ): PrefetchHandle =
        HandleAndRequestImpl(index, constraints, prefetchMetrics).also {
            executor.schedulePrefetch(it)
        }

    fun createNestedPrefetchRequest(
        index: Int,
        constraints: Constraints,
        prefetchMetrics: PrefetchMetrics,
    ): PrefetchRequest = HandleAndRequestImpl(index, constraints = constraints, prefetchMetrics)

    @ExperimentalFoundationApi
    private inner class HandleAndRequestImpl(
        private val index: Int,
        private val constraints: Constraints,
        private val prefetchMetrics: PrefetchMetrics,
    ) : PrefetchHandle, PrefetchRequest {

        private var precomposeHandle: SubcomposeLayoutState.PrecomposedSlotHandle? = null
        private var isMeasured = false
        private var isCanceled = false
        private val isComposed
            get() = precomposeHandle != null

        private var hasResolvedNestedPrefetches = false
        private var nestedPrefetchController: NestedPrefetchController? = null
        private var isUrgent = false

        private val isValid
            get() = !isCanceled && index in 0 until itemContentFactory.itemProvider().itemCount

        override fun cancel() {
            if (!isCanceled) {
                isCanceled = true
                precomposeHandle?.dispose()
                precomposeHandle = null
            }
        }

        override fun markAsUrgent() {
            isUrgent = true
        }

        private fun PrefetchRequestScope.shouldExecute(average: Long): Boolean {
            val available = availableTimeNanos()
            // even for urgent request we only do the work if we have time available, as otherwise
            // it is better to just return early to allow the next frame to start and do the work.
            return (isUrgent && available > 0) || average < available
        }

        override fun PrefetchRequestScope.execute(): Boolean {
            if (!isValid) {
                return false
            }

            val contentType = itemContentFactory.itemProvider().getContentType(index)

            if (!isComposed) {
                val estimatedPrecomposeTime: Long =
                    if (
                        contentType != null &&
                            prefetchMetrics.averageCompositionTimeNanosByContentType.contains(
                                contentType
                            )
                    )
                        prefetchMetrics.averageCompositionTimeNanosByContentType[contentType]
                    else prefetchMetrics.averageCompositionTimeNanos
                if (shouldExecute(estimatedPrecomposeTime)) {
                    prefetchMetrics.recordCompositionTiming(contentType) {
                        trace("compose:lazy:prefetch:compose") { performComposition() }
                    }
                } else {
                    return true
                }
            }

            // if the request is urgent we better proceed with the measuring straight away instead
            // of spending time trying to split the work more via nested prefetch. nested prefetch
            // is always an estimation and it could potentially do work we will not need in the end,
            // but the measuring will only do exactly the needed work (including composing nested
            // lazy layouts)
            if (!isUrgent) {
                // Nested prefetch logic is best-effort: if nested LazyLayout children are
                // added/removed/updated after we've resolved nested prefetch states here or
                // resolved
                // nestedPrefetchRequests below, those changes won't be taken into account.
                if (!hasResolvedNestedPrefetches) {
                    if (availableTimeNanos() > 0) {
                        trace("compose:lazy:prefetch:resolve-nested") {
                            nestedPrefetchController = resolveNestedPrefetchStates()
                            hasResolvedNestedPrefetches = true
                        }
                    } else {
                        return true
                    }
                }

                val hasMoreWork =
                    nestedPrefetchController?.run { executeNestedPrefetches() } ?: false
                if (hasMoreWork) {
                    return true
                }
            }

            if (!isMeasured && !constraints.isZero) {
                val estimatedPremeasureTime: Long =
                    if (
                        contentType != null &&
                            prefetchMetrics.averageMeasureTimeNanosByContentType.contains(
                                contentType
                            )
                    )
                        prefetchMetrics.averageMeasureTimeNanosByContentType[contentType]
                    else prefetchMetrics.averageMeasureTimeNanos
                if (shouldExecute(estimatedPremeasureTime)) {
                    prefetchMetrics.recordMeasureTiming(contentType) {
                        trace("compose:lazy:prefetch:measure") { performMeasure(constraints) }
                    }
                } else {
                    return true
                }
            }

            // All our work is done
            return false
        }

        private fun performComposition() {
            requirePrecondition(isValid) {
                "Callers should check whether the request is still valid before calling " +
                    "performComposition()"
            }
            requirePrecondition(precomposeHandle == null) { "Request was already composed!" }
            val itemProvider = itemContentFactory.itemProvider()
            val key = itemProvider.getKey(index)
            val contentType = itemProvider.getContentType(index)
            val content = itemContentFactory.getContent(index, key, contentType)
            precomposeHandle = subcomposeLayoutState.precompose(key, content)
        }

        private fun performMeasure(constraints: Constraints) {
            requirePrecondition(!isCanceled) {
                "Callers should check whether the request is still valid before calling " +
                    "performMeasure()"
            }
            requirePrecondition(!isMeasured) { "Request was already measured!" }
            isMeasured = true
            val handle =
                requirePreconditionNotNull(precomposeHandle) {
                    "performComposition() must be called before performMeasure()"
                }
            repeat(handle.placeablesCount) { placeableIndex ->
                handle.premeasure(placeableIndex, constraints)
            }
        }

        private fun resolveNestedPrefetchStates(): NestedPrefetchController? {
            val precomposedSlotHandle =
                requirePreconditionNotNull(precomposeHandle) {
                    "Should precompose before resolving nested prefetch states"
                }

            var nestedStates: MutableList? = null
            precomposedSlotHandle.traverseDescendants(TraversablePrefetchStateNodeKey) {
                val prefetchState = (it as TraversablePrefetchStateNode).prefetchState
                nestedStates =
                    nestedStates?.apply { add(prefetchState) } ?: mutableListOf(prefetchState)
                TraverseDescendantsAction.SkipSubtreeAndContinueTraversal
            }
            return nestedStates?.let { NestedPrefetchController(it) }
        }

        override fun toString(): String =
            "HandleAndRequestImpl { index = $index, constraints = $constraints, " +
                "isComposed = $isComposed, isMeasured = $isMeasured, isCanceled = $isCanceled }"

        private inner class NestedPrefetchController(
            private val states: List
        ) {

            // This array is parallel to nestedPrefetchStates, so index 0 in nestedPrefetchStates
            // corresponds to index 0 in this array, etc.
            private val requestsByState: Array?> = arrayOfNulls(states.size)
            private var stateIndex: Int = 0
            private var requestIndex: Int = 0

            init {
                requirePrecondition(states.isNotEmpty()) {
                    "NestedPrefetchController shouldn't be created with no states"
                }
            }

            fun PrefetchRequestScope.executeNestedPrefetches(): Boolean {
                if (stateIndex >= states.size) {
                    return false
                }
                checkPrecondition(!isCanceled) {
                    "Should not execute nested prefetch on canceled request"
                }

                trace("compose:lazy:prefetch:nested") {
                    while (stateIndex < states.size) {
                        if (requestsByState[stateIndex] == null) {
                            if (availableTimeNanos() <= 0) {
                                // When we have time again, we'll resolve nested requests for this
                                // state
                                return true
                            }

                            requestsByState[stateIndex] =
                                states[stateIndex].collectNestedPrefetchRequests()
                        }

                        val nestedRequests = requestsByState[stateIndex]!!
                        while (requestIndex < nestedRequests.size) {
                            val hasMoreWork = with(nestedRequests[requestIndex]) { execute() }
                            if (hasMoreWork) {
                                return true
                            } else {
                                requestIndex++
                            }
                        }

                        requestIndex = 0
                        stateIndex++
                    }
                }

                return false
            }
        }
    }
}

private const val TraversablePrefetchStateNodeKey =
    "androidx.compose.foundation.lazy.layout.TraversablePrefetchStateNode"

/**
 * A modifier which lets the [LazyLayoutPrefetchState] for a [LazyLayout] to be discoverable via
 * [TraversableNode] traversal.
 */
@ExperimentalFoundationApi
internal fun Modifier.traversablePrefetchState(
    lazyLayoutPrefetchState: LazyLayoutPrefetchState?
): Modifier {
    return lazyLayoutPrefetchState?.let { this then TraversablePrefetchStateModifierElement(it) }
        ?: this
}

@ExperimentalFoundationApi
private class TraversablePrefetchStateNode(
    var prefetchState: LazyLayoutPrefetchState,
) : Modifier.Node(), TraversableNode {

    override val traverseKey: String = TraversablePrefetchStateNodeKey
}

@ExperimentalFoundationApi
private data class TraversablePrefetchStateModifierElement(
    private val prefetchState: LazyLayoutPrefetchState,
) : ModifierNodeElement() {
    override fun create() = TraversablePrefetchStateNode(prefetchState)

    override fun update(node: TraversablePrefetchStateNode) {
        node.prefetchState = prefetchState
    }

    override fun InspectorInfo.inspectableProperties() {
        name = "traversablePrefetchState"
        value = prefetchState
    }
}

private val ZeroConstraints = Constraints(maxWidth = 0, maxHeight = 0)