From 7abf0c8805b1d99d4c3a4af35c11f0f3193e303f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Damir=20Jeli=C4=87?= <poljar@termina.org.uk>
Date: Sat, 19 Dec 2020 14:44:46 +0100
Subject: [PATCH] store: Honor state keys for the state storage.

---
 matrix_sdk_base/src/client.rs | 14 +++++---
 matrix_sdk_base/src/store.rs  | 62 ++++++++++++++++++++---------------
 2 files changed, 45 insertions(+), 31 deletions(-)

diff --git a/matrix_sdk_base/src/client.rs b/matrix_sdk_base/src/client.rs
index 64901bca..3fe13aef 100644
--- a/matrix_sdk_base/src/client.rs
+++ b/matrix_sdk_base/src/client.rs
@@ -523,7 +523,7 @@ impl BaseClient {
     ) -> (
         InviteState,
         BTreeMap<UserId, StrippedMemberEvent>,
-        BTreeMap<String, AnyStrippedStateEvent>,
+        BTreeMap<String, BTreeMap<String, AnyStrippedStateEvent>>,
     ) {
         events.into_iter().fold(
             (InviteState::default(), BTreeMap::new(), BTreeMap::new()),
@@ -544,7 +544,10 @@ impl BaseClient {
                             }
                         } else {
                             room_info.handle_state_event(&e);
-                            state_events.insert(e.content().event_type().to_owned(), e);
+                            state_events
+                                .entry(e.content().event_type().to_owned())
+                                .or_insert_with(BTreeMap::new)
+                                .insert(e.state_key().to_owned(), e);
                         }
                     }
                     Err(err) => {
@@ -566,7 +569,7 @@ impl BaseClient {
     ) -> (
         State,
         BTreeMap<UserId, MemberEvent>,
-        BTreeMap<String, AnySyncStateEvent>,
+        BTreeMap<String, BTreeMap<String, AnySyncStateEvent>>,
         BTreeSet<UserId>,
     ) {
         let mut state = State::default();
@@ -610,7 +613,10 @@ impl BaseClient {
                     ),
                 }
             } else {
-                state_events.insert(event.content().event_type().to_owned(), event);
+                state_events
+                    .entry(event.content().event_type().to_owned())
+                    .or_insert_with(BTreeMap::new)
+                    .insert(event.state_key().to_owned(), event);
             }
         }
 
diff --git a/matrix_sdk_base/src/store.rs b/matrix_sdk_base/src/store.rs
index 94eea544..7489a6a4 100644
--- a/matrix_sdk_base/src/store.rs
+++ b/matrix_sdk_base/src/store.rs
@@ -42,11 +42,11 @@ pub struct StateChanges {
     pub presence: BTreeMap<UserId, PresenceEvent>,
 
     pub members: BTreeMap<RoomId, BTreeMap<UserId, MemberEvent>>,
-    pub state: BTreeMap<RoomId, BTreeMap<String, AnySyncStateEvent>>,
+    pub state: BTreeMap<RoomId, BTreeMap<String, BTreeMap<String, AnySyncStateEvent>>>,
     pub room_account_data: BTreeMap<RoomId, BTreeMap<String, AnyBasicEvent>>,
     pub room_infos: BTreeMap<RoomId, RoomInfo>,
 
-    pub stripped_state: BTreeMap<RoomId, BTreeMap<String, AnyStrippedStateEvent>>,
+    pub stripped_state: BTreeMap<RoomId, BTreeMap<String, BTreeMap<String, AnyStrippedStateEvent>>>,
     pub stripped_members: BTreeMap<RoomId, BTreeMap<UserId, StrippedMemberEvent>>,
     pub invited_room_info: BTreeMap<RoomId, RoomInfo>,
 }
@@ -77,6 +77,8 @@ impl StateChanges {
         self.stripped_state
             .entry(room_id.to_owned())
             .or_insert_with(BTreeMap::new)
+            .entry(event.content().event_type().to_string())
+            .or_insert_with(BTreeMap::new)
             .insert(event.state_key().to_string(), event);
     }
 
@@ -92,7 +94,9 @@ impl StateChanges {
         self.state
             .entry(room_id.to_owned())
             .or_insert_with(BTreeMap::new)
-            .insert(event.content().event_type().to_string(), event);
+            .entry(event.content().event_type().to_string())
+            .or_insert_with(BTreeMap::new)
+            .insert(event.state_key().to_string(), event);
     }
 }
 
@@ -242,18 +246,20 @@ impl Store {
                         }
                     }
 
-                    for (room, events) in &changes.state {
-                        for event in events.values() {
-                            state.insert(
-                                format!(
-                                    "{}{}{}",
-                                    room.as_str(),
-                                    event.content().event_type(),
-                                    event.state_key(),
-                                )
-                                .as_bytes(),
-                                serde_json::to_vec(&event).unwrap(),
-                            )?;
+                    for (room, event_types) in &changes.state {
+                        for events in event_types.values() {
+                            for event in events.values() {
+                                state.insert(
+                                    format!(
+                                        "{}{}{}",
+                                        room.as_str(),
+                                        event.content().event_type(),
+                                        event.state_key(),
+                                    )
+                                    .as_bytes(),
+                                    serde_json::to_vec(&event).unwrap(),
+                                )?;
+                            }
                         }
                     }
 
@@ -280,18 +286,20 @@ impl Store {
                         }
                     }
 
-                    for (room, events) in &changes.stripped_state {
-                        for event in events.values() {
-                            stripped_state.insert(
-                                format!(
-                                    "{}{}{}",
-                                    room.as_str(),
-                                    event.content().event_type(),
-                                    event.state_key(),
-                                )
-                                .as_bytes(),
-                                serde_json::to_vec(&event).unwrap(),
-                            )?;
+                    for (room, event_types) in &changes.stripped_state {
+                        for events in event_types.values() {
+                            for event in events.values() {
+                                stripped_state.insert(
+                                    format!(
+                                        "{}{}{}",
+                                        room.as_str(),
+                                        event.content().event_type(),
+                                        event.state_key(),
+                                    )
+                                    .as_bytes(),
+                                    serde_json::to_vec(&event).unwrap(),
+                                )?;
+                            }
                         }
                     }