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

Signed-off-by: Jason Volk <jason@zemos.net>
This commit is contained in:
Jason Volk
2026-02-27 01:54:24 +00:00
parent e9ca105d79
commit d959dd740f
2 changed files with 29 additions and 27 deletions

View File

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

View File

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