Stream::size_hint for mpsc channels (#2660)

Author: stepancheg

futures-channel/src/mpsc/mod.rs

@@ -1072,6 +1072,14 @@ impl<T> Stream for Receiver<T> {
             }
         }
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if let Some(inner) = &self.inner {
+            decode_state(inner.state.load(SeqCst)).size_hint()
+        } else {
+            (0, Some(0))
+        }
+    }
 }
 
 impl<T> Drop for Receiver<T> {
@@ -1216,6 +1224,14 @@ impl<T> Stream for UnboundedReceiver<T> {
             }
         }
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if let Some(inner) = &self.inner {
+            decode_state(inner.state.load(SeqCst)).size_hint()
+        } else {
+            (0, Some(0))
+        }
+    }
 }
 
 impl<T> Drop for UnboundedReceiver<T> {
@@ -1306,6 +1322,14 @@ impl State {
     fn is_closed(&self) -> bool {
         !self.is_open && self.num_messages == 0
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if self.is_open {
+            (self.num_messages, None)
+        } else {
+            (self.num_messages, Some(self.num_messages))
+        }
+    }
 }
 
 /*

futures-channel/tests/mpsc-size_hint.rs

@@ -0,0 +1,40 @@
+use futures::channel::mpsc;
+use futures::stream::Stream;
+
+#[test]
+fn unbounded_size_hint() {
+    let (tx, mut rx) = mpsc::unbounded::<u32>();
+    assert_eq!((0, None), rx.size_hint());
+    tx.unbounded_send(1).unwrap();
+    assert_eq!((1, None), rx.size_hint());
+    rx.try_next().unwrap().unwrap();
+    assert_eq!((0, None), rx.size_hint());
+    tx.unbounded_send(2).unwrap();
+    tx.unbounded_send(3).unwrap();
+    assert_eq!((2, None), rx.size_hint());
+    drop(tx);
+    assert_eq!((2, Some(2)), rx.size_hint());
+    rx.try_next().unwrap().unwrap();
+    assert_eq!((1, Some(1)), rx.size_hint());
+    rx.try_next().unwrap().unwrap();
+    assert_eq!((0, Some(0)), rx.size_hint());
+}
+
+#[test]
+fn channel_size_hint() {
+    let (mut tx, mut rx) = mpsc::channel::<u32>(10);
+    assert_eq!((0, None), rx.size_hint());
+    tx.try_send(1).unwrap();
+    assert_eq!((1, None), rx.size_hint());
+    rx.try_next().unwrap().unwrap();
+    assert_eq!((0, None), rx.size_hint());
+    tx.try_send(2).unwrap();
+    tx.try_send(3).unwrap();
+    assert_eq!((2, None), rx.size_hint());
+    drop(tx);
+    assert_eq!((2, Some(2)), rx.size_hint());
+    rx.try_next().unwrap().unwrap();
+    assert_eq!((1, Some(1)), rx.size_hint());
+    rx.try_next().unwrap().unwrap();
+    assert_eq!((0, Some(0)), rx.size_hint());
+}