Modifier.Node Version of Modifier.drawWithCache

CachedDrawModifierNode.kt

@file:OptIn(ExperimentalComposeUiApi::class)
@file:Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")

package net.peanuuutz.compose.desktop

import androidx.compose.ui.ExperimentalComposeUiApi
import androidx.compose.ui.Modifier
import androidx.compose.ui.draw.BuildDrawCacheParams
import androidx.compose.ui.draw.CacheDrawScope
import androidx.compose.ui.draw.DrawResult
import androidx.compose.ui.geometry.Size
import androidx.compose.ui.graphics.drawscope.ContentDrawScope
import androidx.compose.ui.node.DrawModifierNode
import androidx.compose.ui.node.Nodes
import androidx.compose.ui.node.ObserverNode
import androidx.compose.ui.node.observeReads
import androidx.compose.ui.node.requireCoordinator
import androidx.compose.ui.node.requireLayoutNode
import androidx.compose.ui.unit.Density
import androidx.compose.ui.unit.LayoutDirection
import androidx.compose.ui.unit.toSize

sealed interface CachedDrawModifierNode : DrawModifierNode {
    fun updateCacheBuilder(onBuildCache: CacheDrawScope.() -> DrawResult)

    fun invalidate()
}

fun CachedDrawModifierNode(
    onBuildCache: CacheDrawScope.() -> DrawResult
): CachedDrawModifierNode {
    return CachedDrawModifierNodeImpl(onBuildCache)
}

// ======== Internal ========

private class CachedDrawModifierNodeImpl(
    onBuildCache: CacheDrawScope.() -> DrawResult
) : Modifier.Node(),
    CachedDrawModifierNode,
    ObserverNode
{
    private var onBuildCache: CacheDrawScope.() -> DrawResult = onBuildCache
        set(value) {
            if (field == value) {
                return
            }
            field = value
            invalidate()
        }

    private val cacheDrawScope: CacheDrawScope = CacheDrawScope()

    init {
        cacheDrawScope.cacheParams = NodeBuildDrawCacheParams()
    }

    private var invalidated: Boolean = true

    override fun updateCacheBuilder(onBuildCache: CacheDrawScope.() -> DrawResult) {
        this.onBuildCache = onBuildCache
    }

    override fun invalidate() {
        invalidated = true
    }

    override fun ContentDrawScope.draw() {
        if (invalidated) {
            observeReads {
                with(cacheDrawScope) {
                    drawResult = null
                    onBuildCache()
                    checkNotNull(drawResult) {
                        "DrawResult not defined. Did you forget to call onDraw?"
                    }
                }
            }
            invalidated = false
        }
        cacheDrawScope.drawResult
            ?.block
            ?.let { block -> block() }
    }

    override fun onMeasureResultChanged() {
        invalidate()
    }

    override fun onObservedReadsChanged() {
        invalidate()
    }

    private inner class NodeBuildDrawCacheParams : BuildDrawCacheParams {
        override val density: Density
            get() = requireLayoutNode().density

        override val layoutDirection: LayoutDirection
            get() = requireLayoutNode().layoutDirection

        override val size: Size
            get() = requireCoordinator(Nodes.LayoutAware).size.toSize()
    }
}