Optimize get_shared_rooms()/intersection_sorted_stream2() for tighter loops.

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

src/core/utils/set.rs

@@ -1,12 +1,17 @@
 use std::{
 	cmp::{Eq, Ord},
+	convert::identity,
 	pin::Pin,
 	sync::Arc,
 };
 
-use futures::{Stream, StreamExt};
+use futures::{
+	Stream, StreamExt,
+	stream::{Peekable, unfold},
+};
+use tokio::sync::Mutex;
 
-use crate::{is_equal_to, is_less_than};
+use crate::{is_equal_to, is_less_than, utils::stream::ReadyExt};
 
 /// Intersection of sets
 ///
@@ -57,30 +62,31 @@ where
 /// Intersection of sets
 ///
 /// Outputs the set of elements common to both streams. Streams must be sorted.
-pub fn intersection_sorted_stream2<Item, A, B>(a: A, b: B) -> impl Stream<Item = Item> + Send
+pub fn intersection_sorted_stream2<S, Item>(a: S, b: S) -> impl Stream<Item = Item> + Send
 where
-	A: Stream<Item = Item> + Send,
-	B: Stream<Item = Item> + Send + Unpin,
+	S: Stream<Item = Item> + Send + Unpin,
 	Item: Eq + PartialOrd + Send + Sync,
 {
-	use tokio::sync::Mutex;
+	struct State<S: Stream> {
+		a: S,
+		b: Peekable<S>,
+	}
 
-	let b = Arc::new(Mutex::new(b.peekable()));
-	a.map(move |ai| (ai, b.clone()))
-		.filter_map(async move |(ai, b)| {
-			let mut lock = b.lock().await;
-			while let Some(bi) = Pin::new(&mut *lock)
+	unfold(State { a, b: b.peekable() }, async |mut state| {
+		let ai = state.a.next().await?;
+		while let Some(bi) = Pin::new(&mut state.b)
 			.next_if(|bi| *bi <= ai)
 			.await
 			.as_ref()
 		{
 			if ai == *bi {
-					return Some(ai);
+				return Some((Some(ai), state));
 			}
 		}
 
-			None
+		Some((None, state))
 	})
+	.ready_filter_map(identity)
 }
 
 /// Difference of sets
@@ -93,8 +99,6 @@ where
 	B: Stream<Item = Item> + Send + Unpin,
 	Item: Eq + PartialOrd + Send + Sync,
 {
-	use tokio::sync::Mutex;
-
 	let b = Arc::new(Mutex::new(b.peekable()));
 	a.map(move |ai| (ai, b.clone()))
 		.filter_map(async move |(ai, b)| {

src/service/rooms/state_cache/mod.rs

@@ -17,7 +17,7 @@ use tuwunel_core::{
 	result::LogErr,
 	trace,
 	utils::{
-		BoolExt,
+		self, BoolExt,
 		future::OptionStream,
 		stream::{BroadbandExt, ReadyExt, TryIgnore},
 	},
@@ -207,12 +207,10 @@ pub fn get_shared_rooms<'a>(
 	user_a: &'a UserId,
 	user_b: &'a UserId,
 ) -> impl Stream<Item = &RoomId> + Send + 'a {
-	use tuwunel_core::utils::set;
-
-	let a = self.rooms_joined(user_a);
+	let a = self.rooms_joined(user_a).boxed();
 	let b = self.rooms_joined(user_b).boxed();
 
-	set::intersection_sorted_stream2(a, b)
+	utils::set::intersection_sorted_stream2(a, b)
 }
 
 /// Returns an iterator of all joined members of a room.