mls-memory-storage: save messages in the messages_by_group cache

Co-authored-by: Yuki Kishimoto <[email protected]>
Signed-off-by: Yuki Kishimoto <[email protected]>

Author: erskingardner

crates/nostr-mls-memory-storage/src/messages.rs

@@ -9,8 +9,32 @@ use crate::NostrMlsMemoryStorage;
 
 impl MessageStorage for NostrMlsMemoryStorage {
     fn save_message(&self, message: Message) -> Result<(), MessageError> {
+        // Save in the messages cache
         let mut cache = self.messages_cache.write();
-        cache.put(message.id, message);
+        cache.put(message.id, message.clone());
+
+        // Save in the messages_by_group cache
+        let mut group_cache = self.messages_by_group_cache.write();
+        match group_cache.get_mut(&message.mls_group_id) {
+            Some(group_messages) => {
+                // TODO: time complexity here is O(n). We probably want to use another data struct here.
+
+                // Find and update existing message or add new one
+                match group_messages.iter().position(|m| m.id == message.id) {
+                    Some(idx) => {
+                        group_messages[idx] = message;
+                    }
+                    None => {
+                        group_messages.push(message);
+                    }
+                }
+            }
+            // Not found, insert new
+            None => {
+                group_cache.put(message.mls_group_id.clone(), vec![message]);
+            }
+        }
+
         Ok(())
     }