Reduce the number of retain cycles between host and guest

CHANGELOG.md

@@ -12,6 +12,7 @@ Changed:
 - APIs accepting a `FileSystem` and `Path` now have the `FileSystem` coming before the `Path` in the parameter list. Compatibility functions are retained for this version, but will be removed in the next version.
 
 Fixed:
+- Fix memory leaks caused by reference cycles on iOS. We got into trouble mixing garbage-collected Kotlin objects with reference-counted Swift objects.
 - Work around a problem with our memory-leak fix where our old LazyList code would crash when its placeholders were unexpectedly removed.
 - Avoid calling into the internal Zipline instance from the UI thread on startup. This would manifest as weird native crashes due to multiple threads mutating shared memory.
 - In `UIViewLazyList`, fix `UInt` to `UIColor` conversion math used for  `pullRefreshContentColor`.

redwood-protocol-host/src/commonMain/kotlin/app/cash/redwood/protocol/host/ProtocolBridge.kt

@@ -63,7 +63,11 @@ public class ProtocolBridge<W : Any>(
   /** Nodes available for reuse. */
   private val pool = ArrayDeque<ProtocolNode<W>>()
 
+  private var closed = false
+
   override fun sendChanges(changes: List<Change>) {
+    check(!closed)
+
     @Suppress("NAME_SHADOWING")
     val changes = applyReuse(changes)
 
@@ -154,6 +158,16 @@ public class ProtocolBridge<W : Any>(
     return checkNotNull(nodes[id]) { "Unknown widget ID ${id.value}" }
   }
 
+  /**
+   * Proactively clear held widgets. (This avoids problems when mixing garbage-collected Kotlin
+   * objects with reference-counted Swift objects.)
+   */
+  public fun close() {
+    closed = true
+    nodes.clear()
+    pool.clear()
+  }
+
   private fun poolIfReusable(removedNode: ProtocolNode<W>) {
     if (removedNode.reuse) {
       removedNode.shapeHash = shapeHash(this.factory, removedNode)

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

@@ -263,7 +263,7 @@ internal class ProtocolButton<W : Any>(
       2 -> _widget.enabled(json.decodeFromJsonElement(serializer_1, change.value))
       3 -> {
         val onClick: (() -> Unit)? = if (change.value.jsonPrimitive.boolean) {
-          { eventSink.sendEvent(Event(change.id, EventTag(3))) }
+          OnClick(json, change.id, eventSink) }
         } else {
           null
         }
@@ -338,6 +338,12 @@ internal fun generateProtocolNode(
           var nextSerializerId = 0
           val serializerIds = mutableMapOf<TypeName, Int>()
 
+          for (trait in properties) {
+            if (trait is ProtocolEvent) {
+              addType(generateEventHandler(trait))
+            }
+          }
+
           addFunction(
             FunSpec.builder("apply")
               .addModifiers(OVERRIDE)
@@ -370,31 +376,28 @@ internal fun generateProtocolNode(
                         KotlinxSerialization.jsonPrimitive,
                         KotlinxSerialization.jsonBoolean,
                       )
+
                       val arguments = mutableListOf<CodeBlock>()
-                      for ((index, parameterFqType) in trait.parameterTypes.withIndex()) {
+                      for (parameterFqType in trait.parameterTypes) {
                         val parameterType = parameterFqType.asTypeName()
                         val serializerId = serializerIds.computeIfAbsent(parameterType) {
                           nextSerializerId++
                         }
-                        arguments += CodeBlock.of(
-                          "json.encodeToJsonElement(serializer_%L, arg%L)",
-                          serializerId,
-                          index,
-                        )
+                        arguments += CodeBlock.of("serializer_%L", serializerId)
                       }
+
                       if (trait.parameterTypes.isEmpty()) {
-                        beginControlFlow("{")
+                        addStatement(
+                          "%L(json, change.id, eventSink)",
+                          trait.eventHandlerName,
+                        )
                       } else {
-                        beginControlFlow("{ %L ->", trait.parameterTypes.indices.map { CodeBlock.of("arg$it") }.joinToCode())
+                        addStatement(
+                          "%L(json, change.id, eventSink, %L)",
+                          trait.eventHandlerName,
+                          arguments.joinToCode(),
+                        )
                       }
-                      addStatement(
-                        "eventSink.sendEvent(%T(change.id, %T(%L), listOf(%L)))",
-                        Protocol.Event,
-                        Protocol.EventTag,
-                        trait.tag,
-                        arguments.joinToCode(),
-                      )
-                      endControlFlow()
 
                       nextControlFlow("else")
                       if (trait.isNullable) {
@@ -491,6 +494,115 @@ internal fun generateProtocolNode(
     .build()
 }
 
+/** Returns a class name like "OnClick". */
+private val ProtocolEvent.eventHandlerName: String
+  get() = name.replaceFirstChar { it.uppercase() }
+
+/**
+ * Generates a named event handler class. We do this instead of using a lambda to be explicit in
+ * which variables are captured by the event handler. (This avoids problems when mixing
+ * garbage-collected Kotlin objects with reference-counted Swift objects.)
+ */
+/*
+private class OnClick(
+  private val json: Json,
+  private val id: Id,
+  private val eventSink: EventSink,
+  private val serializer_0: KSerializer<Int>,
+  private val serializer_1: KSerializer<String>,
+) : (Int, String) -> Unit {
+  override fun invoke(arg0: Int, arg1: String) {
+    eventSink.sendEvent(
+      Event(
+        id,
+        EventTag(3),
+        listOf(
+          json.encodeToJsonElement(serializer_0, arg0),
+          json.encodeToJsonElement(serializer_1, arg1),
+        )
+      )
+    )
+  }
+}
+*/
+private fun generateEventHandler(
+  trait: ProtocolEvent,
+): TypeSpec {
+  val constructor = FunSpec.constructorBuilder()
+  val invoke = FunSpec.builder("invoke")
+    .addModifiers(OVERRIDE)
+
+  val classBuilder = TypeSpec.classBuilder(trait.eventHandlerName)
+    .addModifiers(PRIVATE)
+    .addSuperinterface(trait.lambdaType.copy(nullable = false))
+
+  addConstructorParameterAndProperty(classBuilder, constructor, "json", KotlinxSerialization.Json)
+  addConstructorParameterAndProperty(classBuilder, constructor, "id", Protocol.Id)
+  addConstructorParameterAndProperty(classBuilder, constructor, "eventSink", Protocol.EventSink)
+
+  val arguments = mutableListOf<CodeBlock>()
+  for ((index, parameterFqType) in trait.parameterTypes.withIndex()) {
+    val parameterType = parameterFqType.asTypeName()
+    val serializerType = KotlinxSerialization.KSerializer.parameterizedBy(parameterType)
+    val serializerId = "serializer_$index"
+    val parameterName = "arg$index"
+
+    addConstructorParameterAndProperty(classBuilder, constructor, serializerId, serializerType)
+    invoke.addParameter(ParameterSpec(parameterName, parameterType))
+
+    arguments += CodeBlock.of(
+      "json.encodeToJsonElement(%L, %L)",
+      serializerId,
+      parameterName,
+    )
+  }
+
+  if (arguments.isEmpty()) {
+    invoke.addCode(
+      "eventSink.sendEvent(%T(id, %T(%L)))",
+      Protocol.Event,
+      Protocol.EventTag,
+      trait.tag,
+    )
+  } else {
+    invoke.addCode(
+      "eventSink.sendEvent(⇥\n%T(⇥\nid,\n%T(%L),\nlistOf(⇥\n%L,\n⇤),\n⇤),\n⇤)",
+      Protocol.Event,
+      Protocol.EventTag,
+      trait.tag,
+      arguments.joinToCode(separator = ",\n"),
+    )
+  }
+
+  classBuilder.primaryConstructor(constructor.build())
+  classBuilder.addFunction(invoke.build())
+  return classBuilder.build()
+}
+
+/** Adds a constructor parameter and property with the same name. */
+private fun addConstructorParameterAndProperty(
+  classBuilder: TypeSpec.Builder,
+  constructorBuilder: FunSpec.Builder,
+  name: String,
+  type: TypeName,
+) {
+  constructorBuilder.addParameter(
+    ParameterSpec(
+      name = name,
+      type = type,
+    ),
+  )
+
+  classBuilder.addProperty(
+    PropertySpec.builder(
+      name = name,
+      type = type,
+      modifiers = listOf(PRIVATE),
+    ).initializer(name)
+      .build(),
+  )
+}
+
 internal fun generateProtocolModifierImpls(
   schema: ProtocolSchema,
   host: ProtocolSchema = schema,

redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/TreehouseAppContent.kt

@@ -27,7 +27,10 @@ import app.cash.redwood.widget.Widget
 import app.cash.zipline.ZiplineApiMismatchException
 import app.cash.zipline.ZiplineScope
 import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.CoroutineStart
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -302,7 +305,7 @@ private class ViewContentCodeBinding<A : AppService>(
   private val isInitialLaunch: Boolean,
   private val onBackPressedDispatcher: OnBackPressedDispatcher,
   firstUiConfiguration: StateFlow<UiConfiguration>,
-) : EventSink, ChangesSinkService, TreehouseView.SaveCallback, ZiplineTreehouseUi.Host {
+) : ChangesSinkService, TreehouseView.SaveCallback, ZiplineTreehouseUi.Host {
   private val uiConfigurationFlow = SequentialStateFlow(firstUiConfiguration)
 
   private val bindingScope = CoroutineScope(
@@ -324,6 +327,9 @@ private class ViewContentCodeBinding<A : AppService>(
   /** Only accessed on [TreehouseDispatchers.zipline]. Null after [cancel]. */
   private var treehouseUiOrNull: ZiplineTreehouseUi? = null
 
+  /** Note that this is necessary to break the retain cycle between host and guest. */
+  private val eventBridge = EventBridge(dispatchers.zipline, bindingScope)
+
   /** Only accessed on [TreehouseDispatchers.ui]. Empty after [initView]. */
   private val changesAwaitingInitView = ArrayDeque<List<Change>>()
 
@@ -360,15 +366,6 @@ private class ViewContentCodeBinding<A : AppService>(
     }
   }
 
-  /** Send an event from the UI to Zipline. */
-  override fun sendEvent(event: Event) {
-    // Send UI events on the zipline dispatcher.
-    bindingScope.launch(dispatchers.zipline) {
-      val treehouseUi = treehouseUiOrNull ?: return@launch
-      treehouseUi.sendEvent(event)
-    }
-  }
-
   /** Send changes from Zipline to the UI. */
   override fun sendChanges(changes: List<Change>) {
     // Receive UI updates on the UI dispatcher.
@@ -412,7 +409,7 @@ private class ViewContentCodeBinding<A : AppService>(
           json = codeSession.json,
           protocolMismatchHandler = eventPublisher.widgetProtocolMismatchHandler,
         ) as ProtocolFactory<Any>,
-        eventSink = this,
+        eventSink = eventBridge,
       )
       bridgeOrNull = bridge
     }
@@ -430,6 +427,7 @@ private class ViewContentCodeBinding<A : AppService>(
       val scopedAppService = serviceScope.apply(codeSession.appService)
       val treehouseUi = contentSource.get(scopedAppService)
       treehouseUiOrNull = treehouseUi
+      eventBridge.delegate = treehouseUi
       stateSnapshot = viewOrNull?.stateSnapshotId?.let {
         stateStore.get(it.value.orEmpty())
       }
@@ -488,16 +486,47 @@ private class ViewContentCodeBinding<A : AppService>(
     }
   }
 
+  @OptIn(ExperimentalCoroutinesApi::class)
   fun cancel() {
     dispatchers.checkUi()
     canceled = true
     viewOrNull?.saveCallback = null
     viewOrNull = null
+    bridgeOrNull?.close()
     bridgeOrNull = null
-    bindingScope.launch(dispatchers.zipline) {
+    eventBridge.bindingScope = null
+    eventBridge.ziplineDispatcher = null
+    bindingScope.launch(dispatchers.zipline, start = CoroutineStart.ATOMIC) {
       treehouseUiOrNull = null
+      eventBridge.delegate = null
       serviceScope.close()
       bindingScope.cancel()
     }
   }
 }
+
+/**
+ * Bridge events from the UI dispatcher to the Zipline dispatcher.
+ *
+ * Event sinks are in a natural retain cycle between the host and guest. We prevent unwanted
+ * retain cycles by breaking the link to the delegate when the binding is canceled. This avoids
+ * problems when mixing garbage-collected Kotlin objects with reference-counted Swift objects.
+ */
+private class EventBridge(
+  // Both properties are only accessed on the UI dispatcher and null after cancel().
+  var ziplineDispatcher: CoroutineDispatcher?,
+  var bindingScope: CoroutineScope?,
+) : EventSink {
+  // Only accessed on the Zipline dispatcher and null after cancel().
+  var delegate: EventSink? = null
+
+  /** Send an event from the UI to Zipline. */
+  override fun sendEvent(event: Event) {
+    // Send UI events on the zipline dispatcher.
+    val dispatcher = this.ziplineDispatcher ?: return
+    val bindingScope = this.bindingScope ?: return
+    bindingScope.launch(dispatcher) {
+      delegate?.sendEvent(event)
+    }
+  }
+}