Debug State Observation

/*
Copyright (c) 2025 Andrei Shikov

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

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()
    }

// Just wrap the composable that you need more
// information about with the following Composable

DebugStateObservation("MyCustomComposable") {
    MyCustomComposable(...)
}