From 287f6b9163674fbf63e0817105abc746d0ee980c Mon Sep 17 00:00:00 2001
From: mikoto <avdb@keemail.me>
Date: Tue, 17 Sep 2024 18:07:58 +0200
Subject: [PATCH] refactor calculate_invite_state

That was terribly named and terribly implemented.

Co-authored-by: Charles Hall <charles@computer.surgery>
---
 src/api/client_server/membership.rs |  2 +-
 src/service/rooms/state.rs          | 79 ++++++++++++-----------------
 src/service/rooms/timeline.rs       |  2 +-
 3 files changed, 34 insertions(+), 49 deletions(-)

diff --git a/src/api/client_server/membership.rs b/src/api/client_server/membership.rs
index 95aefcf4..00f1db9e 100644
--- a/src/api/client_server/membership.rs
+++ b/src/api/client_server/membership.rs
@@ -1307,7 +1307,7 @@ pub(crate) async fn invite_helper(
                 )?;
 
             let invite_room_state =
-                services().rooms.state.calculate_invite_state(&pdu)?;
+                services().rooms.state.get_helpful_invite_events(&pdu)?;
 
             drop(room_token);
 
diff --git a/src/service/rooms/state.rs b/src/service/rooms/state.rs
index 2ff970bb..f45effd7 100644
--- a/src/service/rooms/state.rs
+++ b/src/service/rooms/state.rs
@@ -1,6 +1,7 @@
 mod data;
 use std::{
     collections::{HashMap, HashSet},
+    iter,
     sync::Arc,
 };
 
@@ -266,58 +267,42 @@ impl Service {
         }
     }
 
+    /// Gather events to help the invited user identify the room
+    ///
+    /// Also includes the invite event itself.
     #[tracing::instrument(skip(self, invite_event))]
-    pub(crate) fn calculate_invite_state(
+    pub(crate) fn get_helpful_invite_events(
         &self,
         invite_event: &PduEvent,
     ) -> Result<Vec<Raw<AnyStrippedStateEvent>>> {
-        let mut state = Vec::new();
-        // Add recommended events
-        if let Some(e) = services().rooms.state_accessor.room_state_get(
-            &invite_event.room_id,
-            &StateEventType::RoomCreate,
-            "",
-        )? {
-            state.push(e.to_stripped_state_event());
-        }
-        if let Some(e) = services().rooms.state_accessor.room_state_get(
-            &invite_event.room_id,
-            &StateEventType::RoomJoinRules,
-            "",
-        )? {
-            state.push(e.to_stripped_state_event());
-        }
-        if let Some(e) = services().rooms.state_accessor.room_state_get(
-            &invite_event.room_id,
-            &StateEventType::RoomCanonicalAlias,
-            "",
-        )? {
-            state.push(e.to_stripped_state_event());
-        }
-        if let Some(e) = services().rooms.state_accessor.room_state_get(
-            &invite_event.room_id,
-            &StateEventType::RoomAvatar,
-            "",
-        )? {
-            state.push(e.to_stripped_state_event());
-        }
-        if let Some(e) = services().rooms.state_accessor.room_state_get(
-            &invite_event.room_id,
-            &StateEventType::RoomName,
-            "",
-        )? {
-            state.push(e.to_stripped_state_event());
-        }
-        if let Some(e) = services().rooms.state_accessor.room_state_get(
-            &invite_event.room_id,
-            &StateEventType::RoomMember,
-            invite_event.sender.as_str(),
-        )? {
-            state.push(e.to_stripped_state_event());
-        }
+        let helpful_events = [
+            (StateEventType::RoomCreate, ""),
+            (StateEventType::RoomJoinRules, ""),
+            (StateEventType::RoomCanonicalAlias, ""),
+            (StateEventType::RoomAvatar, ""),
+            (StateEventType::RoomName, ""),
+            (StateEventType::RoomMember, invite_event.sender.as_str()),
+        ];
 
-        state.push(invite_event.to_stripped_state_event());
-        Ok(state)
+        let helpful_events =
+            helpful_events.into_iter().filter_map(|(state_type, state_key)| {
+                let state = services().rooms.state_accessor.room_state_get(
+                    &invite_event.room_id,
+                    &state_type,
+                    state_key,
+                );
+
+                match state {
+                    Ok(Some(x)) => Some(Ok(x.to_stripped_state_event())),
+                    Err(x) => Some(Err(x)),
+                    Ok(None) => None,
+                }
+            });
+
+        let actual_event =
+            iter::once(Ok(invite_event.to_stripped_state_event()));
+
+        helpful_events.chain(actual_event).collect()
     }
 
     /// Set the state hash to a new version, but does not update state_cache.
diff --git a/src/service/rooms/timeline.rs b/src/service/rooms/timeline.rs
index 53802841..c70c7033 100644
--- a/src/service/rooms/timeline.rs
+++ b/src/service/rooms/timeline.rs
@@ -473,7 +473,7 @@ impl Service {
                             let state = services()
                                 .rooms
                                 .state
-                                .calculate_invite_state(pdu)?;
+                                .get_helpful_invite_events(pdu)?;
                             Some(state)
                         }
                         _ => None,