Optimize get_shared_rooms()/intersection_sorted_stream2() for tighter loops.
Signed-off-by: Jason Volk <jason@zemos.net>
This commit is contained in:
@@ -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)| {
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
Reference in New Issue
Block a user