stream: impl FromStream for std::collections::* (#7966)

tokio-stream/src/stream_ext/collect.rs

@@ -6,7 +6,8 @@ use core::mem;
 use core::pin::Pin;
 use core::task::{ready, Context, Poll};
 use pin_project_lite::pin_project;
-use std::collections::BTreeSet;
+use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet, LinkedList, VecDeque};
+use std::hash::Hash;
 
 // Do not export this struct until `FromStream` can be unsealed.
 pin_project! {
@@ -136,6 +137,44 @@ impl<T> sealed::FromStreamPriv<T> for Vec<T> {
     }
 }
 
+impl<T> FromStream<T> for VecDeque<T> {}
+
+impl<T> sealed::FromStreamPriv<T> for VecDeque<T> {
+    type InternalCollection = VecDeque<T>;
+
+    fn initialize(_: sealed::Internal, lower: usize, _upper: Option<usize>) -> VecDeque<T> {
+        VecDeque::with_capacity(lower)
+    }
+
+    fn extend(_: sealed::Internal, collection: &mut VecDeque<T>, item: T) -> bool {
+        collection.push_back(item);
+        true
+    }
+
+    fn finalize(_: sealed::Internal, collection: &mut VecDeque<T>) -> VecDeque<T> {
+        mem::take(collection)
+    }
+}
+
+impl<T> FromStream<T> for LinkedList<T> {}
+
+impl<T> sealed::FromStreamPriv<T> for LinkedList<T> {
+    type InternalCollection = LinkedList<T>;
+
+    fn initialize(_: sealed::Internal, _lower: usize, _upper: Option<usize>) -> LinkedList<T> {
+        LinkedList::new()
+    }
+
+    fn extend(_: sealed::Internal, collection: &mut LinkedList<T>, item: T) -> bool {
+        collection.push_back(item);
+        true
+    }
+
+    fn finalize(_: sealed::Internal, collection: &mut LinkedList<T>) -> LinkedList<T> {
+        mem::take(collection)
+    }
+}
+
 impl<T: Ord> FromStream<T> for BTreeSet<T> {}
 
 impl<T: Ord> sealed::FromStreamPriv<T> for BTreeSet<T> {
@@ -155,6 +194,82 @@ impl<T: Ord> sealed::FromStreamPriv<T> for BTreeSet<T> {
     }
 }
 
+impl<K: Ord, V> FromStream<(K, V)> for BTreeMap<K, V> {}
+
+impl<K: Ord, V> sealed::FromStreamPriv<(K, V)> for BTreeMap<K, V> {
+    type InternalCollection = BTreeMap<K, V>;
+
+    fn initialize(_: sealed::Internal, _lower: usize, _upper: Option<usize>) -> BTreeMap<K, V> {
+        BTreeMap::new()
+    }
+
+    fn extend(_: sealed::Internal, collection: &mut BTreeMap<K, V>, (key, value): (K, V)) -> bool {
+        collection.insert(key, value);
+        true
+    }
+
+    fn finalize(_: sealed::Internal, collection: &mut BTreeMap<K, V>) -> BTreeMap<K, V> {
+        mem::take(collection)
+    }
+}
+
+impl<T: Eq + Hash> FromStream<T> for HashSet<T> {}
+
+impl<T: Eq + Hash> sealed::FromStreamPriv<T> for HashSet<T> {
+    type InternalCollection = HashSet<T>;
+
+    fn initialize(_: sealed::Internal, lower: usize, _upper: Option<usize>) -> HashSet<T> {
+        HashSet::with_capacity(lower)
+    }
+
+    fn extend(_: sealed::Internal, collection: &mut HashSet<T>, item: T) -> bool {
+        collection.insert(item);
+        true
+    }
+
+    fn finalize(_: sealed::Internal, collection: &mut HashSet<T>) -> HashSet<T> {
+        mem::take(collection)
+    }
+}
+
+impl<K: Eq + Hash, V> FromStream<(K, V)> for HashMap<K, V> {}
+
+impl<K: Eq + Hash, V> sealed::FromStreamPriv<(K, V)> for HashMap<K, V> {
+    type InternalCollection = HashMap<K, V>;
+
+    fn initialize(_: sealed::Internal, lower: usize, _upper: Option<usize>) -> HashMap<K, V> {
+        HashMap::with_capacity(lower)
+    }
+
+    fn extend(_: sealed::Internal, collection: &mut HashMap<K, V>, (key, value): (K, V)) -> bool {
+        collection.insert(key, value);
+        true
+    }
+
+    fn finalize(_: sealed::Internal, collection: &mut HashMap<K, V>) -> HashMap<K, V> {
+        mem::take(collection)
+    }
+}
+
+impl<T: Ord> FromStream<T> for BinaryHeap<T> {}
+
+impl<T: Ord> sealed::FromStreamPriv<T> for BinaryHeap<T> {
+    type InternalCollection = BinaryHeap<T>;
+
+    fn initialize(_: sealed::Internal, lower: usize, _upper: Option<usize>) -> BinaryHeap<T> {
+        BinaryHeap::with_capacity(lower)
+    }
+
+    fn extend(_: sealed::Internal, collection: &mut BinaryHeap<T>, item: T) -> bool {
+        collection.push(item);
+        true
+    }
+
+    fn finalize(_: sealed::Internal, collection: &mut BinaryHeap<T>) -> BinaryHeap<T> {
+        mem::take(collection)
+    }
+}
+
 impl<T> FromStream<T> for Box<[T]> {}
 
 impl<T> sealed::FromStreamPriv<T> for Box<[T]> {

tokio-stream/tests/stream_collect.rs

@@ -1,4 +1,4 @@
-use std::collections::BTreeSet;
+use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet, LinkedList, VecDeque};
 
 use tokio_stream::{self as stream, StreamExt};
 use tokio_test::{assert_pending, assert_ready, assert_ready_err, assert_ready_ok, task};
@@ -63,6 +63,50 @@ async fn collect_vec_items() {
     assert_eq!(vec![1, 2], coll);
 }
 
+#[tokio::test]
+async fn collect_vecdeque_items() {
+    let (tx, rx) = mpsc::unbounded_channel_stream();
+    let mut fut = task::spawn(rx.collect::<VecDeque<i32>>());
+
+    assert_pending!(fut.poll());
+
+    let xs = [1, 2, 42, 3];
+    for x in xs {
+        tx.send(x).unwrap();
+        assert!(fut.is_woken());
+        assert_pending!(fut.poll());
+    }
+
+    drop(tx);
+    assert!(fut.is_woken());
+    let coll = assert_ready!(fut.poll());
+
+    assert_eq!(coll, VecDeque::from(xs));
+    assert_eq!(coll.into_iter().collect::<Vec<_>>(), xs);
+}
+
+#[tokio::test]
+async fn collect_linkedlist_items() {
+    let (tx, rx) = mpsc::unbounded_channel_stream();
+    let mut fut = task::spawn(rx.collect::<LinkedList<i32>>());
+
+    assert_pending!(fut.poll());
+
+    let xs = [1, 2, 42, 3];
+    for x in xs {
+        tx.send(x).unwrap();
+        assert!(fut.is_woken());
+        assert_pending!(fut.poll());
+    }
+
+    drop(tx);
+    assert!(fut.is_woken());
+    let coll = assert_ready!(fut.poll());
+
+    assert_eq!(coll, LinkedList::from(xs));
+    assert_eq!(coll.into_iter().collect::<Vec<_>>(), xs);
+}
+
 #[tokio::test]
 async fn collect_btreeset_items() {
     let (tx, rx) = mpsc::unbounded_channel_stream();
@@ -84,6 +128,90 @@ async fn collect_btreeset_items() {
     assert_eq!(BTreeSet::from([1, 2]), coll);
 }
 
+#[tokio::test]
+async fn collect_btreemap_items() {
+    let (tx, rx) = mpsc::unbounded_channel_stream();
+    let mut fut = task::spawn(rx.collect::<BTreeMap<i32, i32>>());
+
+    assert_pending!(fut.poll());
+
+    tx.send((3, 4)).unwrap();
+    assert!(fut.is_woken());
+    assert_pending!(fut.poll());
+
+    tx.send((1, 2)).unwrap();
+    assert!(fut.is_woken());
+    assert_pending!(fut.poll());
+
+    drop(tx);
+    assert!(fut.is_woken());
+    let coll = assert_ready!(fut.poll());
+    assert_eq!(BTreeMap::from([(1, 2), (3, 4)]), coll);
+}
+
+#[tokio::test]
+async fn collect_hashset_items() {
+    let (tx, rx) = mpsc::unbounded_channel_stream();
+    let mut fut = task::spawn(rx.collect::<HashSet<i32>>());
+
+    assert_pending!(fut.poll());
+
+    tx.send(1).unwrap();
+    assert!(fut.is_woken());
+    assert_pending!(fut.poll());
+
+    tx.send(2).unwrap();
+    assert!(fut.is_woken());
+    assert_pending!(fut.poll());
+
+    drop(tx);
+    assert!(fut.is_woken());
+    let coll = assert_ready!(fut.poll());
+    assert_eq!(HashSet::from([1, 2]), coll);
+}
+
+#[tokio::test]
+async fn collect_hashmap_items() {
+    let (tx, rx) = mpsc::unbounded_channel_stream();
+    let mut fut = task::spawn(rx.collect::<HashMap<i32, i32>>());
+
+    assert_pending!(fut.poll());
+
+    tx.send((1, 2)).unwrap();
+    assert!(fut.is_woken());
+    assert_pending!(fut.poll());
+
+    tx.send((3, 4)).unwrap();
+    assert!(fut.is_woken());
+    assert_pending!(fut.poll());
+
+    drop(tx);
+    assert!(fut.is_woken());
+    let coll = assert_ready!(fut.poll());
+    assert_eq!(HashMap::from([(1, 2), (3, 4)]), coll);
+}
+
+#[tokio::test]
+async fn collect_binaryheap_items() {
+    let (tx, rx) = mpsc::unbounded_channel_stream();
+    let mut fut = task::spawn(rx.collect::<BinaryHeap<i32>>());
+
+    assert_pending!(fut.poll());
+
+    tx.send(2).unwrap();
+    assert!(fut.is_woken());
+    assert_pending!(fut.poll());
+
+    tx.send(1).unwrap();
+    assert!(fut.is_woken());
+    assert_pending!(fut.poll());
+
+    drop(tx);
+    assert!(fut.is_woken());
+    let coll = assert_ready!(fut.poll());
+    assert_eq!(vec![1, 2], coll.into_sorted_vec());
+}
+
 #[tokio::test]
 async fn collect_string_items() {
     let (tx, rx) = mpsc::unbounded_channel_stream();