From 627e2ca5a6b5fe45d4607e59afaeabcf480d0c17 Mon Sep 17 00:00:00 2001
From: Benjamin Bouvier <benjamin@bouvier.cc>
Date: Thu, 27 Mar 2025 14:24:15 +0100
Subject: [PATCH] refactor(linked chunk): rejigger the relational linked chunk
 so it can handle indexed items

This is necessary to save out-of-band items into the relational linked
chunk. I'm not quite sure of the value to keep it generic, at this
point, but at least it makes testing easy.
---
 .../src/event_cache/store/memory_store.rs     |  13 +-
 .../src/linked_chunk/relational.rs            | 199 ++++++++++++------
 2 files changed, 145 insertions(+), 67 deletions(-)

diff --git a/crates/matrix-sdk-base/src/event_cache/store/memory_store.rs b/crates/matrix-sdk-base/src/event_cache/store/memory_store.rs
index 2a7057a55..561f7aa3f 100644
--- a/crates/matrix-sdk-base/src/event_cache/store/memory_store.rs
+++ b/crates/matrix-sdk-base/src/event_cache/store/memory_store.rs
@@ -54,7 +54,7 @@ pub struct MemoryStore {
 struct MemoryStoreInner {
     media: RingBuffer<MediaContent>,
     leases: HashMap<String, (String, Instant)>,
-    events: RelationalLinkedChunk<Event, Gap>,
+    events: RelationalLinkedChunk<OwnedEventId, Event, Gap>,
     media_retention_policy: Option<MediaRetentionPolicy>,
     last_media_cleanup_time: SystemTime,
 }
@@ -209,12 +209,13 @@ impl EventCacheStore for MemoryStore {
     ) -> Result<Option<(Position, Event)>, Self::Error> {
         let inner = self.inner.read().unwrap();
 
-        let event_and_room = inner.events.items().find_map(|(position, event, this_room_id)| {
-            (room_id == this_room_id && event.event_id()? == event_id)
-                .then_some((position, event.clone()))
-        });
+        let pos_and_event =
+            inner.events.items_with_positions().find_map(|(position, event, this_room_id)| {
+                (room_id == this_room_id && event.event_id()? == event_id)
+                    .then_some((position, event.clone()))
+            });
 
-        Ok(event_and_room)
+        Ok(pos_and_event)
     }
 
     async fn add_media_content(
diff --git a/crates/matrix-sdk-common/src/linked_chunk/relational.rs b/crates/matrix-sdk-common/src/linked_chunk/relational.rs
index 361d5c59f..bb3e2acc7 100644
--- a/crates/matrix-sdk-common/src/linked_chunk/relational.rs
+++ b/crates/matrix-sdk-common/src/linked_chunk/relational.rs
@@ -15,10 +15,15 @@
 //! Implementation for a _relational linked chunk_, see
 //! [`RelationalLinkedChunk`].
 
-use ruma::{OwnedRoomId, RoomId};
+use std::{collections::HashMap, hash::Hash};
+
+use ruma::{OwnedEventId, OwnedRoomId, RoomId};
 
 use super::{ChunkContent, ChunkIdentifierGenerator, RawChunk};
-use crate::linked_chunk::{ChunkIdentifier, Position, Update};
+use crate::{
+    deserialized_responses::TimelineEvent,
+    linked_chunk::{ChunkIdentifier, Position, Update},
+};
 
 /// A row of the [`RelationalLinkedChunk::chunks`].
 #[derive(Debug, PartialEq)]
@@ -31,10 +36,10 @@ struct ChunkRow {
 
 /// A row of the [`RelationalLinkedChunk::items`].
 #[derive(Debug, PartialEq)]
-struct ItemRow<Item, Gap> {
+struct ItemRow<ItemId, Gap> {
     room_id: OwnedRoomId,
     position: Position,
-    item: Either<Item, Gap>,
+    item: Either<ItemId, Gap>,
 }
 
 /// Kind of item.
@@ -66,23 +71,49 @@ enum Either<Item, Gap> {
 ///
 /// [`LinkedChunk`]: super::LinkedChunk
 #[derive(Debug)]
-pub struct RelationalLinkedChunk<Item, Gap> {
+pub struct RelationalLinkedChunk<ItemId, Item, Gap> {
     /// Chunks.
     chunks: Vec<ChunkRow>,
 
-    /// Items.
-    items: Vec<ItemRow<Item, Gap>>,
+    /// Items chunks.
+    items_chunks: Vec<ItemRow<ItemId, Gap>>,
+
+    /// The items' content themselves.
+    items: HashMap<OwnedRoomId, HashMap<ItemId, Item>>,
 }
 
-impl<Item, Gap> RelationalLinkedChunk<Item, Gap> {
+/// The [`IndexableItem`] trait is used to mark items that can be indexed into a
+/// [`RelationalLinkedChunk`].
+pub trait IndexableItem {
+    type ItemId: Hash + PartialEq + Eq + Clone;
+
+    /// Return the identifier of the item.
+    fn id(&self) -> Self::ItemId;
+}
+
+impl IndexableItem for TimelineEvent {
+    type ItemId = OwnedEventId;
+
+    fn id(&self) -> Self::ItemId {
+        self.event_id()
+            .expect("all events saved into a relational linked chunk must have a valid event id")
+    }
+}
+
+impl<ItemId, Item, Gap> RelationalLinkedChunk<ItemId, Item, Gap>
+where
+    Item: IndexableItem<ItemId = ItemId>,
+    ItemId: Hash + PartialEq + Eq + Clone,
+{
     /// Create a new relational linked chunk.
     pub fn new() -> Self {
-        Self { chunks: Vec::new(), items: Vec::new() }
+        Self { chunks: Vec::new(), items_chunks: Vec::new(), items: HashMap::new() }
     }
 
     /// Removes all the chunks and items from this relational linked chunk.
     pub fn clear(&mut self) {
         self.chunks.clear();
+        self.items_chunks.clear();
         self.items.clear();
     }
 
@@ -97,7 +128,7 @@ impl<Item, Gap> RelationalLinkedChunk<Item, Gap> {
 
                 Update::NewGapChunk { previous, new, next, gap } => {
                     insert_chunk(&mut self.chunks, room_id, previous, new, next);
-                    self.items.push(ItemRow {
+                    self.items_chunks.push(ItemRow {
                         room_id: room_id.to_owned(),
                         position: Position::new(new, 0),
                         item: Either::Gap(gap),
@@ -108,7 +139,7 @@ impl<Item, Gap> RelationalLinkedChunk<Item, Gap> {
                     remove_chunk(&mut self.chunks, room_id, chunk_identifier);
 
                     let indices_to_remove = self
-                        .items
+                        .items_chunks
                         .iter()
                         .enumerate()
                         .filter_map(
@@ -121,16 +152,21 @@ impl<Item, Gap> RelationalLinkedChunk<Item, Gap> {
                         .collect::<Vec<_>>();
 
                     for index_to_remove in indices_to_remove.into_iter().rev() {
-                        self.items.remove(index_to_remove);
+                        self.items_chunks.remove(index_to_remove);
                     }
                 }
 
                 Update::PushItems { mut at, items } => {
                     for item in items {
-                        self.items.push(ItemRow {
+                        let item_id = item.id();
+                        self.items
+                            .entry(room_id.to_owned())
+                            .or_default()
+                            .insert(item_id.clone(), item);
+                        self.items_chunks.push(ItemRow {
                             room_id: room_id.to_owned(),
                             position: at,
-                            item: Either::Item(item),
+                            item: Either::Item(item_id),
                         });
                         at.increment_index();
                     }
@@ -138,7 +174,7 @@ impl<Item, Gap> RelationalLinkedChunk<Item, Gap> {
 
                 Update::ReplaceItem { at, item } => {
                     let existing = self
-                        .items
+                        .items_chunks
                         .iter_mut()
                         .find(|item| item.position == at)
                         .expect("trying to replace at an unknown position");
@@ -146,14 +182,16 @@ impl<Item, Gap> RelationalLinkedChunk<Item, Gap> {
                         matches!(existing.item, Either::Item(..)),
                         "trying to replace a gap with an item"
                     );
-                    existing.item = Either::Item(item);
+                    let item_id = item.id();
+                    self.items.entry(room_id.to_owned()).or_default().insert(item_id.clone(), item);
+                    existing.item = Either::Item(item_id);
                 }
 
                 Update::RemoveItem { at } => {
                     let mut entry_to_remove = None;
 
                     for (nth, ItemRow { room_id: room_id_candidate, position, .. }) in
-                        self.items.iter_mut().enumerate()
+                        self.items_chunks.iter_mut().enumerate()
                     {
                         // Filter by room ID.
                         if room_id != room_id_candidate {
@@ -175,12 +213,13 @@ impl<Item, Gap> RelationalLinkedChunk<Item, Gap> {
                         }
                     }
 
-                    self.items.remove(entry_to_remove.expect("Remove an unknown item"));
+                    self.items_chunks.remove(entry_to_remove.expect("Remove an unknown item"));
+                    // We deliberately keep the item in the items collection.
                 }
 
                 Update::DetachLastItems { at } => {
                     let indices_to_remove = self
-                        .items
+                        .items_chunks
                         .iter()
                         .enumerate()
                         .filter_map(
@@ -194,7 +233,7 @@ impl<Item, Gap> RelationalLinkedChunk<Item, Gap> {
                         .collect::<Vec<_>>();
 
                     for index_to_remove in indices_to_remove.into_iter().rev() {
-                        self.items.remove(index_to_remove);
+                        self.items_chunks.remove(index_to_remove);
                     }
                 }
 
@@ -202,7 +241,8 @@ impl<Item, Gap> RelationalLinkedChunk<Item, Gap> {
 
                 Update::Clear => {
                     self.chunks.retain(|chunk| chunk.room_id != room_id);
-                    self.items.retain(|chunk| chunk.room_id != room_id);
+                    self.items_chunks.retain(|chunk| chunk.room_id != room_id);
+                    // We deliberately leave the items intact.
                 }
             }
         }
@@ -299,10 +339,12 @@ impl<Item, Gap> RelationalLinkedChunk<Item, Gap> {
         &'a self,
         room_id: &'a RoomId,
     ) -> impl Iterator<Item = (&'a Item, Position)> {
-        self.items.iter().filter_map(move |item_row| {
+        self.items_chunks.iter().filter_map(move |item_row| {
             if item_row.room_id == room_id {
                 match &item_row.item {
-                    Either::Item(item) => Some((item, item_row.position)),
+                    Either::Item(item_id) => {
+                        Some((self.items.get(room_id)?.get(item_id)?, item_row.position))
+                    }
                     Either::Gap(..) => None,
                 }
             } else {
@@ -311,10 +353,26 @@ impl<Item, Gap> RelationalLinkedChunk<Item, Gap> {
         })
     }
 
-    /// Return an iterator over all items of all rooms, in no particular order.
-    pub fn items(&self) -> impl Iterator<Item = (Position, &Item, &RoomId)> {
-        self.items.iter().filter_map(|item_row| {
-            if let Either::Item(item) = &item_row.item {
+    /// Return an iterator over all items of all rooms, without their actual
+    /// positions.
+    ///
+    /// This will include out-of-band items.
+    pub fn items(&self) -> impl Iterator<Item = (&Item, &RoomId)> {
+        self.items
+            .iter()
+            .flat_map(|(room_id, items)| items.values().map(|item| (item, room_id.as_ref())))
+    }
+
+    /// Return an iterator over all items of all rooms, with their actual
+    /// positions, in no particular order.
+    ///
+    /// This will *not* include out of band items.
+    // TODO(bnjbvr): see if I can delete this one — reviewier, if this comment is
+    // still here during review, please let me know.
+    pub fn items_with_positions(&self) -> impl Iterator<Item = (Position, &Item, &RoomId)> {
+        self.items_chunks.iter().filter_map(|item_row| {
+            if let Either::Item(item_id) = &item_row.item {
+                let item = self.items.get(&item_row.room_id)?.get(item_id)?;
                 Some((item_row.position, item, item_row.room_id.as_ref()))
             } else {
                 None
@@ -323,10 +381,11 @@ impl<Item, Gap> RelationalLinkedChunk<Item, Gap> {
     }
 }
 
-impl<Item, Gap> RelationalLinkedChunk<Item, Gap>
+impl<ItemId, Item, Gap> RelationalLinkedChunk<ItemId, Item, Gap>
 where
     Gap: Clone,
     Item: Clone,
+    ItemId: Hash + PartialEq + Eq,
 {
     /// Loads all the chunks.
     ///
@@ -419,24 +478,29 @@ where
     }
 }
 
-impl<Item, Gap> Default for RelationalLinkedChunk<Item, Gap> {
+impl<ItemId, Item, Gap> Default for RelationalLinkedChunk<ItemId, Item, Gap>
+where
+    Item: IndexableItem<ItemId = ItemId>,
+    ItemId: Hash + PartialEq + Eq + Clone,
+{
     fn default() -> Self {
         Self::new()
     }
 }
 
-fn load_raw_chunk<Item, Gap>(
-    relational_linked_chunk: &RelationalLinkedChunk<Item, Gap>,
+fn load_raw_chunk<ItemId, Item, Gap>(
+    relational_linked_chunk: &RelationalLinkedChunk<ItemId, Item, Gap>,
     chunk_row: &ChunkRow,
     room_id: &RoomId,
 ) -> Result<RawChunk<Item, Gap>, String>
 where
     Item: Clone,
     Gap: Clone,
+    ItemId: Hash + PartialEq + Eq,
 {
     // Find all items that correspond to the chunk.
     let mut items = relational_linked_chunk
-        .items
+        .items_chunks
         .iter()
         .filter(|item_row| {
             item_row.room_id == room_id && item_row.position.chunk_identifier() == chunk_row.chunk
@@ -461,8 +525,8 @@ where
 
             for item_row in items {
                 match &item_row.item {
-                    Either::Item(item_value) => {
-                        collected_items.push((item_value.clone(), item_row.position.index()))
+                    Either::Item(item_id) => {
+                        collected_items.push((item_id, item_row.position.index()))
                     }
 
                     Either::Gap(_) => {
@@ -479,7 +543,12 @@ where
 
             RawChunk {
                 content: ChunkContent::Items(
-                    collected_items.into_iter().map(|(item, _index)| item).collect(),
+                    collected_items
+                        .into_iter()
+                        .filter_map(|(item_id, _index)| {
+                            relational_linked_chunk.items.get(room_id)?.get(item_id).cloned()
+                        })
+                        .collect(),
                 ),
                 previous: chunk_row.previous_chunk,
                 identifier: chunk_row.chunk,
@@ -515,10 +584,18 @@ mod tests {
 
     use super::{super::lazy_loader::from_all_chunks, ChunkIdentifier as CId, *};
 
+    impl IndexableItem for char {
+        type ItemId = char;
+
+        fn id(&self) -> Self::ItemId {
+            *self
+        }
+    }
+
     #[test]
     fn test_new_items_chunk() {
         let room_id = room_id!("!r0:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, ()>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, ()>::new();
 
         relational_linked_chunk.apply_updates(
             room_id,
@@ -569,13 +646,13 @@ mod tests {
             ],
         );
         // Items have not been modified.
-        assert!(relational_linked_chunk.items.is_empty());
+        assert!(relational_linked_chunk.items_chunks.is_empty());
     }
 
     #[test]
     fn test_new_gap_chunk() {
         let room_id = room_id!("!r0:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, ()>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, ()>::new();
 
         relational_linked_chunk.apply_updates(
             room_id,
@@ -620,7 +697,7 @@ mod tests {
         );
         // Items contains the gap.
         assert_eq!(
-            relational_linked_chunk.items,
+            relational_linked_chunk.items_chunks,
             &[ItemRow {
                 room_id: room_id.to_owned(),
                 position: Position::new(CId::new(1), 0),
@@ -632,7 +709,7 @@ mod tests {
     #[test]
     fn test_remove_chunk() {
         let room_id = room_id!("!r0:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, ()>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, ()>::new();
 
         relational_linked_chunk.apply_updates(
             room_id,
@@ -672,13 +749,13 @@ mod tests {
             ],
         );
         // Items no longer contains the gap.
-        assert!(relational_linked_chunk.items.is_empty());
+        assert!(relational_linked_chunk.items_chunks.is_empty());
     }
 
     #[test]
     fn test_push_items() {
         let room_id = room_id!("!r0:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, ()>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, ()>::new();
 
         relational_linked_chunk.apply_updates(
             room_id,
@@ -716,7 +793,7 @@ mod tests {
         );
         // Items contains the pushed items.
         assert_eq!(
-            relational_linked_chunk.items,
+            relational_linked_chunk.items_chunks,
             &[
                 ItemRow {
                     room_id: room_id.to_owned(),
@@ -765,7 +842,7 @@ mod tests {
     #[test]
     fn test_remove_item() {
         let room_id = room_id!("!r0:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, ()>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, ()>::new();
 
         relational_linked_chunk.apply_updates(
             room_id,
@@ -796,7 +873,7 @@ mod tests {
         );
         // Items contains the pushed items.
         assert_eq!(
-            relational_linked_chunk.items,
+            relational_linked_chunk.items_chunks,
             &[
                 ItemRow {
                     room_id: room_id.to_owned(),
@@ -820,7 +897,7 @@ mod tests {
     #[test]
     fn test_detach_last_items() {
         let room_id = room_id!("!r0:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, ()>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, ()>::new();
 
         relational_linked_chunk.apply_updates(
             room_id,
@@ -861,7 +938,7 @@ mod tests {
         );
         // Items contains the pushed items.
         assert_eq!(
-            relational_linked_chunk.items,
+            relational_linked_chunk.items_chunks,
             &[
                 ItemRow {
                     room_id: room_id.to_owned(),
@@ -895,21 +972,21 @@ mod tests {
     #[test]
     fn test_start_and_end_reattach_items() {
         let room_id = room_id!("!r0:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, ()>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, ()>::new();
 
         relational_linked_chunk
             .apply_updates(room_id, vec![Update::StartReattachItems, Update::EndReattachItems]);
 
         // Nothing happened.
         assert!(relational_linked_chunk.chunks.is_empty());
-        assert!(relational_linked_chunk.items.is_empty());
+        assert!(relational_linked_chunk.items_chunks.is_empty());
     }
 
     #[test]
     fn test_clear() {
         let r0 = room_id!("!r0:matrix.org");
         let r1 = room_id!("!r1:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, ()>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, ()>::new();
 
         relational_linked_chunk.apply_updates(
             r0,
@@ -952,7 +1029,7 @@ mod tests {
 
         // Items contains the pushed items.
         assert_eq!(
-            relational_linked_chunk.items,
+            relational_linked_chunk.items_chunks,
             &[
                 ItemRow {
                     room_id: r0.to_owned(),
@@ -992,7 +1069,7 @@ mod tests {
         );
 
         assert_eq!(
-            relational_linked_chunk.items,
+            relational_linked_chunk.items_chunks,
             &[ItemRow {
                 room_id: r1.to_owned(),
                 position: Position::new(CId::new(0), 0),
@@ -1006,7 +1083,7 @@ mod tests {
         let room_id = room_id!("!r0:matrix.org");
 
         // When I reload the linked chunk components from an empty store,
-        let relational_linked_chunk = RelationalLinkedChunk::<char, char>::new();
+        let relational_linked_chunk = RelationalLinkedChunk::<_, char, char>::new();
         let result = relational_linked_chunk.load_all_chunks(room_id).unwrap();
         assert!(result.is_empty());
     }
@@ -1015,7 +1092,7 @@ mod tests {
     fn test_load_all_chunks_with_empty_items() {
         let room_id = room_id!("!r0:matrix.org");
 
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, char>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, char>::new();
 
         // When I store an empty items chunks,
         relational_linked_chunk.apply_updates(
@@ -1035,7 +1112,7 @@ mod tests {
     #[test]
     fn test_rebuild_linked_chunk() {
         let room_id = room_id!("!r0:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, char>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, char>::new();
 
         relational_linked_chunk.apply_updates(
             room_id,
@@ -1070,7 +1147,7 @@ mod tests {
     #[test]
     fn test_replace_item() {
         let room_id = room_id!("!r0:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, ()>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, ()>::new();
 
         relational_linked_chunk.apply_updates(
             room_id,
@@ -1097,7 +1174,7 @@ mod tests {
 
         // Items contains the pushed *and* replaced items.
         assert_eq!(
-            relational_linked_chunk.items,
+            relational_linked_chunk.items_chunks,
             &[
                 ItemRow {
                     room_id: room_id.to_owned(),
@@ -1122,7 +1199,7 @@ mod tests {
     fn test_unordered_events() {
         let room_id = room_id!("!r0:matrix.org");
         let other_room_id = room_id!("!r1:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, ()>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, ()>::new();
 
         relational_linked_chunk.apply_updates(
             room_id,
@@ -1156,7 +1233,7 @@ mod tests {
     #[test]
     fn test_load_last_chunk() {
         let room_id = room_id!("!r0:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, ()>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, ()>::new();
 
         // Case #1: no last chunk.
         {
@@ -1230,7 +1307,7 @@ mod tests {
     #[test]
     fn test_load_last_chunk_with_a_cycle() {
         let room_id = room_id!("!r0:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, ()>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, ()>::new();
 
         relational_linked_chunk.apply_updates(
             room_id,
@@ -1253,7 +1330,7 @@ mod tests {
     #[test]
     fn test_load_previous_chunk() {
         let room_id = room_id!("!r0:matrix.org");
-        let mut relational_linked_chunk = RelationalLinkedChunk::<char, ()>::new();
+        let mut relational_linked_chunk = RelationalLinkedChunk::<_, char, ()>::new();
 
         // Case #1: no chunk at all, equivalent to having an inexistent
         // `before_chunk_identifier`.