@Composable
fun LazyTreeDemo() {
    // Create a simple tree structure just to demo.
    // The lambda to rememberLazyTree is a composable lambda, but instead of emitting UI directly,
    // it emits tree nodes. And each tree node takes a composable function that _can_ emit the UI
    // for that node.
    val tree = rememberLazyTree {
        // Note that the child lambdas for Branch nodes aren't until the node is expanded, and they
        // are removed from the composition when it's collapsed.
        Branch({ Text("Root") }) {
            Leaf { Text("Header") }
            Branch({ Text("Second child") }) {
                Leaf { Text("Second-level leaf") }
            }
            Leaf { Text("Footer") }
        }
    }

    // This column will host the UI for the tree created above.
    LazyColumn {
        item {
            Text("This is a lazy tree.")
        }

        // Add a flattened representation of the expanded nodes from the tree to the column.
        // The composable lambda here will be used to wrap each node in the tree.
        treeItems(tree) { content ->
            Row(verticalAlignment = CenterVertically) {
                // Indent each item proportional to its depth. This should match up with the size
                // of the toggle button probably, it doesn't right now to keep the code simpler.
                Spacer(Modifier.width(depth * 48.dp))

                // Only show the expand/collapse toggle button for branch nodes. These properties
                // and the setExpanded function are from the LazyItemTreeScope.
                if (isExpandable) {
                    IconToggleButton(
                        checked = isExpanded,
                        onCheckedChange = ::setExpanded
                    ) {
                        val angle by animateFloatAsState(if (isExpanded) 0f else -90f)
                        Icon(
                            Icons.Default.ArrowDropDown,
                            contentDescription = if (isExpanded) "Collapse node" else "Expand node",
                            modifier = Modifier.graphicsLayer { rotationZ = angle }
                        )
                    }
                    Spacer(Modifier.width(8.dp))
                }

                // Compose the node's actual UI in the rest of the row.
                content()
            }
        }
    }
}

import androidx.compose.animation.core.animateFloatAsState
import androidx.compose.foundation.layout.Row
import androidx.compose.foundation.layout.Spacer
import androidx.compose.foundation.layout.width
import androidx.compose.foundation.lazy.LazyColumn
import androidx.compose.foundation.lazy.LazyListScope
import androidx.compose.foundation.lazy.items
import androidx.compose.material.Icon
import androidx.compose.material.IconToggleButton
import androidx.compose.material.Text
import androidx.compose.material.icons.Icons
import androidx.compose.material.icons.filled.ArrowDropDown
import androidx.compose.runtime.AbstractApplier
import androidx.compose.runtime.Composable
import androidx.compose.runtime.ComposableTargetMarker
import androidx.compose.runtime.ComposeNode
import androidx.compose.runtime.Composition
import androidx.compose.runtime.Stable
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateListOf
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.rememberCompositionContext
import androidx.compose.runtime.saveable.rememberSaveable
import androidx.compose.runtime.setValue
import androidx.compose.ui.Alignment.Companion.CenterVertically
import androidx.compose.ui.Modifier
import androidx.compose.ui.UiComposable
import androidx.compose.ui.graphics.graphicsLayer
import androidx.compose.ui.unit.dp
import androidx.compose.ui.unit.times

/**
 * The receiver type of [TreeBuilderComposable] composable functions. Carries information about the
 * current depth of the tree. This could also be done via composition locals, but since this is just
 * a regular function parameter it's a lot cheaper. It's an API cleanliness/efficiency tradeoff.
 */
@JvmInline
value class TreeBuilderScope(val depth: Int)

/**
 * Indicates that a composable function can only emit [Branch] and [Leaf] nodes, not UI.
 */
@ComposableTargetMarker
@Retention(AnnotationRetention.BINARY)
@Target(
    AnnotationTarget.FUNCTION,
    AnnotationTarget.PROPERTY_GETTER,
    AnnotationTarget.TYPE,
    AnnotationTarget.TYPE_PARAMETER,
)
annotation class TreeBuilderComposable

/**
 * Represents a tree structure that can be displayed inside a lazy list via [treeItems].
 * Created by [rememberLazyTree].
 */
@Stable
class LazyTreeContent internal constructor(internal val roots: List<TreeNode>)

/**
 * Creates a subcomposition that will compose [content] to determine the contents of the tree
 * to show in a lazy list via [treeItems].
 *
 * @param content A composable function that, instead of emitting UI directly, emits [Branch] and
 * [Leaf] nodes that define the structure of a tree.
 */
@Composable
fun rememberLazyTree(
    content: @Composable @TreeBuilderComposable TreeBuilderScope.() -> Unit
): LazyTreeContent {
    val applier = remember { TreeApplier() }
    val compositionContext = rememberCompositionContext()
    val composition = remember(applier, compositionContext) {
        Composition(applier, compositionContext)
    }
    composition.setContent { TreeBuilderScope(0).content() }
    return remember(applier) { LazyTreeContent(applier.children) }
}

/**
 * Contains a number of properties about the current node that can be used to indicate the state of
 * the node in its UI.
 */
interface LazyTreeItemScope {
    /** The depth of the node in the tree. */
    val depth: Int

    /** Whether the node can be expanded. */
    val isExpandable: Boolean

    /**
     * Whether the node is currently expanded and showing its children. Never true if
     * [isExpandable] is false
     */
    val isExpanded: Boolean

    /**
     * Requests a change to the value of [isExpanded].
     */
    fun setExpanded(expanded: Boolean)
}

/**
 * Adds the tree items from [content] to the lazy list.
 *
 * @param content The model of the tree to display, as created by [rememberLazyTree].
 * @param rowContent A wrapper composable that can add UI like branch toggle buttons and indentation
 * around each node, as required.
 */
fun LazyListScope.treeItems(
    content: LazyTreeContent,
    rowContent: @Composable LazyTreeItemScope.(
        content: @Composable LazyTreeItemScope.() -> Unit
    ) -> Unit
) {
    // The compose runtime has already flattened the tree into a flat list, so we can just add that
    // list to the lazy one. This roots property is a snapshot state list, so whenever the tree
    // structure changes the lazy list will automatically update.
    items(content.roots) { item ->
        // TreeNode implements the LazyTreeItemScope interface itself.
        item.rowContent(item.content)
    }
}

/**
 * Emits a branch node of the tree. Branch nodes are expandable and, when expanded, run their
 * [children] function to determine the structure of the subtree beneath them.
 *
 * This composable function can only be called from inside a [rememberLazyTree] composition, not
 * from regular UI composables.
 *
 * This overload manages the expanded/collapsed state internally. There's another overload if you
 * want to hoist it and manage it yourself.
 *
 * @param content A Compose UI composable function that emits the UI for the node. This function
 * can emit things like text, buttons, and anything else that can exist in a regular Compose UI.
 * @param children A composable function that can make other [Branch] and [Leaf] calls to define
 * the structure of the subtree below this node. It can _not_ emit UI composables like text.
 */
@Composable
@TreeBuilderComposable
fun TreeBuilderScope.Branch(
    content: @Composable @UiComposable LazyTreeItemScope.() -> Unit,
    children: @Composable @TreeBuilderComposable TreeBuilderScope.() -> Unit = {}
) {
    var expanded by rememberSaveable { mutableStateOf(false) }
    Branch(
        expanded = expanded,
        onToggleExpanded = { expanded = it },
        content = content,
        children = children
    )
}

/**
 * Emits a branch node of the tree. Branch nodes are expandable and, when expanded, run their
 * [children] function to determine the structure of the subtree beneath them.
 *
 * This composable function can only be called from inside a [rememberLazyTree] composition, not
 * from regular UI composables.
 *
 * @param expanded If true, the node will be shown with a UI indicator that it's expanded, and the
 * [children] function will be ran to compose its children. If false, [children] will not be
 * composed, and the UI will indicate that the node is collapsed.
 * @param onToggleExpanded Called when the node is expanded and collapsed. This callback should
 * update some state that causes [expanded] to be passed as the new value on the next composition.
 * @param content A Compose UI composable function that emits the UI for the node. This function
 * can emit things like text, buttons, and anything else that can exist in a regular Compose UI.
 * @param children A composable function that can make other [Branch] and [Leaf] calls to define
 * the structure of the subtree below this node. It can _not_ emit UI composables like text.
 */
@Composable
@TreeBuilderComposable
fun TreeBuilderScope.Branch(
    expanded: Boolean,
    onToggleExpanded: (Boolean) -> Unit,
    content: @Composable @UiComposable LazyTreeItemScope.() -> Unit,
    children: @Composable @TreeBuilderComposable TreeBuilderScope.() -> Unit = {}
) {
    ComposeNode<TreeNode, TreeApplier>(
        factory = { TreeNode(depth, isExpandable = true) },
        update = {
            set(content) { this.content = it }
            set(expanded) { this._isExpanded = it }
            set(onToggleExpanded) { this.onToggleExpanded = onToggleExpanded }
        }
    )

    if (expanded) {
        // Typically the children of a ComposeNode are passed as the children parameter to the
        // ComposeNode function. That has the effect of making any nodes they emit children of this
        // ComposeNode. However, in this case, we don't want to emit a tree structure, we want to
        // emit a flattened tree where the children of a node show up as its siblings.
        // In other words, we're making the emitted structure mirror the structure of the slot table
        // itself.
        TreeBuilderScope(depth   1).children()
    }
}

/**
 * Emits a leaf node of the tree. Leaf nodes are not expandable and never have children.
 *
 * This composable function can only be called from inside a [rememberLazyTree] composition, not
 * from regular UI composables.
 *
 * @param content A Compose UI composable function that emits the UI for the node. This function
 * can emit things like text, buttons, and anything else that can exist in a regular Compose UI.
 */
@Composable
@TreeBuilderComposable
fun TreeBuilderScope.Leaf(content: @Composable @UiComposable LazyTreeItemScope.() -> Unit) {
    ComposeNode<TreeNode, TreeApplier>(
        factory = { TreeNode(depth, isExpandable = false) },
        update = {
            // Leaf nodes aren't expandable so we don't need to set any of the other properties.
            set(content) { this.content = it }
        }
    )
}

/**
 * Represents information for the UI about each node of the tree. These nodes are emitted by
 * [ComposeNode] calls in the composable functions above that define the tree-building DSL, and are
 * collected by the [TreeApplier] class.
 *
 * This is also internal to the library. The public API for creating these nodes is the [Branch] and
 * [Leaf] composable functions.
 *
 * @param depth The depth of the node in the tree, used to calculate indentation in layout.
 * @param isExpandable Whether an expansion toggle should be shown for this node. If true, it might
 * have children, but that might not be known until the children are actually requested.
 */
internal class TreeNode(
    override val depth: Int,
    override val isExpandable: Boolean
) : LazyTreeItemScope {
    /**
     * Backing property for [isExpanded] so that the setter doesn't clash with the [setExpanded]
     * function.
     */
    @Suppress("PropertyName")
    var _isExpanded: Boolean by mutableStateOf(false)
    override val isExpanded: Boolean get() = _isExpanded

    override fun setExpanded(expanded: Boolean) {
        onToggleExpanded(expanded)
    }

    var onToggleExpanded: (Boolean) -> Unit by mutableStateOf({})
    var content: @Composable @UiComposable LazyTreeItemScope.() -> Unit by mutableStateOf({})

    override fun toString(): String = "TreeNode("  
        "depth=$depth, "  
        "isExpandable=$isExpandable, "  
        "isExpanded=$isExpanded"  
        ")"
}

/**
 * This class is not part of this library's public API. It's used by the Compose runtime to convert
 * [ComposeNode] calls to the actual structure of emitted nodes. Typically that's a tree, but in
 * this case, since we are flattening a tree, we just build a flat list of [TreeNode]s.
 * The root "node" for the applier is simply null, and every method includes an assertion to make
 * sure that nothing is trying to create a tree.
 */
private class TreeApplier : AbstractApplier<TreeNode?>(null) {
    /**
     * This is a mutable _state_ list so that the output from the tree-building subcomposition can
     * be observed by the [LazyColumn]'s layout.
     */
    val children = mutableStateListOf<TreeNode>()

    override fun insertTopDown(index: Int, instance: TreeNode?) {
        checkNotNull(instance)
        check(current == null)
        children.add(index, instance)
    }

    override fun insertBottomUp(index: Int, instance: TreeNode?) {
        // Only this or insertTopDown should be implemented. In our case it doesn't matter since
        // we're building a list, not a tree, so either one works and I just picked topDown
        // arbitrarily.
    }

    override fun remove(index: Int, count: Int) {
        check(current == null)
        children.removeRange(index, index   count)
    }

    override fun move(from: Int, to: Int, count: Int) {
        check(current == null)
        // This move helper function is defined in AbstractApplier, but only for lists of the exact
        // type of the nodes of this applier, so we have to cast our list type to get access to it.
        @Suppress("UNCHECKED_CAST")
        (children as MutableList<TreeNode?>).move(from, to, count)
    }

    override fun onClear() {
        check(current == null)
        children.clear()
    }
}