DebugStateObservation.kt

import androidx.compose.runtime.Composable
import androidx.compose.runtime.Composer
import androidx.compose.runtime.CompositionTracer
import androidx.compose.runtime.DisposableEffect
import androidx.compose.runtime.InternalComposeApi
import androidx.compose.runtime.InternalComposeTracingApi
import androidx.compose.runtime.remember
import androidx.compose.runtime.snapshots.MutableSnapshot
import androidx.compose.runtime.snapshots.Snapshot
import java.util.WeakHashMap

class DebugStateObservation(val id: String) {
    class ObservationEntry(val exception: Exception, val functionName: String?)
    private val map = WeakHashMap<Any, MutableList<ObservationEntry>>()

    val readObserver = { value: Any, function: String? ->
        synchronized(this) {
            val e = Exception()
            val list = map.getOrPut(value) { mutableListOf() }
            list += ObservationEntry(e, function)
        }
    }

    fun print(changes: Set<Any>) {
        synchronized(this) {
            val affected = map.keys.intersect(changes)
            if (affected.isNotEmpty()) {
                affected.forEach {
                    printStateChange(id, it, map[it])
                }
            }
        }
    }

    fun clear() {
        synchronized(this) {
            map.clear()
        }
    }
}

private fun printStateChange(
    id: String,
    state: Any,
    entries: List<DebugStateObservation.ObservationEntry>?
) {
    val traces = entries?.joinToString(separator = "\n") { entry ->
        // remove trace start, sample:
        // at androidx.compose.foundation.demos.DebugStateObservation$readObserver$1.invoke(Test.kt:33)
        // at androidx.compose.foundation.demos.DebugStateObservation$readObserver$1.invoke(Test.kt:31)
        // at androidx.compose.runtime.snapshots.SnapshotKt$mergedReadObserver$1.invoke(Snapshot.kt:1771)
        // at androidx.compose.runtime.snapshots.SnapshotKt$mergedReadObserver$1.invoke(Snapshot.kt:1770)
        // at androidx.compose.runtime.snapshots.SnapshotKt.readable(Snapshot.kt:2003)
        // at androidx.compose.runtime.SnapshotMutableIntStateImpl.getIntValue(SnapshotIntState.kt:138)
        val stackTrace = entry.exception.stackTrace
        buildString {
            if (entry.functionName != null) {
                append("Inside function named: ")
                appendLine(entry.functionName)
            }
            for (i in 10.. minOf(15, stackTrace.size)) {
                appendLine("\tat ${stackTrace[i]}")
            }
            append("\t...")
        }
    } ?: ""
    println("$id might recompose because $state changed, last read at:\n$traces")
}

@PublishedApi
internal val LocalDebugObservationRegistry = ThreadLocal<DebugObservationRegistry?>()

/**
 * Records state observations inside @Composable [block] and prints to [System.out] whenever
 * state mutation is applied.
 *
 * NOTE: This doesn't record recompositions precisely and only uses snapshot system to record state
 * mutations that /might/ invalidate recomposition. Consecutive invocations might result in
 * different results depending on functions that were run / skipped during each execution. To be
 * used directly inside a function scope that recomposes, as Compose might skip inner scopes and
 * reads/mutations are not going to be recorded.
 */
@OptIn(InternalComposeApi::class)
@Composable
inline fun DebugStateChanges(id: String, crossinline block: @Composable () -> Unit) {
    val parentRegistry = LocalDebugObservationRegistry.get()

    var registry = parentRegistry
    val snapshot = if (registry == null) {
        val currentSnapshot = Snapshot.current
        registry = remember { DebugObservationRegistry() }
        if (currentSnapshot is MutableSnapshot) {
            currentSnapshot.takeNestedMutableSnapshot(registry.readObserver)
        } else {
            currentSnapshot.takeNestedSnapshot(registry.readObserver)
        }
    } else {
        Snapshot.current
    }

    val observation = remember(id) { DebugStateObservation(id) }
    DisposableEffect(observation) {
        val disposeHandle = Snapshot.registerApplyObserver { changes, _ ->
            observation.print(changes)
        }
        onDispose {
            observation.clear()
            disposeHandle.dispose()
        }
    }

    observation.clear()

    registry.enter(observation)
    snapshot.runAndDispose {
        if (parentRegistry == null) {
            registry.start()
            LocalDebugObservationRegistry.set(registry)
            block()
            LocalDebugObservationRegistry.set(null)
            registry.end()
        } else {
            block()
        }
    }
    registry.pop()
}

@OptIn(InternalComposeTracingApi::class)
class DebugObservationRegistry : CompositionTracer {
    private val stack = ArrayDeque<DebugStateObservation>()
    private val functionStack = ArrayDeque<String>()

    val readObserver = { it: Any ->
        stack.last().readObserver(it, functionStack.removeLastOrNull())
    }

    fun enter(observation: DebugStateObservation) {
        stack += observation
    }

    fun pop() {
        stack.removeLastOrNull()
    }

    override fun isTraceInProgress(): Boolean = true

    override fun traceEventEnd() {
        functionStack.removeLastOrNull()
    }

    override fun traceEventStart(key: Int, dirty1: Int, dirty2: Int, info: String) {
        functionStack.add(info)
    }

    fun start() { Composer.setTracer(this) }
    fun end() { Composer.setTracer(NoopTracer) }
}

@OptIn(InternalComposeTracingApi::class)
object NoopTracer : CompositionTracer {
    override fun isTraceInProgress(): Boolean = false

    override fun traceEventEnd() {}

    override fun traceEventStart(key: Int, dirty1: Int, dirty2: Int, info: String) {}

}

// Compose doesn't work with try/finally, but we don't really use it for catching things.
@PublishedApi
internal inline fun <T> Snapshot.runAndDispose(block: () -> T): T =
    try {
        enter(block)
    } finally {
        dispose()
    }