@@ -1,6 +1,7 @@
11package soy.iko.opencode.data.repo
22
33import soy.iko.opencode.data.model.BusEvent
4+import soy.iko.opencode.data.model.MessageInfo
45import soy.iko.opencode.data.model.MessagePartRemoved
56import soy.iko.opencode.data.model.MessagePartUpdated
67import soy.iko.opencode.data.model.MessageRemoved
@@ -17,11 +18,13 @@ import soy.iko.opencode.data.network.OpencodeApiClient
1718import soy.iko.opencode.util.runCatchingCancellable
1819import soy.iko.opencode.util.safeExceptionSummary
1920import android.util.Log
21+import kotlinx.coroutines.Dispatchers
2022import kotlinx.coroutines.channels.Channel
2123import kotlinx.coroutines.flow.Flow
2224import kotlinx.coroutines.flow.channelFlow
2325import kotlinx.coroutines.flow.consumeAsFlow
2426import kotlinx.coroutines.flow.conflate
27+import kotlinx.coroutines.flow.flowOn
2528import kotlinx.coroutines.launch
2629import kotlinx.coroutines.sync.Mutex
2730import kotlinx.coroutines.sync.withLock
@@ -159,7 +162,13 @@ open class SessionRepository(
159162 val snapshot = lock.withLock { store.snapshot() }
160163 send(snapshot)
161164 }
162− }.conflate()
165+ // Run the whole reducer pipeline — event reduction, list rebuilds, and the O(N)
166+ // snapshot copy — on Dispatchers.Default rather than the collector's context. The
167+ // sole consumer collects via stateIn(viewModelScope), i.e. Main.immediate, so without
168+ // this every streamed token's reduce()/snapshot() would execute on the UI thread
169+ // (hundreds/sec on a fast run). The seed REST calls here only suspend, so moving them
170+ // off Main is harmless. JSON decoding already runs off-Main upstream in EventStreamClient.
171+ }.flowOn(Dispatchers.Default).conflate()
163172
164173 companion object {
165174 /** Convenience: is this event a run-completion signal for [sessionId]?
@@ -178,8 +187,34 @@ open class SessionRepository(
178187 * Exposed as `internal` so the reducer logic can be unit-tested directly.
179188 */
180189internal class MessageStore {
181− // messageId -> (info + parts), insertion-ordered.
182− private val messages = LinkedHashMap<String, MessageWithParts>()
190+ // messageId -> mutable holder (info + insertion-ordered parts); the map itself is
191+ // insertion-ordered.
192+ private val messages = LinkedHashMap<String, Holder>()
193+
194+ /**
195+ * Mutable per-message holder. Parts live in an insertion-ordered map keyed by part id, so a
196+ * streamed part update is an O(1) put instead of an O(P) list rebuild. The immutable
197+ * [MessageWithParts] view is materialized lazily and cached until the next mutation — so a
198+ * burst of streamed token updates rebuilds a message's part list at most once per published
199+ * [snapshot], not once per token. The cached view is also reference-stable across snapshots
200+ * while unchanged, so Compose skips recomposing bubbles that didn't change.
201+ */
202+ private class Holder(var info: MessageInfo, val parts: LinkedHashMap<String, Part>) {
203+ private var cached: MessageWithParts? = null
204+ fun view(): MessageWithParts =
205+ cached ?: MessageWithParts(info, parts.values.toList()).also { cached = it }
206+ /** Adopt an existing immutable instance as the cache — avoids an immediate rebuild and
207+ * preserves its identity (e.g. when seeding an authoritative REST snapshot). */
208+ fun adopt(m: MessageWithParts) { cached = m }
209+ fun invalidate() { cached = null }
210+ }
211+
212+ /** Build a [Holder] from an immutable message, reusing [m] itself as the cached view. */
213+ private fun holderOf(m: MessageWithParts): Holder {
214+ val parts = LinkedHashMap<String, Part>(m.parts.size)
215+ for (p in m.parts) parts[p.id] = p
216+ return Holder(m.info, parts).also { it.adopt(m) }
217+ }
183218
184219 /** Part ids upserted via [reduce] since the last reseed pivot ([beginReseed]) or seed.
185220 * On a reconnect re-seed these are the parts that streamed in *during* the REST fetch,
@@ -211,7 +246,7 @@ internal class MessageStore {
211246 /** Maximum number of messages to keep in memory; oldest are evicted when exceeded. */
212247 internal val maxMessages = NetworkConfig.maxInMemoryMessages
213248
214− fun snapshot(): List<MessageWithParts> = messages.values.toList()
249+ fun snapshot(): List<MessageWithParts> = messages.values.map { it.view() }
215250
216251 fun seed(initial: List<MessageWithParts>, prune: Boolean = false): Boolean {
217252 var changed = false
@@ -230,7 +265,7 @@ internal class MessageStore {
230265 for (m in initial) {
231266 val existing = messages[m.info.id]
232267 if (existing == null) {
233− messages[m.info.id] = m
268+ messages[m.info.id] = holderOf(m)
234269 changed = true
235270 } else {
236271 // A part streamed in between subscribe and this initial load may already
@@ -238,7 +273,8 @@ internal class MessageStore {
238273 // overwriting so that newer streamed part isn't discarded: take the
239274 // snapshot's part order, swap in the streamed version where ids overlap,
240275 // append any streamed-only parts, and adopt the authoritative REST info.
241− val streamedById = existing.parts.associateBy { it.id }
276+ // existing.parts is already keyed by part id (insertion-ordered).
277+ val streamedById = existing.parts
242278 val snapshotIds = m.parts.mapTo(mutableSetOf()) { it.id }
243279 // Prefer the in-memory version of an overlapping part ONLY if it streamed
244280 // in since the reseed pivot (i.e. during this fetch), so it's newer than
@@ -247,7 +283,7 @@ internal class MessageStore {
247283 val ordered = m.parts.map { p ->
248284 if (p.id in streamedSincePivot) streamedById[p.id] ?: p else p
249285 }.toMutableList()
250− for (p in existing.parts) {
286+ for (p in existing.parts.values) {
251287 if (p.id !in snapshotIds) ordered.add(p)
252288 }
253289 // Adopt the authoritative REST info unless the in-memory info was itself
@@ -256,8 +292,8 @@ internal class MessageStore {
256292 // in-memory info that went stale during a disconnect.
257293 val info = if (existing.info is UnknownMessage || m.info.id !in messageInfoUpdatedSincePivot) m.info else existing.info
258294 val merged = MessageWithParts(info = info, parts = ordered)
259− if (merged != existing) {
260− messages[m.info.id] = merged
295+ if (merged != existing.view()) {
296+ messages[m.info.id] = holderOf(merged)
261297 changed = true
262298 }
263299 }
@@ -284,7 +320,13 @@ internal class MessageStore {
284320 private fun reorderByTime(): Boolean {
285321 val entries = messages.entries.map { it.key to it.value }
286322 val sorted = entries.sortedBy { it.second.info.time?.created ?: Long.MAX_VALUE }
287− if (sorted.map { it.first } == entries.map { it.first }) return false
323+ // Detect a reordering by comparing keys positionally, instead of allocating two more
324+ // N-sized key lists just to compare them.
325+ var changed = false
326+ for (i in sorted.indices) {
327+ if (sorted[i].first != entries[i].first) { changed = true; break }
328+ }
329+ if (!changed) return false
288330 messages.clear()
289331 for ((k, v) in sorted) messages[k] = v
290332 return true
@@ -293,30 +335,7 @@ internal class MessageStore {
293335 /** Returns true if the state changed (and a new snapshot should be published). */
294336 fun reduce(sessionId: String, event: BusEvent): Boolean {
295337 return when (event) {
296− is MessageUpdated -> {
297− val info = event.properties.info
298− if (info.sessionID != sessionId) {
299− false
300− } else {
301− // An UnknownMessage with an empty id (e.g. from an unrecognized
302− // server role with no id field) would collide with other such
303− // messages in the map. Generate a unique synthetic key so each
304− // unknown message gets its own entry instead of overwriting others.
305− val key = if (info.id.isEmpty() && info is UnknownMessage) {
306− "unknown-${unknownCounter.incrementAndGet()}"
307− } else {
308− info.id
309− }
310− val existing = messages[key]
311− messages[key] = existing?.copy(info = info) ?: MessageWithParts(info)
312− // Record that this message's info streamed in live, so a re-seed after a
313− // reconnect knows it's newer than the REST snapshot and keeps it (see
314− // seed()/beginReseed()).
315− messageInfoUpdatedSincePivot.add(key)
316− if (existing == null) evictOldMessages()
317− true
318− }
319− }
338+ is MessageUpdated -> handleMessageUpdated(sessionId, event)
320339
321340 is MessagePartUpdated -> {
322341 val part = event.properties.part
@@ -343,6 +362,34 @@ internal class MessageStore {
343362 }
344363 }
345364
365+ private fun handleMessageUpdated(sessionId: String, event: MessageUpdated): Boolean {
366+ val info = event.properties.info
367+ if (info.sessionID != sessionId) return false
368+ // An UnknownMessage with an empty id (e.g. from an unrecognized server role with no id
369+ // field) would collide with other such messages in the map. Generate a unique synthetic
370+ // key so each unknown message gets its own entry instead of overwriting others.
371+ val key = if (info.id.isEmpty() && info is UnknownMessage) {
372+ "unknown-${unknownCounter.incrementAndGet()}"
373+ } else {
374+ info.id
375+ }
376+ val existing = messages[key]
377+ // A server re-sending byte-identical info is a no-op: skip it so we don't publish a
378+ // redundant snapshot (mirrors upsertPart's identical-part guard).
379+ if (existing != null && existing.info == info) return false
380+ if (existing == null) {
381+ messages[key] = Holder(info, LinkedHashMap())
382+ } else {
383+ existing.info = info
384+ existing.invalidate()
385+ }
386+ // Record that this message's info streamed in live, so a re-seed after a reconnect knows
387+ // it's newer than the REST snapshot and keeps it (see seed()/beginReseed()).
388+ messageInfoUpdatedSincePivot.add(key)
389+ if (existing == null) evictOldMessages()
390+ return true
391+ }
392+
346393 /** Returns false if [eventSession] is non-null and doesn't match [sessionId]. */
347394 private fun matchesSession(eventSession: String?, sessionId: String): Boolean =
348395 eventSession == null || eventSession == sessionId
@@ -359,25 +406,27 @@ internal class MessageStore {
359406 if (!matchesSession(event.properties.sessionID, sessionId)) return false
360407 val removed = messages.remove(id) ?: return false
361408 messageInfoUpdatedSincePivot.remove(id)
362− removed.parts.forEach { streamedSincePivot.remove(it.id) }
409+ removed.parts.keys.forEach { streamedSincePivot.remove(it) }
363410 return true
364411 }
365412
366413 private fun upsertPart(messageId: String, part: Part): Boolean {
367414 val current = messages[messageId]
368− val parts = current?.parts.orEmpty()
369− val idx = parts.indexOfFirst { it.id == part.id }
370− if (idx >= 0 && parts[idx] == part) return false
371− val newParts = if (idx >= 0) {
372− parts.toMutableList().also { it[idx] = part }
415+ if (current != null) {
416+ // O(1) map lookup + put, vs the former O(P) indexOfFirst scan and full-list copy
417+ // on every streamed token. A LinkedHashMap re-put keeps the part's original
418+ // position (replace-in-place); a new id is appended — matching the old list order.
419+ if (current.parts[part.id] == part) return false
420+ current.parts[part.id] = part
421+ current.invalidate()
373422 } else {
374− parts + part
375− }
376− messages[messageId] = current?.copy(parts = newParts)
377− ?: MessageWithParts(
423+ val parts = LinkedHashMap<String, Part>()
424+ parts[part.id] = part
425+ messages[messageId] = Holder(
378426 info = UnknownMessage(id = messageId, sessionID = part.sessionID ?: ""),
379− parts = newParts,
427+ parts = parts,
380428 )
429+ }
381430 // Record that this part streamed in, so a re-seed after a reconnect knows it's
382431 // newer than the REST snapshot and keeps it (see seed()/beginReseed()).
383432 streamedSincePivot.add(part.id)
@@ -387,9 +436,8 @@ internal class MessageStore {
387436
388437 private fun removePart(messageId: String, partId: String): Boolean {
389438 val current = messages[messageId] ?: return false
390− val newParts = current.parts.filterNot { it.id == partId }
391− if (newParts.size == current.parts.size) return false
392− messages[messageId] = current.copy(parts = newParts)
439+ if (current.parts.remove(partId) == null) return false
440+ current.invalidate()
393441 return true
394442 }
395443
@@ -404,7 +452,7 @@ internal class MessageStore {
404452 // unbounded over a long-running session that never reconnects (they're only
405453 // otherwise cleared on reseed/seed).
406454 messageInfoUpdatedSincePivot.remove(oldestKey)
407− removed?.parts?.forEach { streamedSincePivot.remove(it.id) }
455+ removed?.parts?.keys?.forEach { streamedSincePivot.remove(it) }
408456 evicted = true
409457 }
410458 return evicted