Consider supporting a composable tree builder

[zach-klippenstein]

Compose vs Plain Old DSL

The current DSL for building trees looks a lot like compose code:

Branch {
  Leaf()
  Branch {
    Leaf()
  }
  Leaf()
}

Compose is, in fact, a tree-building library. The core runtime knows nothing about UI, it can be used to build any type of tree. For example, the rememberVectorPainter composable has a content parameter that emits not UI nodes, but vector primitives (think SVG elements).

Since one of the main tasks of this library is to express a tree structure, Compose itself might be a good fit to use for the DSL. Some of the advantages over a simpler, non-compose approach are:

• Effects - tree nodes that need to load data asynchronously when expanded can do so by performing the load in an effect that is only composed when expanded.
• State – nodes can store their own state in the composition. Comes in handy with the above.
• Donut-hole skipping – If a node's children can change over time, using Compose for the DSL will ensure only the nodes that change are updated.

There's one other benefit of using Compose, which is that it makes it really easy to make better use of the LazyColumn that is actually used to structure the tree UI.

Tree flattening

To get the full benefit of a LazyColumn, every expanded node in the tree (everything that has a "row" in the lazy column) should have its own lazy item so the column can effectively recycle and reuse bits of UI that aren't on the screen. To do that, you need to flatten a tree structure into a single list that can be processed by the LazyColumn.

This also happens to be one of the core jobs of Compose. It takes the call graph of Composable functions (a tree) and flattens them into a single array-like structure (the slot table). So we can take advantage of that to flatten a tree for the LazyColumn.

Proof-of-concept

I've thrown together a quick sketch of how this could look. I didn't try too hard to make it match all the exact APIs of this library, but just wanted to show how a custom subcomposition could be used and consumed by the LazyColumn.

The main entry point to the API in this demo is the rememberLazyTree function:

val tree = rememberLazyTree {
  Leaf({ /* UI */ })
  Branch({ /* UI */ }) {
    Leaf({ /* UI */ })
  }
  Leaf({ /* UI */ })
}

Then, you can take the tree object returned by that function and pass it to a LazyColumn:

LazyColumn {
  treeItems(tree) { content ->
    // Wrap content with the indent and expand/collapse UI
    …
    content()
}

This is probably something you'd want to hide from the public API of this library, but it shows how flexible and generic this approach can be (it could work in lazy rows, or even grids, I guess).

Code

Usage demo

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

Note the @UiComposable and @TreeNodeComposable annotations that document to the reader, and tell the compiler, what type of nodes can be emitted by each composable function.

Also, I omitted the kotlin DSL annotations from the TreeBuilderScope type, but you'd probably want to have those on there to make sure that children can't accidentally refer to the wrong scope and get the wrong depth.

Implementation

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

[adrielcafe]

Amazing! Thanks for that @zach-klippenstein. I'll learn from your code and incorporate into the library. Stay tuned!

[adrielcafe]

It works great 🎉

Now I'm trying to solve some issues related to expansion:

1. How to get the parent node?

tree.expandNode(node) uses the parent to recursively expand the upper nodes. The solution I came up was:

@Composable
public fun <T> TreeScope<T>.Branch(
    content: T,
    name: String = content.toString(),
    children: @Composable TreeScope<T>.() -> Unit = {}
) {
    var node by remember { mutableStateOf<BranchNode<T>?>(null) }

    ComposeNode<BranchNode<T>, TreeApplier<T>>(
        factory = {
            SimpleBranchNode(...)
            	.also { node = it }
        },
        update = { ... }
    )


    if (isExpanded) {
        TreeScope(depth = depth.inc(), parent = node).children()
    }
}

I believe it will work, but couldn't test it yet because of the next issue. Will keep looking for a solution that doesn't depend on a state.

2. How to expand a node at X depth?

tree.expandNode(node) can expand a node at depth 3 or 10, for example. But now, since the children only exist when they are composed, I still don't know how can I get an instance of node (because it only exists when is visible).

And expandUntil(maxDepth) should expand until X depth, but the function below only expands the root nodes, since it's children wasn't composed yet.

private fun expandDown(nodes: List<Node<T>>, maxDepth: Int) {    
    nodes
        .asSequence()
        .filterIsInstance<BranchNode<T>>()
        .filter { it.depth <= maxDepth }
        .sortedBy { it.depth }
        .forEach { it.setExpanded(true) }
}

Any ideas?

[adrielcafe]

Another issue I'm trying to fix:
Now that I have a custom composition, it isn't possible anymore to animate the children's node because AnimatedVisibility can't be used inside my composition.

When I try to use the following code I get java.lang.ClassCastException: androidx.compose.ui.node.LayoutNode cannot be cast to cafe.adriel.bonsai.core.node.Node

@Composable
public fun <T> TreeScope<T>.Branch(...) {
    ComposeNode<BranchNode<T>, TreeApplier<T>>(...)

    AnimatedVisibility(visible = isExpanded) {
        TreeScope(depth, parent).children()
    }
}

Unfortunatelly LazyColumn still don't supports animated insertions, only reordering.

[adrielcafe]

About the issue #2 (How to expand a node at X depth?), I'm using the TreeScope to propagate the expansion until a maximum depth. It works very well.

@Composable
public fun <T> TreeScope.Branch(...) {
    val (isExpanded, setExpanded) = rememberSaveable { mutableStateOf(isExpanded && depth <= expandMaxDepth) }
    val (expandMaxDepth, setExpandMaxDepth) = rememberSaveable { mutableStateOf(expandMaxDepth) }

    ComposeNode<BranchNode<T>, TreeApplier<T>>(
        factory = {
            BranchNode(...)
        },
        update = {
            set(isExpanded) { this.isExpandedState = isExpanded }
            set(setExpanded) {
                this.onToggleExpanded = { isExpanded, maxDepth ->
                    setExpanded(isExpanded)
                    setExpandMaxDepth(maxDepth)
                }
            }
        }
    )

    if (isExpanded && depth <= expandMaxDepth) {
        TreeScope(
            depth = depth.inc(),
            isExpanded = isExpanded,
            expandMaxDepth = expandMaxDepth
        ).children()
    }
}

[adrielcafe]

Just released v1.2.0 with the custom composition.

Still need to find a way to animate expanded/collapsed nodes (probably will wait for LazyColumn to support animated insertions).

Seems Compose Multiplatform doesn't have @UiComposable and @ComposableTargetMarker annotations, couldn't import them (but could use them on regular Jetpack Compose).

Thanks again @zach-klippenstein, feel free to suggest more enhancements!