Implement chain combinator for Buf

src/buf/buf.rs

@@ -1,4 +1,4 @@
-use super::{Take, Reader, Iter, FromBuf};
+use super::{IntoBuf, Take, Reader, Iter, FromBuf, Chain};
 use byteorder::ByteOrder;
 use iovec::IoVec;
 
@@ -522,6 +522,30 @@ pub trait Buf {
         super::take::new(self, limit)
     }
 
+    /// Creates an adaptor which will chain this buffer with another.
+    ///
+    /// The returned `Buf` instance will first consume all bytes from `self`.
+    /// Afterwards the output is equivalent to the output of next.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::{Bytes, Buf, IntoBuf};
+    /// use bytes::buf::Chain;
+    ///
+    /// let buf = Bytes::from(&b"hello "[..]).into_buf()
+    ///             .chain(Bytes::from(&b"world"[..]));
+    ///
+    /// let full: Bytes = buf.collect();
+    /// assert_eq!(full[..], b"hello world"[..]);
+    /// ```
+    fn chain<U>(self, next: U) -> Chain<Self, U::Buf>
+        where U: IntoBuf,
+              Self: Sized,
+    {
+        Chain::new(self, next.into_buf())
+    }
+
     /// Creates a "by reference" adaptor for this instance of `Buf`.
     ///
     /// The returned adaptor also implements `Buf` and will simply borrow `self`.

src/buf/chain.rs

@@ -0,0 +1,241 @@
+use {Buf, BufMut};
+use iovec::IoVec;
+
+/// A `Chain` sequences two buffers.
+///
+/// `Chain` is an adapter that links two underlying buffers and provides a
+/// continous view across both buffers. It is able to sequence either immutable
+/// buffers ([`Buf`] values) or mutable buffers ([`BufMut`] values).
+///
+/// This struct is generally created by calling [`Buf::chain`]. Please see that
+/// function's documentation for more detail.
+///
+/// # Examples
+///
+/// ```
+/// use bytes::{Bytes, Buf, IntoBuf};
+/// use bytes::buf::Chain;
+///
+/// let buf = Bytes::from(&b"hello "[..]).into_buf()
+///             .chain(Bytes::from(&b"world"[..]));
+///
+/// let full: Bytes = buf.collect();
+/// assert_eq!(full[..], b"hello world"[..]);
+/// ```
+///
+/// [`Buf::chain`]: trait.Buf.html#method.chain
+/// [`Buf`]: trait.Buf.html
+/// [`BufMut`]: trait.BufMut.html
+pub struct Chain<T, U> {
+    a: T,
+    b: U,
+}
+
+impl<T, U> Chain<T, U> {
+    /// Creates a new `Chain` sequencing the provided values.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::BytesMut;
+    /// use bytes::buf::Chain;
+    ///
+    /// let buf = Chain::new(
+    ///     BytesMut::with_capacity(1024),
+    ///     BytesMut::with_capacity(1024));
+    ///
+    /// // Use the chained buffer
+    /// ```
+    pub fn new(a: T, b: U) -> Chain<T, U> {
+        Chain {
+            a: a,
+            b: b,
+        }
+    }
+
+    /// Gets a reference to the first underlying `Buf`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::{Bytes, Buf, IntoBuf};
+    ///
+    /// let buf = Bytes::from(&b"hello"[..]).into_buf()
+    ///             .chain(Bytes::from(&b"world"[..]));
+    ///
+    /// assert_eq!(buf.first_ref().get_ref()[..], b"hello"[..]);
+    /// ```
+    pub fn first_ref(&self) -> &T {
+        &self.a
+    }
+
+    /// Gets a mutable reference to the first underlying `Buf`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::{Bytes, Buf, IntoBuf};
+    ///
+    /// let mut buf = Bytes::from(&b"hello "[..]).into_buf()
+    ///                 .chain(Bytes::from(&b"world"[..]));
+    ///
+    /// buf.first_mut().set_position(1);
+    ///
+    /// let full: Bytes = buf.collect();
+    /// assert_eq!(full[..], b"ello world"[..]);
+    /// ```
+    pub fn first_mut(&mut self) -> &mut T {
+        &mut self.a
+    }
+
+    /// Gets a reference to the last underlying `Buf`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::{Bytes, Buf, IntoBuf};
+    ///
+    /// let buf = Bytes::from(&b"hello"[..]).into_buf()
+    ///             .chain(Bytes::from(&b"world"[..]));
+    ///
+    /// assert_eq!(buf.last_ref().get_ref()[..], b"world"[..]);
+    /// ```
+    pub fn last_ref(&self) -> &U {
+        &self.b
+    }
+
+    /// Gets a mutable reference to the last underlying `Buf`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::{Bytes, Buf, IntoBuf};
+    ///
+    /// let mut buf = Bytes::from(&b"hello "[..]).into_buf()
+    ///                 .chain(Bytes::from(&b"world"[..]));
+    ///
+    /// buf.last_mut().set_position(1);
+    ///
+    /// let full: Bytes = buf.collect();
+    /// assert_eq!(full[..], b"hello orld"[..]);
+    /// ```
+    pub fn last_mut(&mut self) -> &mut U {
+        &mut self.b
+    }
+
+    /// Consumes this `Chain`, returning the underlying values.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::{Bytes, Buf, IntoBuf};
+    ///
+    /// let buf = Bytes::from(&b"hello"[..]).into_buf()
+    ///             .chain(Bytes::from(&b"world"[..]));
+    ///
+    /// let (first, last) = buf.into_inner();
+    /// assert_eq!(first.get_ref()[..], b"hello"[..]);
+    /// assert_eq!(last.get_ref()[..], b"world"[..]);
+    /// ```
+    pub fn into_inner(self) -> (T, U) {
+        (self.a, self.b)
+    }
+}
+
+impl<T, U> Buf for Chain<T, U>
+    where T: Buf,
+          U: Buf,
+{
+    fn remaining(&self) -> usize {
+        self.a.remaining() + self.b.remaining()
+    }
+
+    fn bytes(&self) -> &[u8] {
+        if self.a.has_remaining() {
+            self.a.bytes()
+        } else {
+            self.b.bytes()
+        }
+    }
+
+    fn advance(&mut self, mut cnt: usize) {
+        let a_rem = self.a.remaining();
+
+        if a_rem != 0 {
+            if a_rem >= cnt {
+                self.a.advance(cnt);
+                return;
+            }
+
+            // Consume what is left of a
+            self.a.advance(a_rem);
+
+            cnt -= a_rem;
+        }
+
+        self.b.advance(cnt);
+    }
+
+    fn bytes_vec<'a>(&'a self, dst: &mut [&'a IoVec]) -> usize {
+        if self.a.has_remaining() {
+            let mut n = self.a.bytes_vec(dst);
+
+            if n < dst.len() {
+                n += self.b.bytes_vec(&mut dst[n..]);
+            }
+
+            n
+        } else {
+            self.b.bytes_vec(dst)
+        }
+    }
+}
+
+impl<T, U> BufMut for Chain<T, U>
+    where T: BufMut,
+          U: BufMut,
+{
+    fn remaining_mut(&self) -> usize {
+        self.a.remaining_mut() + self.b.remaining_mut()
+    }
+
+    unsafe fn bytes_mut(&mut self) -> &mut [u8] {
+        if self.a.has_remaining_mut() {
+            self.a.bytes_mut()
+        } else {
+            self.b.bytes_mut()
+        }
+    }
+
+    unsafe fn advance_mut(&mut self, mut cnt: usize) {
+        let a_rem = self.a.remaining_mut();
+
+        if a_rem != 0 {
+            if a_rem >= cnt {
+                self.a.advance_mut(cnt);
+                return;
+            }
+
+            // Consume what is left of a
+            self.a.advance_mut(a_rem);
+
+            cnt -= a_rem;
+        }
+
+        self.b.advance_mut(cnt);
+    }
+
+    unsafe fn bytes_vec_mut<'a>(&'a mut self, dst: &mut [&'a mut IoVec]) -> usize {
+        if self.a.has_remaining_mut() {
+            let mut n = self.a.bytes_vec_mut(dst);
+
+            if n < dst.len() {
+                n += self.b.bytes_vec_mut(&mut dst[n..]);
+            }
+
+            n
+        } else {
+            self.b.bytes_vec_mut(dst)
+        }
+    }
+}

src/buf/mod.rs

@@ -19,6 +19,7 @@
 mod buf;
 mod buf_mut;
 mod from_buf;
+mod chain;
 mod into_buf;
 mod iter;
 mod reader;
@@ -29,6 +30,7 @@ mod writer;
 pub use self::buf::Buf;
 pub use self::buf_mut::BufMut;
 pub use self::from_buf::FromBuf;
+pub use self::chain::Chain;
 pub use self::into_buf::IntoBuf;
 pub use self::iter::Iter;
 pub use self::reader::Reader;

tests/test_chain.rs

@@ -0,0 +1,93 @@
+extern crate bytes;
+extern crate iovec;
+
+use bytes::{Buf, BufMut, Bytes, BytesMut};
+use bytes::buf::Chain;
+use iovec::IoVec;
+use std::io::Cursor;
+
+#[test]
+fn collect_two_bufs() {
+    let a = Cursor::new(Bytes::from(&b"hello"[..]));
+    let b = Cursor::new(Bytes::from(&b"world"[..]));
+
+    let res: Vec<u8> = a.chain(b).collect();
+    assert_eq!(res, &b"helloworld"[..]);
+}
+
+#[test]
+fn writing_chained() {
+    let mut a = BytesMut::with_capacity(64);
+    let mut b = BytesMut::with_capacity(64);
+
+    {
+        let mut buf = Chain::new(&mut a, &mut b);
+
+        for i in 0..128 {
+            buf.put(i as u8);
+        }
+    }
+
+    assert_eq!(64, a.len());
+    assert_eq!(64, b.len());
+
+    for i in 0..64 {
+        let expect = i as u8;
+        assert_eq!(expect, a[i]);
+        assert_eq!(expect + 64, b[i]);
+    }
+}
+
+#[test]
+fn iterating_two_bufs() {
+    let a = Cursor::new(Bytes::from(&b"hello"[..]));
+    let b = Cursor::new(Bytes::from(&b"world"[..]));
+
+    let res: Vec<u8> = a.chain(b).iter().collect();
+    assert_eq!(res, &b"helloworld"[..]);
+}
+
+#[test]
+fn vectored_read() {
+    let a = Cursor::new(Bytes::from(&b"hello"[..]));
+    let b = Cursor::new(Bytes::from(&b"world"[..]));
+
+    let mut buf = a.chain(b);
+    let mut iovecs: [&IoVec; 4] = Default::default();
+
+    assert_eq!(2, buf.bytes_vec(&mut iovecs));
+    assert_eq!(iovecs[0][..], b"hello"[..]);
+    assert_eq!(iovecs[1][..], b"world"[..]);
+    assert!(iovecs[2].is_empty());
+    assert!(iovecs[3].is_empty());
+
+    buf.advance(2);
+
+    iovecs = Default::default();
+
+    assert_eq!(2, buf.bytes_vec(&mut iovecs));
+    assert_eq!(iovecs[0][..], b"llo"[..]);
+    assert_eq!(iovecs[1][..], b"world"[..]);
+    assert!(iovecs[2].is_empty());
+    assert!(iovecs[3].is_empty());
+
+    buf.advance(3);
+
+    iovecs = Default::default();
+
+    assert_eq!(1, buf.bytes_vec(&mut iovecs));
+    assert_eq!(iovecs[0][..], b"world"[..]);
+    assert!(iovecs[1].is_empty());
+    assert!(iovecs[2].is_empty());
+    assert!(iovecs[3].is_empty());
+
+    buf.advance(3);
+
+    iovecs = Default::default();
+
+    assert_eq!(1, buf.bytes_vec(&mut iovecs));
+    assert_eq!(iovecs[0][..], b"ld"[..]);
+    assert!(iovecs[1].is_empty());
+    assert!(iovecs[2].is_empty());
+    assert!(iovecs[3].is_empty());
+}