From d650404bb89da3c12529dc3ffbb7d4ffb7147506 Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Sat, 23 Jul 2016 09:49:59 -0700
Subject: [PATCH] Add an AppendBuf
---
src/alloc/mod.rs | 13 ++----
src/buf/append.rs | 111 ++++++++++++++++++++++++++++++++++++++++++++
src/buf/byte.rs | 14 ++++--
src/buf/mod.rs | 2 +
src/buf/ring.rs | 2 +-
src/str/seq.rs | 2 +-
test/test.rs | 1 +
test/test_append.rs | 30 ++++++++++++
8 files changed, 159 insertions(+), 16 deletions(-)
create mode 100644 src/buf/append.rs
create mode 100644 test/test_append.rs
diff --git a/src/alloc/mod.rs b/src/alloc/mod.rs
index 891412f..820ede7 100644
--- a/src/alloc/mod.rs
+++ b/src/alloc/mod.rs
@@ -49,20 +49,15 @@ impl MemRef {
}
#[inline]
- pub fn bytes(&self) -> &[u8] {
+ pub unsafe fn bytes(&self) -> &[u8] {
use std::slice;
-
- unsafe {
- slice::from_raw_parts(self.bytes_ptr(), self.len())
- }
+ slice::from_raw_parts(self.bytes_ptr(), self.len())
}
#[inline]
- pub fn bytes_mut(&mut self) -> &mut [u8] {
+ pub unsafe fn bytes_mut(&mut self) -> &mut [u8] {
use std::slice;
- unsafe {
- slice::from_raw_parts_mut(self.bytes_ptr(), self.len())
- }
+ slice::from_raw_parts_mut(self.bytes_ptr(), self.len())
}
#[inline]
diff --git a/src/buf/append.rs b/src/buf/append.rs
new file mode 100644
index 0000000..6528eda
--- /dev/null
+++ b/src/buf/append.rs
@@ -0,0 +1,111 @@
+use alloc;
+use buf::{MutBuf};
+use str::{ByteStr, Bytes, SeqByteStr, SmallByteStr};
+use std::cell::Cell;
+use std::cmp;
+
+/// A `Buf` backed by a contiguous region of memory.
+///
+/// This buffer can only be written to once. Byte strings (immutable views) can
+/// be created at any time, not just when the writing is complete.
+pub struct AppendBuf {
+ mem: alloc::MemRef,
+ rd: Cell<u32>, // Read cursor
+ wr: u32, // Write cursor
+ cap: u32,
+}
+
+impl AppendBuf {
+ pub fn with_capacity(mut capacity: u32) -> AppendBuf {
+ // Handle 0 capacity case
+ if capacity == 0 {
+ return AppendBuf::none();
+ }
+
+ // Round the capacity to the closest power of 2
+ capacity = capacity.next_power_of_two();
+
+ // Allocate the memory
+ let mem = alloc::heap(capacity as usize);
+
+ // If the allocation failed, return a blank buf
+ if mem.is_none() {
+ return AppendBuf::none();
+ }
+
+ AppendBuf {
+ mem: mem,
+ rd: Cell::new(0),
+ wr: 0,
+ cap: capacity,
+ }
+ }
+
+ /// Returns an AppendBuf with no capacity
+ pub fn none() -> AppendBuf {
+ AppendBuf {
+ mem: alloc::MemRef::none(),
+ rd: Cell::new(0),
+ wr: 0,
+ cap: 0,
+ }
+ }
+
+ pub unsafe fn from_mem_ref(mem: alloc::MemRef, cap: u32, pos: u32) -> AppendBuf {
+ AppendBuf {
+ mem: mem,
+ rd: Cell::new(pos),
+ wr: pos,
+ cap: cap,
+ }
+ }
+
+ pub fn bytes(&self) -> &[u8] {
+ let rd = self.rd.get() as usize;
+ let wr = self.wr as usize;
+ unsafe { &self.mem.bytes()[rd..wr] }
+ }
+
+ pub fn shift(&self, n: usize) -> Bytes {
+ let ret = self.slice(0, n);
+ self.rd.set(self.rd.get() + ret.len() as u32);
+ ret
+ }
+
+ pub fn slice(&self, begin: usize, end: usize) -> Bytes {
+ if end <= begin {
+ return Bytes::of(SmallByteStr::zero());
+ }
+
+ if let Some(bytes) = SmallByteStr::from_slice(&self.bytes()[begin..end]) {
+ return Bytes::of(bytes);
+ }
+
+ let begin = cmp::min(self.wr, begin as u32 + self.rd.get());
+ let end = cmp::min(self.wr, end as u32 + self.rd.get());
+
+ let bytes = unsafe { SeqByteStr::from_mem_ref(self.mem.clone(), begin, end - begin) };
+
+ Bytes::of(bytes)
+ }
+}
+
+impl MutBuf for AppendBuf {
+ fn remaining(&self) -> usize {
+ (self.cap - self.wr) as usize
+ }
+
+ unsafe fn advance(&mut self, cnt: usize) {
+ self.wr += cnt as u32;
+
+ if self.wr > self.cap {
+ self.wr = self.cap;
+ }
+ }
+
+ unsafe fn mut_bytes<'a>(&'a mut self) -> &'a mut [u8] {
+ let wr = self.wr as usize;
+ let cap = self.cap as usize;
+ &mut self.mem.bytes_mut()[wr..cap]
+ }
+}
diff --git a/src/buf/byte.rs b/src/buf/byte.rs
index 661abf9..0bfe6d6 100644
--- a/src/buf/byte.rs
+++ b/src/buf/byte.rs
@@ -103,7 +103,9 @@ impl ByteBuf {
let cnt = len as u32;
let pos = self.pos as usize;
- dst[0..len].copy_from_slice(&self.mem.bytes()[pos..pos+len]);
+ unsafe {
+ dst[0..len].copy_from_slice(&self.mem.bytes()[pos..pos+len]);
+ }
self.pos += cnt;
len
@@ -168,7 +170,7 @@ impl Buf for ByteBuf {
#[inline]
fn bytes<'a>(&'a self) -> &'a [u8] {
- &self.mem.bytes()[self.pos()..self.lim()]
+ unsafe { &self.mem.bytes()[self.pos()..self.lim()] }
}
#[inline]
@@ -294,8 +296,10 @@ impl MutByteBuf {
let cnt = cmp::min(src.len(), self.buf.remaining());
let pos = self.buf.pos as usize;
- self.buf.mem.bytes_mut()[pos..pos+cnt]
- .copy_from_slice(&src[0..cnt]);
+ unsafe {
+ self.buf.mem.bytes_mut()[pos..pos+cnt]
+ .copy_from_slice(&src[0..cnt]);
+ }
self.buf.pos += cnt as u32;
@@ -303,7 +307,7 @@ impl MutByteBuf {
}
pub fn bytes<'a>(&'a self) -> &'a [u8] {
- &self.buf.mem.bytes()[..self.buf.pos()]
+ unsafe { &self.buf.mem.bytes()[..self.buf.pos()] }
}
}
diff --git a/src/buf/mod.rs b/src/buf/mod.rs
index 1ce2c00..4b1e3eb 100644
--- a/src/buf/mod.rs
+++ b/src/buf/mod.rs
@@ -1,3 +1,4 @@
+mod append;
mod byte;
mod ring;
mod sink;
@@ -5,6 +6,7 @@ mod slice;
mod source;
mod take;
+pub use self::append::AppendBuf;
pub use self::byte::{ByteBuf, MutByteBuf, ROByteBuf};
pub use self::ring::RingBuf;
pub use self::slice::{SliceBuf, MutSliceBuf};
diff --git a/src/buf/ring.rs b/src/buf/ring.rs
index 891eb12..e799e2a 100644
--- a/src/buf/ring.rs
+++ b/src/buf/ring.rs
@@ -156,7 +156,7 @@ impl Buf for RingBuf {
to = self.cap
}
- &self.ptr.bytes()[self.pos .. to]
+ unsafe { &self.ptr.bytes()[self.pos .. to] }
}
fn advance(&mut self, cnt: usize) {
diff --git a/src/str/seq.rs b/src/str/seq.rs
index 044f507..b36456a 100644
--- a/src/str/seq.rs
+++ b/src/str/seq.rs
@@ -81,7 +81,7 @@ impl ops::Index<usize> for SeqByteStr {
fn index(&self, index: usize) -> &u8 {
assert!(index < self.len());
- self.mem.bytes().index(index + self.pos as usize)
+ unsafe { self.mem.bytes().index(index + self.pos as usize) }
}
}
diff --git a/test/test.rs b/test/test.rs
index 9fc3f42..d8172d6 100644
--- a/test/test.rs
+++ b/test/test.rs
@@ -4,6 +4,7 @@ extern crate bytes;
extern crate rand;
extern crate byteorder;
+mod test_append;
mod test_buf;
mod test_buf_fill;
mod test_buf_take;
diff --git a/test/test_append.rs b/test/test_append.rs
new file mode 100644
index 0000000..227e9ae
--- /dev/null
+++ b/test/test_append.rs
@@ -0,0 +1,30 @@
+use bytes::{ByteStr, Buf, MutBuf};
+use bytes::buf::AppendBuf;
+
+#[test]
+pub fn test_initial_buf_empty() {
+ // Run in a loop a bunch in hope that if there is a memory issue, it will
+ // be exposed
+ for _ in 0..1000 {
+ let mut buf = AppendBuf::with_capacity(100);
+ let mut dst: Vec<u8> = vec![];
+
+ assert_eq!(buf.remaining(), 128);
+
+ buf.write_slice(b"hello world");
+ assert_eq!(buf.remaining(), 117);
+ assert_eq!(buf.bytes(), b"hello world");
+
+ let view1 = buf.slice(0, 11);
+ view1.buf().copy_to(&mut dst).unwrap();
+
+ assert_eq!(dst, b"hello world");
+ assert_eq!(view1, buf.slice(0, 11));
+
+ drop(buf);
+ let mut buf = AppendBuf::with_capacity(100);
+ buf.write_slice(b"zomg no no no no");
+
+ assert_eq!(dst, b"hello world");
+ }
+}
--
GitLab