Change our tree to be custom nodes instead of raw widgets (#1017)

We already were faking half of the tree with the synthetic children widgets before. Now we keep children nodes which do the same thing as before but are joined by widget nodes which wrap a widget and provide access to the parent children. This is used to notify the parent children when any child's layout modifiers have changed.

redwood-compose/src/commonMain/kotlin/app/cash/redwood/compose/RedwoodComposition.kt

@@ -28,7 +28,6 @@ import androidx.compose.runtime.currentComposer
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.runtime.snapshots.Snapshot
-import app.cash.redwood.LayoutModifier
 import app.cash.redwood.RedwoodCodegenApi
 import app.cash.redwood.widget.Widget
 import app.cash.redwood.widget.compose.ComposeWidgetChildren
@@ -73,12 +72,12 @@ public fun <W : Any> RedwoodComposition(
   provider: Widget.Provider<W>,
   onEndChanges: () -> Unit = {},
 ): RedwoodComposition {
-  return WidgetRedwoodComposition(scope, WidgetApplier(provider, container, onEndChanges))
+  return WidgetRedwoodComposition(scope, NodeApplier(provider, container, onEndChanges))
 }
 
 private class WidgetRedwoodComposition(
   private val scope: CoroutineScope,
-  applier: WidgetApplier<*>,
+  applier: NodeApplier<*>,
 ) : RedwoodComposition {
   private val recomposer = Recomposer(scope.coroutineContext)
   private val composition = Composition(applier, recomposer)
@@ -128,7 +127,7 @@ public interface RedwoodApplier<W : Any> {
 @RedwoodCodegenApi
 public inline fun <P : Widget.Provider<*>, W : Widget<*>> RedwoodComposeNode(
   crossinline factory: (P) -> W,
-  update: @DisallowComposableCalls Updater<W>.() -> Unit,
+  update: @DisallowComposableCalls Updater<WidgetNode<W, *>>.() -> Unit,
   content: @Composable RedwoodComposeContent<W>.() -> Unit,
 ) {
   // NOTE: You MUST keep the implementation of this function (or more specifically, the interaction
@@ -140,13 +139,13 @@ public inline fun <P : Widget.Provider<*>, W : Widget<*>> RedwoodComposeNode(
     val applier = currentComposer.applier as RedwoodApplier<P>
     currentComposer.createNode {
       @Suppress("UNCHECKED_CAST") // Safe so long as you use generated composition function.
-      factory(applier.provider as P)
+      WidgetNode(factory(applier.provider as P) as Widget<Any>)
     }
   } else {
     currentComposer.useNode()
   }
 
-  Updater<W>(currentComposer).update()
+  Updater<WidgetNode<W, *>>(currentComposer).update()
   RedwoodComposeContent.Instance.content()
 
   currentComposer.endNode()
@@ -162,10 +161,10 @@ public class RedwoodComposeContent<out W : Widget<*>> {
     accessor: (W) -> Widget.Children<*>,
     content: @Composable () -> Unit,
   ) {
-    ComposeNode<ChildrenWidget<*>, Applier<*>>(
+    ComposeNode<ChildrenNode<*>, Applier<*>>(
       factory = {
         @Suppress("UNCHECKED_CAST")
-        ChildrenWidget(accessor as (Widget<Any>) -> Widget.Children<Any>)
+        ChildrenNode(accessor as (Widget<Any>) -> Widget.Children<Any>)
       },
       update = {},
       content = content,
@@ -176,26 +175,3 @@ public class RedwoodComposeContent<out W : Widget<*>> {
     public val Instance: RedwoodComposeContent<Nothing> = RedwoodComposeContent()
   }
 }
-
-/**
- * A synthetic widget which allows the applier to differentiate between multiple groups of children.
- *
- * Compose's tree assumes each node only has single list of children. Or, put another way, even if
- * you apply multiple children Compose treats them as a single list of child nodes. In order to
- * differentiate between these children lists we introduce synthetic nodes. Every real node which
- * supports one or more groups of children will have one or more of these synthetic nodes as its
- * direct descendants. The nodes which are produced by each group of children will then become the
- * descendants of those synthetic nodes.
- */
-internal class ChildrenWidget<W : Any> private constructor(
-  var accessor: ((Widget<W>) -> Widget.Children<W>)?,
-  var children: Widget.Children<W>?,
-) : Widget<W> {
-  constructor(accessor: (Widget<W>) -> Widget.Children<W>) : this(accessor, null)
-  constructor(children: Widget.Children<W>) : this(null, children)
-
-  override val value: Nothing get() = throw AssertionError()
-  override var layoutModifiers: LayoutModifier
-    get() = throw AssertionError()
-    set(_) { throw AssertionError() }
-}

redwood-compose/src/commonMain/kotlin/app/cash/redwood/compose/WidgetApplier.kt

@@ -17,40 +17,45 @@ package app.cash.redwood.compose
 
 import androidx.compose.runtime.AbstractApplier
 import androidx.compose.runtime.Applier
+import app.cash.redwood.LayoutModifier
 import app.cash.redwood.RedwoodCodegenApi
 import app.cash.redwood.widget.Widget
 
 /**
- * An [Applier] for a tree of [Widget]s.
+ * An [Applier] for Redwood's tree of nodes.
  *
- * This applier has special handling for emulating nodes which contain multiple children. Nodes in
- * the tree are required to alternate between [ChildrenWidget] instances and user [Widget] subtypes
- * starting at the root. This invariant is maintained by virtue of the fact that all of the input
- * `@Composables` should be generated Redwood code.
+ * This applier has special handling for emulating widgets which contain multiple children.
+ * Nodes in the tree are required to alternate between [WidgetNode] and [ChildrenNode] starting
+ * from the root. This invariant is maintained by virtue of the fact that all of the input
+ * `@Composable`s should bottom out in generated Redwood code.
  *
  * For example, a widget tree may look like this:
  * ```
- *                    Children(tag=1)
- *                     /          \
- *                    /            \
- *            ToolbarNode        ListNode
- *             ·     ·                 ·
- *            ·       ·                 ·
- * Children(tag=1)  Children(tag=2)   Children(tag=1)
- *        |              |               /       \
- *        |              |              /         \
- *   ButtonNode     ButtonNode     TextNode     TextNode
+ *                        ChildrenNode(root)
+ *                          /           \
+ *                         /             \
+ *          WidgetNode<Toolbar>       WidgetNode<List>
+ *             /          \                      \
+ *            /            \                      \
+ * ChildrenNode(tag=1)  ChildrenNode(tag=2)     ChildrenNode(tag=1)
+ *        |                   |                      /       \
+ *        |                   |                     /         \
+ * WidgetNode<Button>  WidgetNode<Button>  WidgetNode<Text>  WidgetNode<Text>
  * ```
- * The tree produced by this applier is not a real tree. We do not maintain any relationship from
- * user widgets to the synthetic children widgets as they can never be individually moved/removed.
- * The hierarchy is maintained by Compose's slot table and is represented by dotted lines above.
+ * The node tree produced by this applier is not actually a tree. We do not maintain a relationship
+ * from each [WidgetNode] to their [ChildrenNode]s as they can never be individually moved/removed.
+ * Similarly, no relationship is maintained from a [ChildrenNode] to their [WidgetNode]s. Instead,
+ * the [WidgetNode.widget] is what's added to the parent [ChildrenNode.children].
+ *
+ * Compose maintains the tree structure internally. All non-insert operations are performed
+ * using indexes and counts rather than references which are forwarded to [ChildrenNode.children].
  */
 @OptIn(RedwoodCodegenApi::class)
-internal class WidgetApplier<W : Any>(
+internal class NodeApplier<W : Any>(
   override val provider: Widget.Provider<W>,
   root: Widget.Children<W>,
   private val onEndChanges: () -> Unit,
-) : AbstractApplier<Widget<W>>(ChildrenWidget(root)), RedwoodApplier<W> {
+) : AbstractApplier<Node<W>>(ChildrenNode(root)), RedwoodApplier<W> {
   private var closed = false
 
   override fun onEndChanges() {
@@ -59,37 +64,102 @@ internal class WidgetApplier<W : Any>(
     onEndChanges.invoke()
   }
 
-  override fun insertTopDown(index: Int, instance: Widget<W>) {
+  override fun insertTopDown(index: Int, instance: Node<W>) {
     check(!closed)
 
-    if (instance is ChildrenWidget) {
-      instance.children = instance.accessor!!.invoke(current)
-      instance.accessor = null
+    if (instance is ChildrenNode) {
+      @Suppress("UNCHECKED_CAST") // Guaranteed by generated code.
+      val widgetNode = current as WidgetNode<Widget<W>, W>
+      instance.attachTo(widgetNode.widget)
     } else {
-      val current = current as ChildrenWidget
-      current.children!!.insert(index, instance)
+      @Suppress("UNCHECKED_CAST") // Guaranteed by generated code.
+      val widgetNode = instance as WidgetNode<Widget<W>, W>
+      val children = (current as ChildrenNode<W>).children
+
+      widgetNode.container = children
+      children.insert(index, widgetNode.widget)
     }
   }
 
-  override fun insertBottomUp(index: Int, instance: Widget<W>) {
+  override fun insertBottomUp(index: Int, instance: Node<W>) {
     // Ignored, we insert top-down.
   }
 
   override fun remove(index: Int, count: Int) {
     check(!closed)
 
-    val current = current as ChildrenWidget
-    current.children!!.remove(index, count)
+    val current = current as ChildrenNode
+    current.children.remove(index, count)
   }
 
   override fun move(from: Int, to: Int, count: Int) {
     check(!closed)
 
-    val current = current as ChildrenWidget
-    current.children!!.move(from, to, count)
+    val current = current as ChildrenNode
+    current.children.move(from, to, count)
   }
 
   override fun onClear() {
     closed = true
   }
 }
+
+/** @suppress For generated code usage only. */
+@RedwoodCodegenApi
+public sealed interface Node<W : Any>
+
+/**
+ * A node which wraps a [Widget] and also holds onto the [Widget.Children] in which the widget
+ * is placed. The generics of this class are a little funky in order to avoid casts of the widget
+ * in the generated composables.
+ *
+ * @suppress For generated code usage only.
+ */
+@RedwoodCodegenApi
+public class WidgetNode<W : Widget<V>, V : Any>(
+  public val widget: W,
+) : Node<W> {
+  public var container: Widget.Children<V>? = null
+
+  public companion object {
+    public val SetLayoutModifiers: WidgetNode<*, *>.(LayoutModifier) -> Unit = {
+      widget.layoutModifiers = it
+      container?.onLayoutModifierUpdated()
+    }
+  }
+}
+
+/**
+ * A synthetic widget which allows the applier to differentiate between multiple groups of children.
+ *
+ * Compose's tree assumes each node only has single list of children. Or, put another way, even if
+ * you apply multiple children Compose treats them as a single list of child nodes. In order to
+ * differentiate between these children lists we introduce synthetic nodes. Every real node which
+ * supports one or more groups of children will have one or more of these synthetic nodes as its
+ * direct descendants. The nodes which are produced by each group of children will then become the
+ * descendants of those synthetic nodes.
+ *
+ * This node has two valid states:
+ * 1. Non-null accessor and null children. This is the state when created but not inserted in
+ *    the node tree.
+ * 2. Null accessor and non-null children. Once inserted into the tree, the accessor is used to
+ *    fetch the appropriate children reference from the parent widget.
+ *
+ * Transition from 1 to 2 by calling [attachTo]. This may only be done once, and you
+ * cannot transition back to 1 afterwards.
+ */
+@RedwoodCodegenApi
+internal class ChildrenNode<W : Any> private constructor(
+  private var accessor: ((Widget<W>) -> Widget.Children<W>)?,
+  private var _children: Widget.Children<W>?,
+) : Node<W> {
+  constructor(accessor: (Widget<W>) -> Widget.Children<W>) : this(accessor, null)
+  constructor(children: Widget.Children<W>) : this(null, children)
+
+  val children: Widget.Children<W> get() = _children!!
+
+  fun attachTo(parent: Widget<W>) {
+    _children = checkNotNull(accessor).invoke(parent)
+    accessor = null
+  }
+}

redwood-protocol-compose/src/commonMain/kotlin/app/cash/redwood/protocol/compose/ProtocolWidgetChildren.kt

@@ -48,6 +48,5 @@ internal class ProtocolWidgetChildren(
   }
 
   override fun onLayoutModifierUpdated() {
-    throw AssertionError()
   }
 }

redwood-tooling-codegen/src/main/kotlin/app/cash/redwood/tooling/codegen/composeGeneration.kt

@@ -47,12 +47,12 @@ fun Row(
   layoutModifier: LayoutModifier = LayoutModifier,
   children: @Composable @SunspotComposable RowScope.() -> Unit,
 ): Unit {
-  _RedwoodComposeNode<SunspotWidgetFactoryProvider<*>, Row<*>>(
+  RedwoodComposeNode<SunspotWidgetFactoryProvider<*>, Row<*>>(
     factory = { it.RedwoodLayout.Row() },
     update = {
-      set(layoutModifier, Widget.SetLayoutModifiers)
-      set(margin, Row<*>::margin)
-      set(overflow, Row<*>::overflow)
+      set(layoutModifier, WidgetNode.SetLayoutModifiers)
+      set(margin) { widget.margin(it) }
+      set(overflow) { widget.overflow(it) }
     },
     content = {
       into(Row<*>::children) {
@@ -127,15 +127,15 @@ internal fun generateComposable(
           }
 
           val updateLambda = CodeBlock.builder()
-            .add("set(layoutModifier, %T.SetLayoutModifiers)\n", RedwoodWidget.Widget)
+            .add("set(layoutModifier, %T.SetLayoutModifiers)\n", RedwoodCompose.WidgetNode)
 
           val childrenLambda = CodeBlock.builder()
           for (trait in widget.traits) {
             when (trait) {
               is Property,
               is Event,
               -> {
-                updateLambda.add("set(%1N, %2T::%1N)\n", trait.name, widgetType)
+                updateLambda.add("set(%1N) { widget.%1N(it) }\n", trait.name)
               }
               is Children -> {
                 childrenLambda.apply {

redwood-tooling-codegen/src/main/kotlin/app/cash/redwood/tooling/codegen/types.kt

@@ -85,6 +85,7 @@ internal object RedwoodWidget {
 
 internal object RedwoodCompose {
   val RedwoodComposeNode = MemberName("app.cash.redwood.compose", "RedwoodComposeNode")
+  val WidgetNode = ClassName("app.cash.redwood.compose", "WidgetNode")
 }
 
 internal object ComposeRuntime {

redwood-widget/src/commonMain/kotlin/app/cash/redwood/widget/Widget.kt

@@ -16,7 +16,6 @@
 package app.cash.redwood.widget
 
 import app.cash.redwood.LayoutModifier
-import app.cash.redwood.RedwoodCodegenApi
 import kotlin.native.ObjCName
 
 @ObjCName("Widget", exact = true)
@@ -32,12 +31,6 @@ public interface Widget<W : Any> {
    */
   public var layoutModifiers: LayoutModifier
 
-  public companion object {
-    /** @suppress Optimization for generated code to avoid generating/allocating many lambdas. */
-    @RedwoodCodegenApi
-    public val SetLayoutModifiers: Widget<*>.(LayoutModifier) -> Unit = { layoutModifiers = it }
-  }
-
   /** Marker interface for types whose properties expose factories of [Widget]s. */
   @Suppress("unused") // This type parameter used to match against other types like Children.
   public interface Provider<W : Any>