Change runMosaic to support effects

mosaic-runtime/src/commonMain/kotlin/com/jakewharton/mosaic/mosaic.kt

@@ -11,15 +11,16 @@ import androidx.compose.runtime.snapshots.Snapshot
 import com.github.ajalt.mordant.terminal.Terminal as MordantTerminal
 import com.jakewharton.mosaic.layout.MosaicNode
 import com.jakewharton.mosaic.ui.BoxMeasurePolicy
+import kotlin.coroutines.CoroutineContext
 import kotlin.time.ExperimentalTime
-import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.CoroutineStart.UNDISPATCHED
+import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.Job
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.consumeEach
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.yield
 
 /**
  * True for a debug-like output that renders each "frame" on its own with a timestamp delta.
@@ -50,59 +51,40 @@ public fun renderMosaic(content: @Composable () -> Unit): String {
 	return render.toString()
 }
 
-public interface MosaicScope : CoroutineScope {
-	public fun setContent(content: @Composable () -> Unit)
-}
-
-public suspend fun runMosaic(body: suspend MosaicScope.() -> Unit): Unit = coroutineScope {
+public suspend fun runMosaic(content: @Composable () -> Unit): Unit = coroutineScope {
 	val rendering = if (debugOutput) {
 		@OptIn(ExperimentalTime::class) // Not used in production.
 		DebugRendering()
 	} else {
 		AnsiRendering()
 	}
 
-	var hasFrameWaiters = false
-	val clock = BroadcastFrameClock {
-		hasFrameWaiters = true
-	}
-
+	val clock = BroadcastFrameClock()
 	val job = Job(coroutineContext[Job])
 	val composeContext = coroutineContext + clock + job
 
-	val rootNode = createRootNode()
-	var displaySignal: CompletableDeferred<Unit>? = null
-	val applier = MosaicNodeApplier(rootNode) {
-		val render = rendering.render(rootNode)
-		platformDisplay(render)
+	GlobalSnapshotManager.ensureStarted(composeContext)
 
-		displaySignal?.complete(Unit)
-		hasFrameWaiters = false
+	launch(composeContext) {
+		while (true) {
+			clock.sendFrame(0L) // Frame time value is not used by Compose runtime.
+			delay(50L)
+		}
 	}
 
 	val recomposer = Recomposer(composeContext)
-	val composition = Composition(applier, recomposer)
-
-	// Start undispatched to ensure we can use suspending things inside the content.
-	launch(start = UNDISPATCHED, context = composeContext) {
+	launch(composeContext, start = CoroutineStart.UNDISPATCHED) {
 		recomposer.runRecomposeAndApplyChanges()
 	}
 
-	launch(context = composeContext) {
-		while (true) {
-			clock.sendFrame(0L) // Frame time value is not used by Compose runtime.
-			delay(50)
-		}
-	}
-
 	val terminal = MordantTerminal()
 	val terminalInfo = mutableStateOf(
 		Terminal(
 			size = Terminal.Size(terminal.info.width, terminal.info.height),
 		),
 	)
 
-	launch(context = composeContext) {
+	launch(composeContext) {
 		while (true) {
 			val currentTerminalInfo = terminalInfo.value
 			if (terminal.info.updateTerminalSize() &&
@@ -115,59 +97,32 @@ public suspend fun runMosaic(body: suspend MosaicScope.() -> Unit): Unit = corou
 					size = Terminal.Size(terminal.info.width, terminal.info.height),
 				)
 			}
-			delay(50)
+			delay(50L)
 		}
 	}
 
-	coroutineScope {
-		val scope = object : MosaicScope, CoroutineScope by this {
-			override fun setContent(content: @Composable () -> Unit) {
-				composition.setContent {
-					CompositionLocalProvider(LocalTerminal provides terminalInfo.value) {
-						content()
-					}
-				}
-			}
-		}
+	val rootNode = createRootNode()
+	val applier = MosaicNodeApplier(rootNode) {
+		val render = rendering.render(rootNode)
+		platformDisplay(render)
+	}
 
-		var snapshotNotificationsPending = false
-		val observer: (state: Any) -> Unit = {
-			if (!snapshotNotificationsPending) {
-				snapshotNotificationsPending = true
-				launch {
-					snapshotNotificationsPending = false
-					Snapshot.sendApplyNotifications()
-				}
-			}
-		}
-		val snapshotObserverHandle = Snapshot.registerGlobalWriteObserver(observer)
-		try {
-			scope.body()
-		} finally {
-			snapshotObserverHandle.dispose()
+	Composition(applier, recomposer).setContent {
+		CompositionLocalProvider(LocalTerminal provides terminalInfo.value) {
+			content()
 		}
 	}
 
-	// Ensure the final state modification is discovered. We need to ensure that the coroutine
-	// which is running the recomposition loop wakes up, notices the changes, and waits for the
-	// next frame. If you are using snapshots this only requires a single yield. If you are not
-	// then it requires two yields. THIS IS NOT GREAT! But at least it's implementation detail...
-	// TODO https://issuetracker.google.com/issues/169425431
-	yield()
-	yield()
-	Snapshot.sendApplyNotifications()
-	yield()
-	yield()
-
-	if (hasFrameWaiters) {
-		CompletableDeferred<Unit>().also {
-			displaySignal = it
-			it.await()
-		}
+	val effectJob = checkNotNull(recomposer.effectCoroutineContext[Job]) {
+		"No Job in effectCoroutineContext of recomposer"
 	}
+	effectJob.children.forEach { it.join() }
+	recomposer.awaitIdle()
+
+	recomposer.close()
+	recomposer.join()
 
 	job.cancel()
-	composition.dispose()
 }
 
 internal fun createRootNode(): MosaicNode {
@@ -204,3 +159,25 @@ internal class MosaicNodeApplier(
 
 	override fun onClear() {}
 }
+
+internal object GlobalSnapshotManager {
+	private val started = AtomicBoolean(false)
+	private val sent = AtomicBoolean(false)
+
+	fun ensureStarted(coroutineContext: CoroutineContext) {
+		if (started.compareAndSet(expect = false, update = true)) {
+			val channel = Channel<Unit>(1)
+			CoroutineScope(coroutineContext).launch {
+				channel.consumeEach {
+					sent.set(false)
+					Snapshot.sendApplyNotifications()
+				}
+			}
+			Snapshot.registerGlobalWriteObserver {
+				if (sent.compareAndSet(expect = false, update = true)) {
+					channel.trySend(Unit)
+				}
+			}
+		}
+	}
+}

mosaic-runtime/src/commonMain/kotlin/com/jakewharton/mosaic/platform.kt

@@ -1,3 +1,10 @@
 package com.jakewharton.mosaic
 
 internal expect fun platformDisplay(chars: CharSequence)
+
+internal expect class AtomicBoolean(initialValue: Boolean) {
+
+	fun set(value: Boolean)
+
+	fun compareAndSet(expect: Boolean, update: Boolean): Boolean
+}

mosaic-runtime/src/concurrentMain/kotlin/com/jakewharton/mosaic/blocking.kt

@@ -1,9 +1,10 @@
 package com.jakewharton.mosaic
 
+import androidx.compose.runtime.Composable
 import kotlinx.coroutines.runBlocking
 
-public fun runMosaicBlocking(body: suspend MosaicScope.() -> Unit) {
+public fun runMosaicBlocking(content: @Composable () -> Unit) {
 	runBlocking {
-		runMosaic(body)
+		runMosaic(content)
 	}
 }

mosaic-runtime/src/jvmMain/kotlin/com/jakewharton/mosaic/platform.kt

@@ -2,6 +2,7 @@ package com.jakewharton.mosaic
 
 import java.nio.CharBuffer
 import java.nio.charset.StandardCharsets.UTF_8
+import java.util.concurrent.atomic.AtomicBoolean as JavaAtomicBoolean
 import org.fusesource.jansi.AnsiConsole
 
 private val out = AnsiConsole.out()!!.also { AnsiConsole.systemInstall() }
@@ -20,3 +21,15 @@ internal actual fun platformDisplay(chars: CharSequence) {
 	// buffered and not processed until the next frame, or not at all on the final frame.
 	out.flush()
 }
+
+internal actual class AtomicBoolean actual constructor(initialValue: Boolean) {
+	private val delegate = JavaAtomicBoolean(initialValue)
+
+	actual fun set(value: Boolean) {
+		delegate.set(value)
+	}
+
+	actual fun compareAndSet(expect: Boolean, update: Boolean): Boolean {
+		return delegate.compareAndSet(expect, update)
+	}
+}

mosaic-runtime/src/nonJvmMain/kotlin/com/jakewharton/mosaic/platform.kt

@@ -3,3 +3,17 @@ package com.jakewharton.mosaic
 internal actual fun platformDisplay(chars: CharSequence) {
 	print(chars.toString())
 }
+
+internal actual class AtomicBoolean actual constructor(initialValue: Boolean) {
+	private var value: Boolean = initialValue
+
+	actual fun set(value: Boolean) {
+		this.value = value
+	}
+
+	actual fun compareAndSet(expect: Boolean, update: Boolean): Boolean {
+		if (value != expect) return false
+		value = update
+		return true
+	}
+}

samples/counter/src/commonMain/kotlin/example/counter.kt

@@ -1,22 +1,24 @@
 package example
 
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
 import androidx.compose.runtime.setValue
-import com.jakewharton.mosaic.MosaicScope
 import com.jakewharton.mosaic.ui.Text
 import kotlinx.coroutines.delay
 
-suspend fun MosaicScope.runCounter() {
-	// TODO https://github.com/JakeWharton/mosaic/issues/3
-	var count by mutableStateOf(0)
+@Composable
+fun Counter() {
+	var count by remember { mutableStateOf(0) }
 
-	setContent {
-		Text("The count is: $count")
-	}
+	Text("The count is: $count")
 
-	for (i in 1..20) {
-		delay(250)
-		count = i
+	LaunchedEffect(Unit) {
+		for (i in 1..20) {
+			delay(250)
+			count = i
+		}
 	}
 }

samples/counter/src/jsMain/kotlin/example/main.kt

@@ -3,5 +3,5 @@ package example
 import com.jakewharton.mosaic.runMosaic
 
 suspend fun main() = runMosaic {
-	runCounter()
+	Counter()
 }

samples/counter/src/jvmMain/kotlin/example/main.kt

@@ -5,5 +5,5 @@ package example
 import com.jakewharton.mosaic.runMosaic
 
 suspend fun main() = runMosaic {
-	runCounter()
+	Counter()
 }

samples/counter/src/nativeMain/kotlin/example/main.kt

@@ -3,5 +3,5 @@ package example
 import com.jakewharton.mosaic.runMosaicBlocking
 
 fun main() = runMosaicBlocking {
-	runCounter()
+	Counter()
 }

samples/demo/src/main/kotlin/example/demo.kt

@@ -1,5 +1,6 @@
 package example
 
+import androidx.compose.runtime.LaunchedEffect
 import com.jakewharton.mosaic.LocalTerminal
 import com.jakewharton.mosaic.runMosaicBlocking
 import com.jakewharton.mosaic.text.SpanStyle
@@ -10,29 +11,29 @@ import com.jakewharton.mosaic.ui.Text
 import kotlinx.coroutines.suspendCancellableCoroutine
 
 fun main() = runMosaicBlocking {
-	setContent {
-		val terminal = LocalTerminal.current
-		Text(
-			buildAnnotatedString {
-				append("Terminal(")
-				withStyle(SpanStyle(color = Color.BrightGreen)) {
-					append("width=")
-				}
-				withStyle(SpanStyle(color = Color.BrightBlue)) {
-					append(terminal.size.width.toString())
-				}
-				append(", ")
-				withStyle(SpanStyle(color = Color.BrightGreen)) {
-					append("height=")
-				}
-				withStyle(SpanStyle(color = Color.BrightBlue)) {
-					append(terminal.size.height.toString())
-				}
-				append(")")
-			},
-		)
-	}
+	val terminal = LocalTerminal.current
+	Text(
+		buildAnnotatedString {
+			append("Terminal(")
+			withStyle(SpanStyle(color = Color.BrightGreen)) {
+				append("width=")
+			}
+			withStyle(SpanStyle(color = Color.BrightBlue)) {
+				append(terminal.size.width.toString())
+			}
+			append(", ")
+			withStyle(SpanStyle(color = Color.BrightGreen)) {
+				append("height=")
+			}
+			withStyle(SpanStyle(color = Color.BrightBlue)) {
+				append(terminal.size.height.toString())
+			}
+			append(")")
+		},
+	)
 
-	// Run forever!
-	suspendCancellableCoroutine { }
+	LaunchedEffect(Unit) {
+		// Run forever!
+		suspendCancellableCoroutine<Unit> {}
+	}
 }