Timeline: refactor factory/cache

features/messages/src/main/java/io/element/android/x/features/messages/MessageTimelineItemStateFactory.kt

@@ -10,11 +10,15 @@ import io.element.android.x.features.messages.model.MessagesItemGroupPosition
 import io.element.android.x.features.messages.model.MessagesItemReactionState
 import io.element.android.x.features.messages.model.MessagesTimelineItemState
 import io.element.android.x.features.messages.model.content.*
+import io.element.android.x.features.messages.util.invalidateLast
 import io.element.android.x.matrix.MatrixClient
 import io.element.android.x.matrix.media.MediaResolver
 import io.element.android.x.matrix.room.MatrixRoom
 import io.element.android.x.matrix.timeline.MatrixTimelineItem
 import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 import kotlinx.coroutines.withContext
@@ -32,51 +36,66 @@ class MessageTimelineItemStateFactory(
     private val dispatcher: CoroutineDispatcher,
 ) {
 
-    private val timelineItemCaches = arrayListOf<MessagesTimelineItemState?>()
-    private var currentSnapshot: List<MatrixTimelineItem> = emptyList()
+    private val timelineItemStates = MutableStateFlow<List<MessagesTimelineItemState>>(emptyList())
+    private val timelineItemStatesCache = arrayListOf<MessagesTimelineItemState?>()
+
+    // Items from rust sdk, used for diffing
+    private var timelineItems: List<MatrixTimelineItem> = emptyList()
 
     private val lock = Mutex()
-    private val cacheInvalidator = CacheInvalidator(timelineItemCaches)
+    private val cacheInvalidator = CacheInvalidator(timelineItemStatesCache)
+
+    fun flow(): StateFlow<List<MessagesTimelineItemState>> = timelineItemStates.asStateFlow()
 
-    suspend fun create(
+    suspend fun replaceWith(
         timelineItems: List<MatrixTimelineItem>,
-    ): List<MessagesTimelineItemState> =
-        withContext(dispatcher) {
-            lock.withLock {
-                calculateAndApplyDiff(timelineItems)
-                getOrCreateFromCache(timelineItems)
-            }
+    ) = withContext(dispatcher) {
+        lock.withLock {
+            calculateAndApplyDiff(timelineItems)
+            buildAndEmitTimelineItemStates(timelineItems)
         }
+    }
 
-    private suspend fun getOrCreateFromCache(timelineItems: List<MatrixTimelineItem>): List<MessagesTimelineItemState> {
-        val messagesTimelineItemState = ArrayList<MessagesTimelineItemState>()
-        for (index in timelineItemCaches.indices.reversed()) {
-            val cacheItem = timelineItemCaches[index]
+    suspend fun pushItem(
+        timelineItem: MatrixTimelineItem,
+    ) = withContext(dispatcher) {
+        lock.withLock {
+            // Makes sure to invalidate last as we need to recompute some data (like groupPosition)
+            timelineItemStatesCache.invalidateLast()
+            timelineItemStatesCache.add(null)
+            timelineItems = timelineItems + timelineItem
+            buildAndEmitTimelineItemStates(timelineItems)
+        }
+    }
+
+    private suspend fun buildAndEmitTimelineItemStates(timelineItems: List<MatrixTimelineItem>) {
+        val newTimelineItemStates = ArrayList<MessagesTimelineItemState>()
+        for (index in timelineItemStatesCache.indices.reversed()) {
+            val cacheItem = timelineItemStatesCache[index]
             if (cacheItem == null) {
                 buildAndCacheItem(timelineItems, index)?.also { timelineItemState ->
-                    messagesTimelineItemState.add(timelineItemState)
+                    newTimelineItemStates.add(timelineItemState)
                 }
             } else {
-                messagesTimelineItemState.add(cacheItem)
+                newTimelineItemStates.add(cacheItem)
             }
         }
-        return messagesTimelineItemState
+        timelineItemStates.emit(newTimelineItemStates)
     }
 
 
-    private fun calculateAndApplyDiff(timelineItems: List<MatrixTimelineItem>) {
+    private fun calculateAndApplyDiff(newTimelineItems: List<MatrixTimelineItem>) {
         val timeToDiff = measureTimeMillis {
             val diffCallback =
                 MatrixTimelineItemsDiffCallback(
-                    oldList = currentSnapshot,
-                    newList = timelineItems
+                    oldList = timelineItems,
+                    newList = newTimelineItems
                 )
-
             val diffResult = DiffUtil.calculateDiff(diffCallback, false)
-            currentSnapshot = timelineItems
+            timelineItems = newTimelineItems
             diffResult.dispatchUpdatesTo(cacheInvalidator)
         }
-        Timber.v("Time to apply diff on new list of ${timelineItems.size} items: $timeToDiff ms")
+        Timber.v("Time to apply diff on new list of ${newTimelineItems.size} items: $timeToDiff ms")
     }
 
     private suspend fun buildAndCacheItem(
@@ -97,7 +116,7 @@ class MessageTimelineItemStateFactory(
                 )
                 MatrixTimelineItem.Other -> null
             }
-        timelineItemCaches[index] = timelineItemState
+        timelineItemStatesCache[index] = timelineItemState
         return timelineItemState
     }
 
@@ -213,4 +232,5 @@ class MessageTimelineItemStateFactory(
             .resolve(url, kind = MediaResolver.Kind.Thumbnail(size.value))
         return AvatarData(name, model, size)
     }
+
 }

features/messages/src/main/java/io/element/android/x/features/messages/MessagesViewModel.kt

@@ -17,13 +17,12 @@ import io.element.android.x.matrix.MatrixInstance
 import io.element.android.x.matrix.media.MediaResolver
 import io.element.android.x.matrix.room.MatrixRoom
 import io.element.android.x.matrix.timeline.MatrixTimeline
+import io.element.android.x.matrix.timeline.MatrixTimelineItem
 import io.element.android.x.textcomposer.MessageComposerMode
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.launchIn
-import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.launch
-import timber.log.Timber
 
 private const val PAGINATION_COUNT = 50
 
@@ -57,6 +56,14 @@ class MessagesViewModel(
         }
     }
 
+    private val timelineCallback = object : MatrixTimeline.Callback {
+        override fun onPushedTimelineItem(timelineItem: MatrixTimelineItem) {
+            viewModelScope.launch {
+                messageTimelineItemStateFactory.pushItem(timelineItem)
+            }
+        }
+    }
+
     init {
         handleInit()
     }
@@ -93,7 +100,10 @@ class MessagesViewModel(
         return currentState.itemActionsSheetState.invoke()?.targetItem
     }
 
-    fun handleItemAction(action: MessagesItemAction, targetEvent: MessagesTimelineItemState.MessageEvent) {
+    fun handleItemAction(
+        action: MessagesItemAction,
+        targetEvent: MessagesTimelineItemState.MessageEvent
+    ) {
         viewModelScope.launch(Dispatchers.Default) {
             when (action) {
                 MessagesItemAction.Copy -> notImplementedYet()
@@ -109,46 +119,10 @@ class MessagesViewModel(
         setComposerMode(MessageComposerMode.Normal(""))
     }
 
-    private fun handleActionEdit(targetEvent: MessagesTimelineItemState.MessageEvent) {
-        setComposerMode(
-            MessageComposerMode.Edit(
-                targetEvent.id,
-                (targetEvent.content as? MessagesTimelineItemTextBasedContent)?.body.orEmpty()
-            )
-        )
-    }
-
-    private fun handleActionReply(targetEvent: MessagesTimelineItemState.MessageEvent) {
-        setComposerMode(MessageComposerMode.Reply(targetEvent.safeSenderName, targetEvent.id, ""))
-    }
-
-    private fun setComposerMode(mode: MessageComposerMode) {
-        setState {
-            copy(
-                composerMode = mode,
-                highlightedEventId = mode.relatedEventId
-            )
-        }
-    }
-
-    private fun notImplementedYet() {
-        setSnackbarContent("Not implemented yet!")
-    }
-
     fun onSnackbarShown() {
         setSnackbarContent(null)
     }
 
-    private fun setSnackbarContent(message: String?) {
-        setState { copy(snackbarContent = message) }
-    }
-
-    private fun handleActionRedact(event: MessagesTimelineItemState.MessageEvent) {
-        viewModelScope.launch {
-            room.redactEvent(event.id)
-        }
-    }
-
     fun computeActionsSheetState(messagesTimelineItemState: MessagesTimelineItemState.MessageEvent?) {
         if (messagesTimelineItemState == null) {
             setState { copy(itemActionsSheetState = Uninitialized) }
@@ -181,6 +155,7 @@ class MessagesViewModel(
 
     private fun handleInit() {
         timeline.initialize()
+        timeline.callback = timelineCallback
         room.syncUpdateFlow()
             .onEach {
                 val avatarData =
@@ -194,12 +169,52 @@ class MessagesViewModel(
 
         timeline
             .timelineItems()
-            .map(messageTimelineItemStateFactory::create)
+            .onEach(messageTimelineItemStateFactory::replaceWith)
+            .launchIn(viewModelScope)
+
+        messageTimelineItemStateFactory
+            .flow()
             .execute {
                 copy(timelineItems = it)
             }
     }
 
+    private fun setSnackbarContent(message: String?) {
+        setState { copy(snackbarContent = message) }
+    }
+
+    private fun handleActionRedact(event: MessagesTimelineItemState.MessageEvent) {
+        viewModelScope.launch {
+            room.redactEvent(event.id)
+        }
+    }
+
+    private fun handleActionEdit(targetEvent: MessagesTimelineItemState.MessageEvent) {
+        setComposerMode(
+            MessageComposerMode.Edit(
+                targetEvent.id,
+                (targetEvent.content as? MessagesTimelineItemTextBasedContent)?.body.orEmpty()
+            )
+        )
+    }
+
+    private fun handleActionReply(targetEvent: MessagesTimelineItemState.MessageEvent) {
+        setComposerMode(MessageComposerMode.Reply(targetEvent.safeSenderName, targetEvent.id, ""))
+    }
+
+    private fun setComposerMode(mode: MessageComposerMode) {
+        setState {
+            copy(
+                composerMode = mode,
+                highlightedEventId = mode.relatedEventId
+            )
+        }
+    }
+
+    private fun notImplementedYet() {
+        setSnackbarContent("Not implemented yet!")
+    }
+
     private suspend fun loadAvatarData(
         name: String,
         url: String?,
@@ -212,6 +227,7 @@ class MessagesViewModel(
 
     override fun onCleared() {
         super.onCleared()
+        timeline.callback = null
         timeline.dispose()
     }
 }

features/messages/src/main/java/io/element/android/x/features/messages/diff/CacheInvalidator.kt

@@ -2,36 +2,39 @@ package io.element.android.x.features.messages.diff
 
 import androidx.recyclerview.widget.ListUpdateCallback
 import io.element.android.x.features.messages.model.MessagesTimelineItemState
+import io.element.android.x.features.messages.util.invalidateLast
 import timber.log.Timber
 
-internal class CacheInvalidator(private val timelineItemCache: MutableList<MessagesTimelineItemState?>) :
+internal class CacheInvalidator(private val itemStatesCache: MutableList<MessagesTimelineItemState?>) :
     ListUpdateCallback {
 
     override fun onChanged(position: Int, count: Int, payload: Any?) {
         Timber.v("onChanged(position= $position, count= $count")
         (position until position + count).forEach {
             // Invalidate cache
-            timelineItemCache[it] = null
+            itemStatesCache[it] = null
         }
     }
 
     override fun onMoved(fromPosition: Int, toPosition: Int) {
         Timber.v("onMoved(fromPosition= $fromPosition, toPosition= $toPosition")
-        val model = timelineItemCache.removeAt(fromPosition)
-        timelineItemCache.add(toPosition, model)
+        val model = itemStatesCache.removeAt(fromPosition)
+        itemStatesCache.add(toPosition, model)
     }
 
     override fun onInserted(position: Int, count: Int) {
         Timber.v("onInserted(position= $position, count= $count")
+        itemStatesCache.invalidateLast()
         repeat(count) {
-            timelineItemCache.add(position, null)
+            itemStatesCache.add(position, null)
         }
     }
 
     override fun onRemoved(position: Int, count: Int) {
         Timber.v("onRemoved(position= $position, count= $count")
+        itemStatesCache.invalidateLast()
         repeat(count) {
-            timelineItemCache.removeAt(position)
+            itemStatesCache.removeAt(position)
         }
     }
 

features/messages/src/main/java/io/element/android/x/features/messages/util/MutableListExt.kt

@@ -0,0 +1,8 @@
+package io.element.android.x.features.messages.util
+
+internal inline fun <reified T> MutableList<T?>.invalidateLast() {
+    val indexOfLast = size
+    if (indexOfLast > 0) {
+        set(indexOfLast - 1, null)
+    }
+}

libraries/matrix/src/main/java/io/element/android/x/matrix/timeline/MatrixTimeline.kt

@@ -1,7 +1,6 @@
 package io.element.android.x.matrix.timeline
 
 import io.element.android.x.core.coroutine.CoroutineDispatchers
-import io.element.android.x.matrix.core.EventId
 import io.element.android.x.matrix.room.MatrixRoom
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
@@ -20,10 +19,10 @@ class MatrixTimeline(
     private val coroutineScope: CoroutineScope,
     private val coroutineDispatchers: CoroutineDispatchers,
 ) : TimelineListener {
+
     interface Callback {
-        fun onUpdatedTimelineItem(eventId: EventId)
-        fun onStartedBackPaginating()
-        fun onFinishedBackPaginating()
+        fun onUpdatedTimelineItem(timelineItem: MatrixTimelineItem) = Unit
+        fun onPushedTimelineItem(timelineItem: MatrixTimelineItem) = Unit
     }
 
     var callback: Callback? = null
@@ -48,12 +47,14 @@ class MatrixTimeline(
             TimelineChange.PUSH -> {
                 Timber.v("Apply push on list with size: $size")
                 val item = diff.push()?.asMatrixTimelineItem() ?: return
+                callback?.onPushedTimelineItem(item)
                 add(item)
             }
             TimelineChange.UPDATE_AT -> {
                 val updateAtData = diff.updateAt() ?: return
                 Timber.v("Apply $updateAtData on list with size: $size")
                 val item = updateAtData.item.asMatrixTimelineItem()
+                callback?.onUpdatedTimelineItem(item)
                 set(updateAtData.index.toInt(), item)
             }
             TimelineChange.INSERT_AT -> {