flatten state_full_shortids

Signed-off-by: Jason Volk <jason@zemos.net>

src/service/rooms/state/mod.rs

@@ -429,13 +429,14 @@ impl Service {
 		sender: &UserId,
 		state_key: Option<&str>,
 		content: &serde_json::value::RawValue,
-	) -> Result<StateMap<Arc<PduEvent>>> {
+	) -> Result<StateMap<PduEvent>> {
 		let Ok(shortstatehash) = self.get_room_shortstatehash(room_id).await else {
 			return Ok(HashMap::new());
 		};
 
-		let mut sauthevents: HashMap<_, _> =
+		let auth_types = state_res::auth_types_for_event(kind, sender, state_key, content)?;
-			state_res::auth_types_for_event(kind, sender, state_key, content)?
+
+		let sauthevents: HashMap<_, _> = auth_types
 			.iter()
 			.stream()
 			.broad_filter_map(|(event_type, state_key)| {
@@ -445,9 +446,6 @@ impl Service {
 					.map_ok(move |ssk| (ssk, (event_type, state_key)))
 					.map(Result::ok)
 			})
-				.map(|(ssk, (event_type, state_key))| {
-					(ssk, (event_type.to_owned(), state_key.to_owned()))
-				})
 			.collect()
 			.await;
 
@@ -455,34 +453,30 @@ impl Service {
 			.services
 			.state_accessor
 			.state_full_shortids(shortstatehash)
-			.await
+			.ready_filter_map(Result::ok)
-			.map_err(|e| err!(Database(error!(?room_id, ?shortstatehash, "{e:?}"))))?
+			.ready_filter_map(|(shortstatekey, shorteventid)| {
-			.into_iter()
-			.filter_map(|(shortstatekey, shorteventid)| {
 				sauthevents
-					.remove(&shortstatekey)
+					.get(&shortstatekey)
-					.map(|(event_type, state_key)| ((event_type, state_key), shorteventid))
+					.map(|(ty, sk)| ((ty, sk), shorteventid))
 			})
-			.unzip();
+			.unzip()
+			.await;
 
-		let auth_pdus = self
+		self.services
-			.services
 			.short
 			.multi_get_eventid_from_short(event_ids.into_iter().stream())
 			.zip(state_keys.into_iter().stream())
-			.ready_filter_map(|(event_id, tsk)| Some((tsk, event_id.ok()?)))
+			.ready_filter_map(|(event_id, (ty, sk))| Some(((ty, sk), event_id.ok()?)))
-			.broad_filter_map(|(tsk, event_id): (_, OwnedEventId)| async move {
+			.broad_filter_map(|((ty, sk), event_id): (_, OwnedEventId)| async move {
 				self.services
 					.timeline
 					.get_pdu(&event_id)
 					.await
-					.map(Arc::new)
+					.map(move |pdu| (((*ty).clone(), (*sk).clone()), pdu))
-					.map(move |pdu| (tsk, pdu))
 					.ok()
 			})
 			.collect()
-			.await;
+			.map(Ok)
-
+			.await
-		Ok(auth_pdus)
 	}
 }

src/service/rooms/state_accessor/mod.rs

@@ -1,6 +1,7 @@
 use std::{
 	borrow::Borrow,
 	fmt::Write,
+	ops::Deref,
 	sync::{Arc, Mutex as StdMutex, Mutex},
 };
 
@@ -10,8 +11,7 @@ use conduwuit::{
 	utils,
 	utils::{
 		math::{usize_from_f64, Expected},
-		stream::BroadbandExt,
+		stream::{BroadbandExt, IterStream, ReadyExt, TryExpect},
-		IterStream, ReadyExt,
 	},
 	Err, Error, PduEvent, Result,
 };
@@ -158,12 +158,8 @@ impl Service {
 	) -> impl Stream<Item = PduEvent> + Send + '_ {
 		let short_ids = self
 			.state_full_shortids(shortstatehash)
-			.map(|result| result.expect("missing shortstatehash"))
+			.expect_ok()
-			.map(Vec::into_iter)
+			.map(at!(1));
-			.map(|iter| iter.map(at!(1)))
-			.map(IterStream::stream)
-			.flatten_stream()
-			.boxed();
 
 		self.services
 			.short
@@ -187,9 +183,8 @@ impl Service {
 	{
 		let shortids = self
 			.state_full_shortids(shortstatehash)
-			.map(|result| result.expect("missing shortstatehash"))
+			.expect_ok()
-			.map(|vec| vec.into_iter().unzip())
+			.unzip()
-			.boxed()
 			.shared();
 
 		let shortstatekeys = shortids
@@ -255,25 +250,25 @@ impl Service {
 	}
 
 	#[inline]
-	pub async fn state_full_shortids(
+	pub fn state_full_shortids(
 		&self,
 		shortstatehash: ShortStateHash,
-	) -> Result<Vec<(ShortStateKey, ShortEventId)>> {
+	) -> impl Stream<Item = Result<(ShortStateKey, ShortEventId)>> + Send + '_ {
-		let shortids = self
+		self.services
-			.services
 			.state_compressor
 			.load_shortstatehash_info(shortstatehash)
-			.await
+			.map_err(|e| err!(Database("Missing state IDs: {e}")))
-			.map_err(|e| err!(Database("Missing state IDs: {e}")))?
+			.map_ok(|vec| vec.last().expect("at least one layer").full_state.clone())
-			.pop()
+			.map_ok(|full_state| {
-			.expect("there is always one layer")
+				full_state
-			.full_state
+					.deref()
 					.iter()
 					.copied()
 					.map(parse_compressed_state_event)
-			.collect();
+					.collect()
-
+			})
-		Ok(shortids)
+			.map_ok(|vec: Vec<_>| vec.into_iter().try_stream())
+			.try_flatten_stream()
 	}
 
 	/// Returns a single PDU from `room_id` with key (`event_type`,