Repository: cashapp/redwood (kotlin)

Issue:
Content.preload()
For LazyColumn it’ll be a performance win to run asynchronous guest code eagerly, before the target `TreehouseView` is attached. One possible way to implement this is a `preload()` function on `Content`:


```
public interface Content {
  /**
   * Prepare the content for an upcoming call to [bind]. This may yield
   * lower latency for lazy layouts.
   * 
   * To avoid resource leaks every call to [preload] must be followed by
   * a call to [unbind], whether or not there is a call to [bind].
   */
  public fun preload(view: TreehouseView)

  ...
}
```

The `preload` function doesn’t _really_ need the `TreehouseView` yet, but accepting it keeps the `hostConfiguration` quite simple. Otherwise we’d need to accept a `HostConfiguration` and fuse it to the `StateFlow<HostConfiguration>` that follows.

Pull request:
Closes #861.

Patch:
diff --git a/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/Content.kt b/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/Content.kt
index 33ca642122..1188cf4c8d 100644
--- a/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/Content.kt
+++ b/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/Content.kt
@@ -16,6 +16,7 @@
 package app.cash.redwood.treehouse
 
 import kotlin.native.ObjCName
+import kotlinx.coroutines.CancellationException
 
 /**
  * A UI built as an interactive widget tree, that may or may not be actively running, or bound to an
@@ -35,6 +36,11 @@ import kotlin.native.ObjCName
  */
 @ObjCName("Content", exact = true)
 public interface Content {
+  /**
+   * Immediately begins preparing the widget tree.
+   */
+  public fun preload(hostConfiguration: HostConfiguration)
+
   /**
    * It is an error to bind multiple views simultaneously.
    *
@@ -42,6 +48,14 @@ public interface Content {
    */
   public fun bind(view: TreehouseView)
 
+  /**
+   * Suspends until content is available; either it is already in the view or it is preloaded and a call to [bind]
+   * will immediately show this content.
+   *
+   * @throws [CancellationException] if it's unbound before it returns.
+   */
+  public suspend fun awaitContent()
+
   /**
    * Calling [unbind] without a bound view is safe.
    *
diff --git a/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/SequentialStateFlow.kt b/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/SequentialStateFlow.kt
new file mode 100644
index 0000000000..76549fbbcf
--- /dev/null
+++ b/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/SequentialStateFlow.kt
@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 2023 Square, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package app.cash.redwood.treehouse
+
+import kotlinx.coroutines.flow.FlowCollector
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.collectLatest
+
+internal class SequentialStateFlow<T>(
+  first: StateFlow<T>,
+) : StateFlow<T> {
+  public val stateFlowValue: MutableStateFlow<StateFlow<T>> = MutableStateFlow(first)
+
+  override val replayCache: List<T> = listOf()
+
+  override suspend fun collect(collector: FlowCollector<T>): Nothing {
+    stateFlowValue.collectLatest { stateFlow ->
+      stateFlow.collect {
+        collector.emit(it)
+      }
+    }
+    error("stateFlow.collect never returns")
+  }
+
+  override val value: T
+    get() = stateFlowValue.value.value
+}
diff --git a/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/TreehouseApp.kt b/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/TreehouseApp.kt
index d7fd0e8071..c0d843c41c 100644
--- a/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/TreehouseApp.kt
+++ b/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/TreehouseApp.kt
@@ -179,7 +179,7 @@ public class TreehouseApp<A : AppService> private constructor(
       next.startFrameClock()
 
       for (content in boundContents) {
-        content.receiveZiplineSession(next, codeChanged = true)
+        content.receiveZiplineSession(next)
       }
 
       if (previous != null) {
@@ -263,7 +263,7 @@ public class TreehouseApp<A : AppService> private constructor(
       get() = EmptySerializersModule()
 
     /**
-     * Returns true to only load code from the network. Otherwise this will recover from
+     * Returns true to only load code from the network. Otherwise, this will recover from
      * unreachable network code by loading code from the cache or the embedded file system.
      *
      * This is false by default. Override it to return true in development, where loading code from
diff --git a/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/TreehouseAppContent.kt b/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/TreehouseAppContent.kt
index 9ac2d9ed8e..d792fb91d7 100644
--- a/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/TreehouseAppContent.kt
+++ b/redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/TreehouseAppContent.kt
@@ -23,12 +23,42 @@ import app.cash.redwood.protocol.widget.ProtocolBridge
 import app.cash.redwood.widget.Widget
 import app.cash.zipline.ZiplineScope
 import app.cash.zipline.withScope
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.cancel
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.job
 import kotlinx.coroutines.launch
 
+private class State<A : AppService>(
+  val viewState: ViewState,
+  val codeState: CodeState<A>,
+)
+
+private sealed interface ViewState {
+  object None : ViewState
+
+  class Preloading(
+    val hostConfiguration: HostConfiguration,
+  ) : ViewState
+
+  class Bound(
+    val view: TreehouseView,
+  ) : ViewState
+}
+
+private sealed interface CodeState<A : AppService> {
+  class Idle<A : AppService> : CodeState<A>
+
+  class Running<A : AppService>(
+    val viewContentCodeBinding: ViewContentCodeBinding<A>,
+    val hasDiffs: Boolean = false,
+  ) : CodeState<A>
+}
+
 internal class TreehouseAppContent<A : AppService>(
   private val treehouseApp: TreehouseApp<A>,
   private val source: TreehouseContentSource<A>,
@@ -36,77 +66,158 @@ internal class TreehouseAppContent<A : AppService>(
 ) : Content {
   private val dispatchers = treehouseApp.dispatchers
 
-  /** How many times this content has received fresh code. */
-  private var codeLoadCount = 0
+  private var stateFlow = MutableStateFlow<State<A>>(State(ViewState.None, CodeState.Idle()))
+
+  override fun preload(hostConfiguration: HostConfiguration) {
+    treehouseApp.dispatchers.checkUi()
+    val previousState = stateFlow.value
+    val previousViewState = previousState.viewState
+
+    check(previousState.viewState is ViewState.None)
+
+    if (previousViewState !is ViewState.Preloading) {
+      val nextViewState = ViewState.Preloading(hostConfiguration)
+
+      // Start the code if necessary.
+      val ziplineSession = treehouseApp.ziplineSession
+      val nextCodeState = when {
+        previousState.codeState is CodeState.Idle && ziplineSession != null -> {
+          CodeState.Running(
+            startViewCodeContentBinding(ziplineSession, MutableStateFlow(nextViewState.hostConfiguration)),
+          )
+        }
+        else -> previousState.codeState
+      }
 
-  private var view: TreehouseView? = null
-  private var viewContentCodeBinding: ViewContentCodeBinding<A>? = null
+      // Ask to get notified when code is ready.
+      treehouseApp.boundContents += this
+
+      stateFlow.value = State(nextViewState, nextCodeState)
+    }
+  }
 
   override fun bind(view: TreehouseView) {
     treehouseApp.dispatchers.checkUi()
+    val previousState = stateFlow.value
+    val previousViewState = previousState.viewState
 
-    // Binding the bound view does nothing. This is necessary so that listeners don't need to
-    // independently track the bound/unbound state.
-    if (this.view == view) return
+    if (previousViewState is ViewState.Bound && previousViewState.view == view) return // Idempotent.
+    check(previousViewState is ViewState.None || previousViewState is ViewState.Preloading)
 
-    require(this.view == null)
-    this.view = view
+    val nextViewState = ViewState.Bound(view)
 
-    treehouseApp.boundContents += this
-    receiveZiplineSession(treehouseApp.ziplineSession, false)
+    // Start the code if necessary.
+    val ziplineSession = treehouseApp.ziplineSession
+    val nextCodeState = when {
+      previousState.codeState is CodeState.Idle && ziplineSession != null -> {
+        CodeState.Running(
+          startViewCodeContentBinding(ziplineSession, nextViewState.view.hostConfiguration),
+        )
+      }
+      else -> previousState.codeState
+    }
+
+    // Ask to get notified when code is ready.
+    if (previousViewState is ViewState.None) {
+      treehouseApp.boundContents += this
+    }
+
+    // Make sure we're showing something in the view; either loaded code or a spinner to show that code is coming.
+    when (nextCodeState) {
+      is CodeState.Idle -> codeListener.onInitialCodeLoading(view)
+      is CodeState.Running -> nextCodeState.viewContentCodeBinding.initView(view)
+    }
+
+    stateFlow.value = State(nextViewState, nextCodeState)
+  }
+
+  override suspend fun awaitContent() {
+    stateFlow.first {
+      if (it.viewState is ViewState.None) {
+        throw CancellationException("unbound while awaiting content")
+      }
+
+      it.codeState is CodeState.Running && it.codeState.hasDiffs
+    }
   }
 
   override fun unbind() {
     treehouseApp.dispatchers.checkUi()
 
-    if (view == null) return // unbind() is idempotent.
+    val previousState = stateFlow.value
+    val previousViewState = previousState.viewState
+
+    if (previousViewState is ViewState.None) return // Idempotent.
+
+    val nextViewState = ViewState.None
+    val nextCodeState = CodeState.Idle<A>()
 
-    treehouseApp.boundContents.remove(this)
-    viewContentCodeBinding?.cancel()
-    view = null
-    viewContentCodeBinding = null
+    // Cancel the code if necessary.
+    if (previousState.codeState is CodeState.Running) {
+      treehouseApp.boundContents.remove(this)
+      previousState.codeState.viewContentCodeBinding.cancel()
+    }
+
+    stateFlow.value = State(nextViewState, nextCodeState)
+  }
+
+  internal fun receiveZiplineSession(next: ZiplineSession<A>) {
+    treehouseApp.dispatchers.checkUi()
+
+    val previousState = stateFlow.value
+    val previousViewState = previousState.viewState
+    val previousCodeState = previousState.codeState
+
+    val hostConfiguration = when (previousViewState) {
+      is ViewState.Preloading -> MutableStateFlow(previousViewState.hostConfiguration)
+      is ViewState.Bound -> previousViewState.view.hostConfiguration
+      else -> error("unexpected receiveZiplineSession with no view bound and no preload")
+    }
+
+    val nextViewState = ViewState.None
+    val nextCodeState = CodeState.Running(
+      startViewCodeContentBinding(next, hostConfiguration),
+    )
+
+    if (previousViewState is ViewState.Bound) {
+      nextCodeState.viewContentCodeBinding.initView(previousViewState.view)
+    }
+
+    // Cancel the old code if necessary.
+    if (previousCodeState is CodeState.Running) {
+      previousCodeState.viewContentCodeBinding.cancel()
+    }
+
+    stateFlow.value = State(nextViewState, nextCodeState)
   }
 
   /** This function may only be invoked on [TreehouseDispatchers.ui]. */
-  internal fun receiveZiplineSession(
-    ziplineSession: ZiplineSession<A>?,
-    codeChanged: Boolean,
-  ) {
+  private fun startViewCodeContentBinding(
+    ziplineSession: ZiplineSession<A>,
+    firstHostConfiguration: StateFlow<HostConfiguration>,
+  ): ViewContentCodeBinding<A> {
     dispatchers.checkUi()
 
-    val previous = viewContentCodeBinding
-    if (!codeChanged && previous != null) return // No change.
-
-    val view = this.view!!
-    if (ziplineSession != null) {
-      // We have code. Launch the treehouse UI.
-      viewContentCodeBinding = ViewContentCodeBinding(
-        app = treehouseApp,
-        appScope = treehouseApp.appScope,
-        eventPublisher = treehouseApp.eventPublisher,
-        contentSource = source,
-        codeListener = codeListener,
-        session = ziplineSession,
-        view = view,
-      ).apply {
-        start(ziplineSession, view)
-      }
-    } else {
-      // We don't have code yet. Let the CodeListener show a loading spinner or something.
-      if (codeLoadCount == 0) {
-        codeListener.onInitialCodeLoading(view)
-      }
-      viewContentCodeBinding = null
+    return ViewContentCodeBinding(
+      app = treehouseApp,
+      appScope = treehouseApp.appScope,
+      eventPublisher = treehouseApp.eventPublisher,
+      contentSource = source,
+      codeListener = codeListener,
+      stateFlow = stateFlow,
+      session = ziplineSession,
+      firstHostConfiguration = firstHostConfiguration,
+    ).apply {
+      start(ziplineSession)
     }
-
-    codeLoadCount++
-    previous?.cancel()
   }
 }
 
 /**
  * Connects a [TreehouseView], a [TreehouseContentSource], and a [ZiplineSession].
  *
+ * The TreehouseView may not be known immediately, as in [Content.preload].
+ *
  * Canceled by [TreehouseAppContent] if the view is unbound from its content, or if the code is
  * updated.
  *
@@ -120,26 +231,23 @@ private class ViewContentCodeBinding<A : AppService>(
   val eventPublisher: EventPublisher,
   val contentSource: TreehouseContentSource<A>,
   val codeListener: CodeListener,
+  val stateFlow: MutableStateFlow<State<A>>,
   session: ZiplineSession<A>,
-  view: TreehouseView,
+  firstHostConfiguration: StateFlow<HostConfiguration>,
 ) : EventSink, DiffSinkService {
+  private val hostConfigurationFlow = SequentialStateFlow(firstHostConfiguration)
+
   private val isInitialLaunch: Boolean = session.isInitialLaunch
 
+  private val json = session.zipline.json
+
   private val bindingScope = CoroutineScope(SupervisorJob(appScope.coroutineContext.job))
 
-  /** Only accessed on [TreehouseDispatchers.ui]. Null after [cancel]. */
-  private var viewOrNull: TreehouseView? = view
+  /** Only accessed on [TreehouseDispatchers.ui]. Null before [initView] and after [cancel]. */
+  private var viewOrNull: TreehouseView? = null
 
-  /** Only accessed on [TreehouseDispatchers.ui]. Null after [cancel]. */
-  @Suppress("UNCHECKED_CAST") // We don't have a type parameter for the widget type.
-  private var bridgeOrNull: ProtocolBridge<*>? = ProtocolBridge(
-    container = view.children as Widget.Children<Any>,
-    factory = view.widgetSystem.widgetFactory(
-      json = session.zipline.json,
-      protocolMismatchHandler = eventPublisher.protocolMismatchHandler(app),
-    ) as DiffConsumingNode.Factory<Any>,
-    eventSink = this,
-  )
+  /** Only accessed on [TreehouseDispatchers.ui]. Null before [initView] and after [cancel]. */
+  private var bridgeOrNull: ProtocolBridge<*>? = null
 
   /** Only accessed on [TreehouseDispatchers.zipline]. */
   private val ziplineScope = ZiplineScope()
@@ -147,8 +255,42 @@ private class ViewContentCodeBinding<A : AppService>(
   /** Only accessed on [TreehouseDispatchers.zipline]. Null after [cancel]. */
   private var treehouseUiOrNull: ZiplineTreehouseUi? = null
 
+  /** Only accessed on [TreehouseDispatchers.ui]. Empty after [initView]. */
+  private val diffsAwaitingInitView = ArrayDeque<Diff>()
+
+  /** Diffs applied to the UI. Only accessed on [TreehouseDispatchers.ui]. */
+  var diffCount = 0
+
   /** Only accessed on [TreehouseDispatchers.ui]. */
-  private var diffCount = 0
+  private var canceled = false
+
+  private var initViewCalled: Boolean = false
+
+  fun initView(view: TreehouseView) {
+    app.dispatchers.checkUi()
+
+    require(!initViewCalled)
+    initViewCalled = true
+
+    if (canceled) return
+
+    viewOrNull = view
+    @Suppress("UNCHECKED_CAST") // We don't have a type parameter for the widget type.
+    bridgeOrNull = ProtocolBridge(
+      container = view.children as Widget.Children<Any>,
+      factory = view.widgetSystem.widgetFactory(
+        json = json,
+        protocolMismatchHandler = eventPublisher.protocolMismatchHandler(app),
+      ) as DiffConsumingNode.Factory<Any>,
+      eventSink = this,
+    )
+
+    // Apply all the diffs received before we had a view to apply them to.
+    while (true) {
+      val diff = diffsAwaitingInitView.removeFirstOrNull() ?: break
+      receiveDiffOnUiDispatcher(diff)
+    }
+  }
 
   /** Send an event from the UI to Zipline. */
   override fun sendEvent(event: Event) {
@@ -163,31 +305,54 @@ private class ViewContentCodeBinding<A : AppService>(
   override fun sendDiff(diff: Diff) {
     // Receive UI updates on the UI dispatcher.
     bindingScope.launch(app.dispatchers.ui) {
-      val view = viewOrNull ?: return@launch
-      val bridge = bridgeOrNull ?: return@launch
+      receiveDiffOnUiDispatcher(diff)
+    }
+  }
+
+  private fun receiveDiffOnUiDispatcher(diff: Diff) {
+    val view = viewOrNull
+    val bridge = bridgeOrNull
+
+    if (canceled) {
+      return
+    }
 
-      if (diffCount++ == 0) {
-        view.reset()
-        codeListener.onCodeLoaded(view, isInitialLaunch)
+    if (view == null || bridge == null) {
+      if (diffsAwaitingInitView.isEmpty()) {
+        // Update the TreehouseAppContent's ViewState on the first diff.
+        val currentState = stateFlow.value
+        if (currentState.codeState is CodeState.Running && currentState.codeState.viewContentCodeBinding == this) {
+          stateFlow.value = State(currentState.viewState, CodeState.Running(this, hasDiffs = true))
+        }
       }
 
-      bridge.sendDiff(diff)
+      diffsAwaitingInitView += diff
+      return
+    }
+
+    if (diffCount++ == 0) {
+      view.reset()
+      codeListener.onCodeLoaded(view, isInitialLaunch)
     }
+
+    bridge.sendDiff(diff)
   }
 
-  fun start(session: ZiplineSession<A>, view: TreehouseView) {
+  fun start(session: ZiplineSession<A>) {
     bindingScope.launch(app.dispatchers.zipline) {
       val scopedAppService = session.appService.withScope(ziplineScope)
       val treehouseUi = contentSource.get(scopedAppService)
       treehouseUiOrNull = treehouseUi
       treehouseUi.start(
         diffSink = this@ViewContentCodeBinding,
-        hostConfigurations = view.hostConfiguration.toFlowWithInitialValue(),
+        hostConfigurations = hostConfigurationFlow.toFlowWithInitialValue(),
       )
     }
   }
 
   fun cancel() {
+    app.dispatchers.checkUi()
+    canceled = true
     viewOrNull = null
     bridgeOrNull = null
     appScope.launch(app.dispatchers.zipline) {


Test patch:
diff --git a/redwood-treehouse-host/src/commonTest/kotlin/app/cash/redwood/treehouse/SequentialStateFlowTest.kt b/redwood-treehouse-host/src/commonTest/kotlin/app/cash/redwood/treehouse/SequentialStateFlowTest.kt
new file mode 100644
index 0000000000..7dae76b7c3
--- /dev/null
+++ b/redwood-treehouse-host/src/commonTest/kotlin/app/cash/redwood/treehouse/SequentialStateFlowTest.kt
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2023 Square, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package app.cash.redwood.treehouse
+
+import app.cash.turbine.test
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.runTest
+
+@OptIn(ExperimentalCoroutinesApi::class)
+class SequentialStateFlowTest {
+
+  @Test
+  fun happyPath() = runTest {
+    val ab = MutableStateFlow("a")
+    val sequentialStateFlow = SequentialStateFlow(ab)
+    assertEquals("a", sequentialStateFlow.value)
+
+    sequentialStateFlow.test {
+      assertEquals("a", awaitItem())
+      assertEquals("a", sequentialStateFlow.value)
+
+      ab.value = "b"
+      assertEquals("b", sequentialStateFlow.value)
+      assertEquals("b", awaitItem())
+
+      val cd = MutableStateFlow("c")
+      sequentialStateFlow.stateFlowValue.value = cd
+      assertEquals("c", sequentialStateFlow.value)
+      assertEquals("c", awaitItem())
+
+      cd.value = "d"
+      assertEquals("d", sequentialStateFlow.value)
+      assertEquals("d", awaitItem())
+    }
+  }
+
+  @Test
+  fun valuesNotEmittedAfterStateFlowChanges() = runTest {
+    val ab = MutableStateFlow("a")
+    val sequentialStateFlow = SequentialStateFlow(ab)
+    sequentialStateFlow.test {
+      assertEquals("a", awaitItem())
+
+      val cd = MutableStateFlow("c")
+      sequentialStateFlow.stateFlowValue.value = cd
+      assertEquals("c", awaitItem())
+
+      ab.value = "b" // The flow should not emit this value!
+      cd.value = "d"
+      assertEquals("d", awaitItem())
+    }
+  }
+
+  @Test
+  fun stateFlowChanges() = runTest {
+    val ab = MutableStateFlow("a")
+    val sequentialStateFlow = SequentialStateFlow(ab)
+    sequentialStateFlow.test {
+      assertEquals("a", awaitItem())
+
+      val cd = MutableStateFlow("c")
+      sequentialStateFlow.stateFlowValue.value = cd
+
+      val ef = MutableStateFlow("e")
+      sequentialStateFlow.stateFlowValue.value = ef
+
+      // Because Turbine receives without suspending, emit("c") completes before that collect is canceled.
+      assertEquals("c", awaitItem())
+      assertEquals("e", awaitItem())
+    }
+  }
+}