base: Add state store method to fetch several display_names at once

Signed-off-by: Kévin Commaille <zecakeh@tedomum.fr>

crates/matrix-sdk-base/Changelog.md

@@ -25,6 +25,11 @@
   `get_room_infos`.
 - `BaseClient::get_stripped_rooms` is deprecated. Use `get_rooms_filtered` with
   `RoomStateFilter::INVITED` instead.
+- Add methods to `StateStore` to be able to retrieve data in batch
+  - `get_state_events_for_keys`
+  - `get_profiles`
+  - `get_presence_events`
+  - `get_users_with_display_names`
 
 ## 0.5.1
 

crates/matrix-sdk-base/src/store/integration_tests.rs

@@ -75,6 +75,8 @@ pub trait StateStoreIntegrationTests {
     async fn test_room_removal(&self) -> Result<()>;
     /// Test presence saving.
     async fn test_presence_saving(&self);
+    /// Test display names saving.
+    async fn test_display_names_saving(&self);
 }
 
 #[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
@@ -1061,6 +1063,65 @@ impl StateStoreIntegrationTests for DynStateStore {
         let presence_events = self.get_presence_events(&[]).await;
         assert!(presence_events.unwrap().is_empty());
     }
+
+    async fn test_display_names_saving(&self) {
+        let room_id = room_id!("!test_display_names_saving:localhost");
+        let user_id = user_id();
+        let user_display_name = "User";
+        let second_user_id = user_id!("@second:localhost");
+        let third_user_id = user_id!("@third:localhost");
+        let other_display_name = "Raoul";
+        let unknown_display_name = "Unknown";
+
+        // No event in store.
+        let mut display_names = vec![user_display_name.to_owned()];
+        let users = self.get_users_with_display_name(room_id, user_display_name).await.unwrap();
+        assert!(users.is_empty());
+        let names = self.get_users_with_display_names(room_id, &display_names).await.unwrap();
+        assert!(names.is_empty());
+
+        // One event in store.
+        let mut changes = StateChanges::default();
+        changes
+            .ambiguity_maps
+            .entry(room_id.to_owned())
+            .or_default()
+            .insert(user_display_name.to_owned(), [user_id.to_owned()].into());
+        self.save_changes(&changes).await.unwrap();
+
+        let users = self.get_users_with_display_name(room_id, user_display_name).await.unwrap();
+        assert_eq!(users.len(), 1);
+        let names = self.get_users_with_display_names(room_id, &display_names).await.unwrap();
+        assert_eq!(names.len(), 1);
+        assert_eq!(names.get(&user_display_name).unwrap().len(), 1);
+
+        // Several events in store.
+        let mut changes = StateChanges::default();
+        changes.ambiguity_maps.entry(room_id.to_owned()).or_default().insert(
+            other_display_name.to_owned(),
+            [second_user_id.to_owned(), third_user_id.to_owned()].into(),
+        );
+        self.save_changes(&changes).await.unwrap();
+
+        display_names.push(other_display_name.to_owned());
+        let users = self.get_users_with_display_name(room_id, user_display_name).await.unwrap();
+        assert_eq!(users.len(), 1);
+        let users = self.get_users_with_display_name(room_id, other_display_name).await.unwrap();
+        assert_eq!(users.len(), 2);
+        let names = self.get_users_with_display_names(room_id, &display_names).await.unwrap();
+        assert_eq!(names.len(), 2);
+        assert_eq!(names.get(&user_display_name).unwrap().len(), 1);
+        assert_eq!(names.get(&other_display_name).unwrap().len(), 2);
+
+        // Several events in store with one unknown.
+        display_names.push(unknown_display_name.to_owned());
+        let names = self.get_users_with_display_names(room_id, &display_names).await.unwrap();
+        assert_eq!(names.len(), 2);
+
+        // Empty user IDs list.
+        let names = self.get_users_with_display_names(room_id, &[]).await;
+        assert!(names.unwrap().is_empty());
+    }
 }
 
 /// Macro building to allow your StateStore implementation to run the entire
@@ -1191,6 +1252,12 @@ macro_rules! statestore_integration_tests {
             let store = get_store().await.expect("creating store failed").into_state_store();
             store.test_presence_saving().await;
         }
+
+        #[async_test]
+        async fn test_display_names_saving() {
+            let store = get_store().await.expect("creating store failed").into_state_store();
+            store.test_display_names_saving().await;
+        }
     };
 }
 

crates/matrix-sdk-base/src/store/memory_store.rs

@@ -488,6 +488,27 @@ impl MemoryStore {
             .collect()
     }
 
+    async fn get_users_with_display_names<'a, I>(
+        &self,
+        room_id: &RoomId,
+        display_names: I,
+    ) -> Result<BTreeMap<&'a str, BTreeSet<OwnedUserId>>>
+    where
+        I: IntoIterator<Item = &'a str>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        let display_names = display_names.into_iter();
+        if display_names.len() == 0 {
+            return Ok(BTreeMap::new());
+        }
+
+        let Some(room_names) = self.display_names.get(room_id) else {
+            return Ok(BTreeMap::new());
+        };
+
+        Ok(display_names.filter_map(|n| room_names.get(n).map(|d| (n, d.clone()))).collect())
+    }
+
     async fn get_account_data_event(
         &self,
         event_type: GlobalAccountDataEventType,
@@ -696,12 +717,23 @@ impl StateStore for MemoryStore {
         display_name: &str,
     ) -> Result<BTreeSet<OwnedUserId>> {
         Ok(self
-            .display_names
-            .get(room_id)
-            .and_then(|d| d.get(display_name).map(|d| d.clone()))
+            .get_users_with_display_names(room_id, iter::once(display_name))
+            .await?
+            .into_values()
+            .next()
             .unwrap_or_default())
     }
 
+    async fn get_users_with_display_names<'a>(
+        &self,
+        room_id: &RoomId,
+        display_names: &'a [String],
+    ) -> Result<BTreeMap<&'a str, BTreeSet<OwnedUserId>>> {
+        Ok(self
+            .get_users_with_display_names(room_id, display_names.iter().map(AsRef::as_ref))
+            .await?)
+    }
+
     async fn get_account_data_event(
         &self,
         event_type: GlobalAccountDataEventType,

crates/matrix-sdk-base/src/store/traits.rs

@@ -215,6 +215,19 @@ pub trait StateStore: AsyncTraitDeps {
         display_name: &str,
     ) -> Result<BTreeSet<OwnedUserId>, Self::Error>;
 
+    /// Get all the users that use the given display names in the given room.
+    ///
+    /// # Arguments
+    ///
+    /// * `room_id` - The ID of the room to fetch the display names for.
+    ///
+    /// * `display_names` - The display names that the users use.
+    async fn get_users_with_display_names<'a>(
+        &self,
+        room_id: &RoomId,
+        display_names: &'a [String],
+    ) -> Result<BTreeMap<&'a str, BTreeSet<OwnedUserId>>, Self::Error>;
+
     /// Get an event out of the account data store.
     ///
     /// # Arguments
@@ -483,6 +496,14 @@ impl<T: StateStore> StateStore for EraseStateStoreError<T> {
         self.0.get_users_with_display_name(room_id, display_name).await.map_err(Into::into)
     }
 
+    async fn get_users_with_display_names<'a>(
+        &self,
+        room_id: &RoomId,
+        display_names: &'a [String],
+    ) -> Result<BTreeMap<&'a str, BTreeSet<OwnedUserId>>, Self::Error> {
+        self.0.get_users_with_display_names(room_id, display_names).await.map_err(Into::into)
+    }
+
     async fn get_account_data_event(
         &self,
         event_type: GlobalAccountDataEventType,

crates/matrix-sdk-indexeddb/src/state_store/mod.rs

@@ -1029,6 +1029,35 @@ impl_state_store!({
             .unwrap_or_else(|| Ok(Default::default()))
     }
 
+    async fn get_users_with_display_names<'a>(
+        &self,
+        room_id: &RoomId,
+        display_names: &'a [String],
+    ) -> Result<BTreeMap<&'a str, BTreeSet<OwnedUserId>>> {
+        if display_names.is_empty() {
+            return Ok(BTreeMap::new());
+        }
+
+        let txn = self
+            .inner
+            .transaction_on_one_with_mode(keys::DISPLAY_NAMES, IdbTransactionMode::Readonly)?;
+        let store = txn.object_store(keys::DISPLAY_NAMES)?;
+
+        let mut map = BTreeMap::new();
+        for display_name in display_names {
+            if let Some(user_ids) = store
+                .get(&self.encode_key(keys::DISPLAY_NAMES, (room_id, display_name)))?
+                .await?
+                .map(|f| self.deserialize_event::<BTreeSet<OwnedUserId>>(&f))
+                .transpose()?
+            {
+                map.insert(display_name.as_ref(), user_ids);
+            }
+        }
+
+        Ok(map)
+    }
+
     async fn get_account_data_event(
         &self,
         event_type: GlobalAccountDataEventType,

crates/matrix-sdk-sqlite/src/state_store.rs

@@ -696,15 +696,24 @@ trait SqliteObjectStateStoreExt: SqliteObjectExt {
             .optional()?)
     }
 
-    async fn get_display_name(&self, room_id: Key, name: Key) -> Result<Option<Vec<u8>>> {
+    async fn get_display_names(
+        &self,
+        room_id: Key,
+        names: Vec<Key>,
+    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
+        let sql_params = vec!["?"; names.len()].join(", ");
+        let sql = format!(
+            "SELECT name, data FROM display_name WHERE room_id = ? AND name IN ({sql_params})"
+        );
+        let params = iter::once(room_id).chain(names);
+
         Ok(self
-            .query_row(
-                "SELECT data FROM display_name WHERE room_id = ? AND name = ?",
-                (room_id, name),
-                |row| row.get(0),
-            )
-            .await
-            .optional()?)
+            .prepare(sql, move |mut stmt| {
+                stmt.query(rusqlite::params_from_iter(params))?
+                    .mapped(|row| Ok((row.get(0)?, row.get(1)?)))
+                    .collect()
+            })
+            .await?)
     }
 
     async fn get_user_receipt(
@@ -1312,17 +1321,52 @@ impl StateStore for SqliteStateStore {
         display_name: &str,
     ) -> Result<BTreeSet<OwnedUserId>> {
         let room_id = self.encode_key(keys::DISPLAY_NAME, room_id);
-        let name = self.encode_key(keys::DISPLAY_NAME, display_name);
+        let names = vec![self.encode_key(keys::DISPLAY_NAME, display_name)];
+
         Ok(self
             .acquire()
             .await?
-            .get_display_name(room_id, name)
+            .get_display_names(room_id, names)
             .await?
-            .map(|data| self.deserialize_json(&data))
+            .into_iter()
+            .next()
+            .map(|(_, data)| self.deserialize_json(&data))
             .transpose()?
             .unwrap_or_default())
     }
 
+    async fn get_users_with_display_names<'a>(
+        &self,
+        room_id: &RoomId,
+        display_names: &'a [String],
+    ) -> Result<BTreeMap<&'a str, BTreeSet<OwnedUserId>>> {
+        if display_names.is_empty() {
+            return Ok(BTreeMap::new());
+        }
+
+        let room_id = self.encode_key(keys::DISPLAY_NAME, room_id);
+        let mut names_map = display_names
+            .iter()
+            .map(|n| (self.encode_key(keys::DISPLAY_NAME, n), n.as_ref()))
+            .collect::<BTreeMap<_, _>>();
+        let names = names_map.keys().cloned().collect();
+
+        self.acquire()
+            .await?
+            .get_display_names(room_id, names)
+            .await?
+            .into_iter()
+            .map(|(name, data)| {
+                Ok((
+                    names_map
+                        .remove(name.as_slice())
+                        .expect("returned display names were requested"),
+                    self.deserialize_json(&data)?,
+                ))
+            })
+            .collect::<Result<BTreeMap<_, _>>>()
+    }
+
     async fn get_account_data_event(
         &self,
         event_type: GlobalAccountDataEventType,