From ed087be8532dba1e698e57cbb99a8e5428244985 Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Tue, 28 Jul 2015 12:51:24 -0700
Subject: [PATCH] Provide `Take` decorator that limits `Buf` size
---
src/buf/mod.rs | 2 ++
src/buf/take.rs | 79 +++++++++++++++++++++++++++++++++++++++++++
src/lib.rs | 3 +-
test/test.rs | 1 +
test/test_buf_take.rs | 12 +++++++
5 files changed, 96 insertions(+), 1 deletion(-)
create mode 100644 src/buf/take.rs
create mode 100644 test/test_buf_take.rs
diff --git a/src/buf/mod.rs b/src/buf/mod.rs
index 07de5f8..2498b19 100644
--- a/src/buf/mod.rs
+++ b/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};
diff --git a/src/buf/take.rs b/src/buf/take.rs
new file mode 100644
index 0000000..9995152
--- /dev/null
+++ b/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);
+ }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 42a9eab..820f49b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -17,7 +17,8 @@ pub use buf::{
SliceBuf,
MutSliceBuf,
Source,
- Sink
+ Sink,
+ Take,
};
pub use str::{
ByteStr,
diff --git a/test/test.rs b/test/test.rs
index 283685f..77fcd57 100644
--- a/test/test.rs
+++ b/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;
diff --git a/test/test_buf_take.rs b/test/test_buf_take.rs
new file mode 100644
index 0000000..b2300c1
--- /dev/null
+++ b/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");
+}
--
GitLab