Repository: alexcrichton/futures-rs
Pull request #544 (2017-07-28 21:09:17)

Add `reunite` for BiLock and Split{Stream,Sink}
In both cases, consumes two halves of a split-up value and recovers the original.

```rust
let foo: Foo = Foo();
let (sink, stream) = foo.split();
let foo: Foo = sink.reunite(stream).expect("sink and stream aren't from the same original Stream + Sink");
```


Patch:
diff --git a/.travis.yml b/.travis.yml
index bbd15e98..cd932ad4 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -12,7 +12,7 @@ matrix:
       after_success:
         - travis-cargo doc-upload
     - os: linux
-      rust: 1.10.0
+      rust: 1.17.0
       script: cargo test
 sudo: false
 script:
diff --git a/src/sink/mod.rs b/src/sink/mod.rs
index 37e8cf73..4e31b979 100644
--- a/src/sink/mod.rs
+++ b/src/sink/mod.rs
@@ -12,7 +12,6 @@ use {IntoFuture, Poll, StartSend};
 use stream::Stream;
 
 mod with;
-mod with_flat_map;
 // mod with_map;
 // mod with_filter;
 // mod with_filter_map;
@@ -74,7 +73,6 @@ if_std! {
 }
 
 pub use self::with::With;
-pub use self::with_flat_map::WithFlatMap;
 pub use self::flush::Flush;
 pub use self::send::Send;
 pub use self::send_all::SendAll;
@@ -317,44 +315,6 @@ pub trait Sink {
         with::new(self, f)
     }
 
-    /// Composes a function *in front of* the sink.
-    ///
-    /// This adapter produces a new sink that passes each value through the
-    /// given function `f` before sending it to `self`.
-    ///
-    /// To process each value, `f` produces a *stream*, of which each value
-    /// is passed to the underlying sink. A new value will not be accepted until
-    /// the stream has been drained
-    ///
-    /// Note that this function consumes the given sink, returning a wrapped
-    /// version, much like `Iterator::flat_map`.
-    ///
-    /// # Examples
-    /// ---
-    /// Using this function with an iterator through use of the `stream::iter()`
-    /// function
-    ///
-    /// ```
-    /// use futures::{Sink,Future,Stream};
-    /// use futures::stream;
-    /// use futures::sync::mpsc;
-    /// 
-    /// let (tx, rx) = mpsc::channel::<i32>(5);
-    ///
-    /// let tx = tx.with_flat_map(|x| {
-    ///     stream::iter(vec![42; x].into_iter().map(|y|Ok(y)))
-    /// });
-    /// tx.send(5).wait().unwrap();
-    /// assert_eq!(rx.collect().wait(), Ok(vec![42, 42, 42, 42, 42]))
-    /// ```
-    fn with_flat_map<U, F, St>(self, f: F) -> WithFlatMap<Self, U, F, St>
-        where F: FnMut(U) -> St,
-              St: Stream<Item = Self::SinkItem, Error=Self::SinkError>,
-              Self: Sized
-        {
-            with_flat_map::new(self, f)
-        }
-
     /*
     fn with_map<U, F>(self, f: F) -> WithMap<Self, U, F>
         where F: FnMut(U) -> Self::SinkItem,
diff --git a/src/sink/with_flat_map.rs b/src/sink/with_flat_map.rs
deleted file mode 100644
index ec28358b..00000000
--- a/src/sink/with_flat_map.rs
+++ /dev/null
@@ -1,126 +0,0 @@
-use core::marker::PhantomData;
-
-use {Poll, Async, StartSend, AsyncSink};
-use sink::Sink;
-use stream::Stream;
-
-/// Sink for the `Sink::with_flat_map` combinator, chaining a computation that returns an iterator
-/// to run prior to pushing a value into the underlying sink
-#[derive(Debug)]
-#[must_use = "sinks do nothing unless polled"]
-pub struct WithFlatMap<S, U, F, St>
-where
-    S: Sink,
-    F: FnMut(U) -> St,
-    St: Stream<Item = S::SinkItem, Error=S::SinkError>,
-{
-    sink: S,
-    f: F,
-    stream: Option<St>,
-    buffer: Option<S::SinkItem>,
-    _phantom: PhantomData<fn(U)>,
-}
-
-pub fn new<S, U, F, St>(sink: S, f: F) -> WithFlatMap<S, U, F, St>
-where
-    S: Sink,
-    F: FnMut(U) -> St,
-    St: Stream<Item = S::SinkItem, Error=S::SinkError>,
-{
-    WithFlatMap {
-        sink: sink,
-        f: f,
-        stream: None,
-        buffer: None,
-        _phantom: PhantomData,
-    }
-}
-
-impl<S, U, F, St> WithFlatMap<S, U, F, St>
-where
-    S: Sink,
-    F: FnMut(U) -> St,
-    St: Stream<Item = S::SinkItem, Error=S::SinkError>,
-{
-    /// Get a shared reference to the inner sink.
-    pub fn get_ref(&self) -> &S {
-        &self.sink
-    }
-
-    /// Get a mutable reference to the inner sink.
-    pub fn get_mut(&mut self) -> &mut S {
-        &mut self.sink
-    }
-
-    /// Consumes this combinator, returning the underlying sink.
-    ///
-    /// Note that this may discard intermediate state of this combinator, so
-    /// care should be taken to avoid losing resources when this is called.
-    pub fn into_inner(self) -> S {
-        self.sink
-    }
-
-    fn try_empty_stream(&mut self) -> Poll<(), S::SinkError> {
-        if let Some(x) = self.buffer.take() {
-            if let AsyncSink::NotReady(x) = try!(self.sink.start_send(x)) {
-                self.buffer = Some(x);
-                return Ok(Async::NotReady);
-            }
-        }
-        if let Some(mut stream) = self.stream.take() {
-            while let Some(x) = try_ready!(stream.poll()) {
-                if let AsyncSink::NotReady(x) = try!(self.sink.start_send(x)) {
-                    self.stream = Some(stream);
-                    self.buffer = Some(x);
-                    return Ok(Async::NotReady);
-                }
-            }
-        }
-        Ok(Async::Ready(()))
-    }
-}
-
-impl<S, U, F, St> Stream for WithFlatMap<S, U, F, St>
-where
-    S: Stream + Sink,
-    F: FnMut(U) -> St,
-    St: Stream<Item = S::SinkItem, Error=S::SinkError>,
-{
-    type Item = S::Item;
-    type Error = S::Error;
-    fn poll(&mut self) -> Poll<Option<S::Item>, S::Error> {
-        self.sink.poll()
-    }
-}
-
-impl<S, U, F, St> Sink for WithFlatMap<S, U, F, St>
-where
-    S: Sink,
-    F: FnMut(U) -> St,
-    St: Stream<Item = S::SinkItem, Error=S::SinkError>,
-{
-    type SinkItem = U;
-    type SinkError = S::SinkError;
-    fn start_send(&mut self, i: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
-        if try!(self.try_empty_stream()).is_not_ready() {
-            return Ok(AsyncSink::NotReady(i));
-        }
-        assert!(self.stream.is_none());
-        self.stream = Some((self.f)(i));
-        try!(self.try_empty_stream());
-        Ok(AsyncSink::Ready)
-    }
-    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
-        if try!(self.try_empty_stream()).is_not_ready() {
-            return Ok(Async::NotReady);
-        }
-        self.sink.poll_complete()
-    }
-    fn close(&mut self) -> Poll<(), Self::SinkError> {
-        if try!(self.try_empty_stream()).is_not_ready() {
-            return Ok(Async::NotReady);
-        }
-        assert!(self.stream.is_none());
-        self.sink.close()
-    }
-}
diff --git a/src/stream/split.rs b/src/stream/split.rs
index 6e53b978..edfa81b1 100644
--- a/src/stream/split.rs
+++ b/src/stream/split.rs
@@ -1,10 +1,21 @@
 use {StartSend, Sink, Stream, Poll, Async, AsyncSink};
+use std::error::Error;
+use std::fmt;
 use sync::BiLock;
 
 /// A `Stream` part of the split pair
 #[derive(Debug)]
 pub struct SplitStream<S>(BiLock<S>);
 
+impl<S> SplitStream<S> {
+    /// Attempts to put the two "halves" of a split `Stream + Sink` back
+    /// together. Succeeds only if the `SplitStream<S>` and `SplitSink<S>` are
+    /// a matching pair originating from the same call to `Stream::split`.
+    pub fn reunite(self, other: SplitSink<S>) -> Result<S, ReuniteError<S>> {
+        other.reunite(self)
+    }
+}
+
 impl<S: Stream> Stream for SplitStream<S> {
     type Item = S::Item;
     type Error = S::Error;
@@ -21,6 +32,17 @@ impl<S: Stream> Stream for SplitStream<S> {
 #[derive(Debug)]
 pub struct SplitSink<S>(BiLock<S>);
 
+impl<S> SplitSink<S> {
+    /// Attempts to put the two "halves" of a split `Stream + Sink` back
+    /// together. Succeeds only if the `SplitStream<S>` and `SplitSink<S>` are
+    /// a matching pair originating from the same call to `Stream::split`.
+    pub fn reunite(self, other: SplitStream<S>) -> Result<S, ReuniteError<S>> {
+        self.0.reunite(other.0).map_err(|err| {
+            ReuniteError(SplitSink(err.0), SplitStream(err.1))
+        })
+    }
+}
+
 impl<S: Sink> Sink for SplitSink<S> {
     type SinkItem = S::SinkItem;
     type SinkError = S::SinkError;
@@ -55,3 +77,27 @@ pub fn split<S: Stream + Sink>(s: S) -> (SplitSink<S>, SplitStream<S>) {
     let write = SplitSink(b);
     (write, read)
 }
+
+/// Error indicating a `SplitSink<S>` and `SplitStream<S>` were not two halves
+/// of a `Stream + Split`, and thus could not be `reunite`d.
+pub struct ReuniteError<T>(pub SplitSink<T>, pub SplitStream<T>);
+
+impl<T> fmt::Debug for ReuniteError<T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        fmt.debug_tuple("ReuniteError")
+            .field(&"...")
+            .finish()
+    }
+}
+
+impl<T> fmt::Display for ReuniteError<T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        write!(fmt, "tried to reunite a SplitStream and SplitSink that don't form a pair")
+    }
+}
+
+impl<T> Error for ReuniteError<T> {
+    fn description(&self) -> &str {
+        "tried to reunite a SplitStream and SplitSink that don't form a pair"
+    }
+}
diff --git a/src/sync/bilock.rs b/src/sync/bilock.rs
index 06b39fe6..62a6178b 100644
--- a/src/sync/bilock.rs
+++ b/src/sync/bilock.rs
@@ -1,5 +1,7 @@
 use std::boxed::Box;
 use std::cell::UnsafeCell;
+use std::error::Error;
+use std::fmt;
 use std::mem;
 use std::ops::{Deref, DerefMut};
 use std::sync::Arc;
@@ -35,7 +37,7 @@ pub struct BiLock<T> {
 #[derive(Debug)]
 struct Inner<T> {
     state: AtomicUsize,
-    inner: UnsafeCell<T>,
+    inner: Option<UnsafeCell<T>>,
 }
 
 unsafe impl<T: Send> Send for Inner<T> {}
@@ -50,7 +52,7 @@ impl<T> BiLock<T> {
     pub fn new(t: T) -> (BiLock<T>, BiLock<T>) {
         let inner = Arc::new(Inner {
             state: AtomicUsize::new(0),
-            inner: UnsafeCell::new(t),
+            inner: Some(UnsafeCell::new(t)),
         });
 
         (BiLock { inner: inner.clone() }, BiLock { inner: inner })
@@ -131,6 +133,21 @@ impl<T> BiLock<T> {
         }
     }
 
+    /// Attempts to put the two "halves" of a `BiLock<T>` back together and
+    /// recover the original value. Succeeds only if the two `BiLock<T>`s
+    /// originated from the same call to `BiLock::new`.
+    pub fn reunite(self, other: Self) -> Result<T, ReuniteError<T>> {
+        if Arc::ptr_eq(&self.inner, &other.inner) {
+            drop(other);
+            let inner = Arc::try_unwrap(self.inner)
+                .ok()
+                .expect("futures: try_unwrap failed in BiLock<T>::reunite");
+            Ok(unsafe { inner.into_inner() })
+        } else {
+            Err(ReuniteError(self, other))
+        }
+    }
+
     fn unlock(&self) {
         match self.inner.state.swap(0, SeqCst) {
             // we've locked the lock, shouldn't be possible for us to see an
@@ -149,12 +166,42 @@ impl<T> BiLock<T> {
     }
 }
 
+impl<T> Inner<T> {
+    unsafe fn into_inner(mut self) -> T {
+        mem::replace(&mut self.inner, None).unwrap().into_inner()
+    }
+}
+
 impl<T> Drop for Inner<T> {
     fn drop(&mut self) {
         assert_eq!(self.state.load(SeqCst), 0);
     }
 }
 
+/// Error indicating two `BiLock<T>`s were not two halves of a whole, and
+/// thus could not be `reunite`d.
+pub struct ReuniteError<T>(pub BiLock<T>, pub BiLock<T>);
+
+impl<T> fmt::Debug for ReuniteError<T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        fmt.debug_tuple("ReuniteError")
+            .field(&"...")
+            .finish()
+    }
+}
+
+impl<T> fmt::Display for ReuniteError<T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        write!(fmt, "tried to reunite two BiLocks that don't form a pair")
+    }
+}
+
+impl<T> Error for ReuniteError<T> {
+    fn description(&self) -> &str {
+        "tried to reunite two BiLocks that don't form a pair"
+    }
+}
+
 /// Returned RAII guard from the `poll_lock` method.
 ///
 /// This structure acts as a sentinel to the data in the `BiLock<T>` itself,
@@ -168,13 +215,13 @@ pub struct BiLockGuard<'a, T: 'a> {
 impl<'a, T> Deref for BiLockGuard<'a, T> {
     type Target = T;
     fn deref(&self) -> &T {
-        unsafe { &*self.inner.inner.inner.get() }
+        unsafe { &*self.inner.inner.inner.as_ref().unwrap().get() }
     }
 }
 
 impl<'a, T> DerefMut for BiLockGuard<'a, T> {
     fn deref_mut(&mut self) -> &mut T {
-        unsafe { &mut *self.inner.inner.inner.get() }
+        unsafe { &mut *self.inner.inner.inner.as_ref().unwrap().get() }
     }
 }
 
@@ -231,13 +278,13 @@ impl<T> BiLockAcquired<T> {
 impl<T> Deref for BiLockAcquired<T> {
     type Target = T;
     fn deref(&self) -> &T {
-        unsafe { &*self.inner.as_ref().unwrap().inner.inner.get() }
+        unsafe { &*self.inner.as_ref().unwrap().inner.inner.as_ref().unwrap().get() }
     }
 }
 
 impl<T> DerefMut for BiLockAcquired<T> {
     fn deref_mut(&mut self) -> &mut T {
-        unsafe { &mut *self.inner.as_mut().unwrap().inner.inner.get() }
+        unsafe { &mut *self.inner.as_mut().unwrap().inner.inner.as_ref().unwrap().get() }
     }
 }
 


Test patch:
diff --git a/tests/bilock.rs b/tests/bilock.rs
index 0458779a..bc1bafa5 100644
--- a/tests/bilock.rs
+++ b/tests/bilock.rs
@@ -15,25 +15,31 @@ use support::*;
 fn smoke() {
     let future = future::lazy(|| {
         let (a, b) = BiLock::new(1);
-        let mut lock = match a.poll_lock() {
-            Async::Ready(l) => l,
-            Async::NotReady => panic!("poll not ready"),
-        };
-        assert_eq!(*lock, 1);
-        *lock = 2;
 
-        assert!(b.poll_lock().is_not_ready());
-        assert!(a.poll_lock().is_not_ready());
-        drop(lock);
+        {
+            let mut lock = match a.poll_lock() {
+                Async::Ready(l) => l,
+                Async::NotReady => panic!("poll not ready"),
+            };
+            assert_eq!(*lock, 1);
+            *lock = 2;
+
+            assert!(b.poll_lock().is_not_ready());
+            assert!(a.poll_lock().is_not_ready());
+        }
 
         assert!(b.poll_lock().is_ready());
         assert!(a.poll_lock().is_ready());
 
-        let lock = match b.poll_lock() {
-            Async::Ready(l) => l,
-            Async::NotReady => panic!("poll not ready"),
-        };
-        assert_eq!(*lock, 2);
+        {
+            let lock = match b.poll_lock() {
+                Async::Ready(l) => l,
+                Async::NotReady => panic!("poll not ready"),
+            };
+            assert_eq!(*lock, 2);
+        }
+
+        assert_eq!(a.reunite(b).expect("bilock/smoke: reunite error"), 2);
 
         Ok::<(), ()>(())
     });
@@ -73,6 +79,8 @@ fn concurrent() {
         Async::NotReady => panic!("poll not ready"),
     }
 
+    assert_eq!(a.reunite(b).expect("bilock/concurrent: reunite error"), 2 * N);
+
     struct Increment {
         remaining: usize,
         a: Option<BiLock<usize>>,
diff --git a/tests/sink.rs b/tests/sink.rs
index ad34fbdb..54a4a0e1 100644
--- a/tests/sink.rs
+++ b/tests/sink.rs
@@ -167,19 +167,6 @@ fn with_as_map() {
     assert_eq!(sink.get_ref(), &[0, 2, 4]);
 }
 
-#[test]
-// test simple use of with_flat_map
-fn with_flat_map() {
-    let sink = Vec::new().with_flat_map(|item| {
-        stream::iter(vec!(item; item).into_iter().map(Ok))
-    });
-    let sink = sink.send(0).wait().unwrap();
-    let sink = sink.send(1).wait().unwrap();
-    let sink = sink.send(2).wait().unwrap();
-    let sink = sink.send(3).wait().unwrap();
-    assert_eq!(sink.get_ref(), &[1,2,2,3,3,3]);
-}
-
 // Immediately accepts all requests to start pushing, but completion is managed
 // by manually flushing
 struct ManualFlush<T> {
diff --git a/tests/split.rs b/tests/split.rs
index b9309c7f..e00c72fd 100644
--- a/tests/split.rs
+++ b/tests/split.rs
@@ -39,6 +39,8 @@ fn test_split() {
     {
         let j = Join(iter(vec![Ok(10), Ok(20), Ok(30)]), &mut dest);
         let (sink, stream) = j.split();
+        let j = sink.reunite(stream).expect("test_split: reunite error");
+        let (sink, stream) = j.split();
         sink.send_all(stream).wait().unwrap();
     }
     assert_eq!(dest, vec![10, 20, 30]);