Provide Take decorator that limits Buf size

src/buf/mod.rs

@@ -3,10 +3,12 @@ mod ring;
 mod sink;
 mod slice;
 mod source;
+mod take;
 
 pub use self::byte::{ByteBuf, MutByteBuf, ROByteBuf};
 pub use self::ring::RingBuf;
 pub use self::slice::{SliceBuf, MutSliceBuf};
+pub use self::take::Take;
 
 use {BufError, RopeBuf};
 use std::{cmp, fmt, io, ptr, usize};

src/buf/take.rs

@@ -0,0 +1,79 @@
+use buf::{Buf, MutBuf};
+use std::{cmp, io};
+
+#[derive(Debug)]
+pub struct Take<T> {
+    inner: T,
+    limit: usize,
+}
+
+impl<T> Take<T> {
+    pub fn new(inner: T, limit: usize) -> Take<T> {
+        Take {
+            inner: inner,
+            limit: limit,
+        }
+    }
+
+    pub fn into_inner(self) -> T {
+        self.inner
+    }
+
+    pub fn get_ref(&self) -> &T {
+        &self.inner
+    }
+
+    pub fn get_mut(&mut self) -> &mut T {
+        &mut self.inner
+    }
+
+    pub fn limit(&self) -> usize {
+        self.limit
+    }
+
+    pub fn set_limit(&mut self, lim: usize) {
+        self.limit = lim
+    }
+}
+
+impl<T: Buf> Buf for Take<T> {
+    fn remaining(&self) -> usize {
+        cmp::min(self.inner.remaining(), self.limit)
+    }
+
+    fn bytes<'a>(&'a self) -> &'a [u8] {
+        &self.inner.bytes()[..self.limit]
+    }
+
+    fn advance(&mut self, cnt: usize) {
+        let cnt = cmp::min(cnt, self.limit);
+        self.limit -= cnt;
+        self.inner.advance(cnt);
+    }
+}
+
+impl<T: Buf> io::Read for Take<T> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        if !self.has_remaining() {
+            return Ok(0);
+        }
+
+        Ok(self.read_slice(buf))
+    }
+}
+
+impl<T: MutBuf> MutBuf for Take<T> {
+    fn remaining(&self) -> usize {
+        cmp::min(self.inner.remaining(), self.limit)
+    }
+
+    fn mut_bytes<'a>(&'a mut self) -> &'a mut [u8] {
+        &mut self.inner.mut_bytes()[..self.limit]
+    }
+
+    fn advance(&mut self, cnt: usize) {
+        let cnt = cmp::min(cnt, self.limit);
+        self.limit -= cnt;
+        self.inner.advance(cnt);
+    }
+}

src/lib.rs

@@ -17,7 +17,8 @@ pub use buf::{
     SliceBuf,
     MutSliceBuf,
     Source,
-    Sink
+    Sink,
+    Take,
 };
 pub use str::{
     ByteStr,

test/test.rs

@@ -5,6 +5,7 @@ extern crate rand;
 
 mod test_buf;
 mod test_buf_fill;
+mod test_buf_take;
 mod test_byte_buf;
 mod test_bytes;
 mod test_ring;

test/test_buf_take.rs

@@ -0,0 +1,12 @@
+use bytes::*;
+use std::io::{Cursor, Read};
+
+#[test]
+pub fn test_take_from_buf() {
+    let mut buf = Take::new(Cursor::new(b"hello world".to_vec()), 5);
+    let mut res = vec![];
+
+    buf.read_to_end(&mut res);
+
+    assert_eq!(&res, b"hello");
+}