Add an AppendBuf

This commit is contained in:
Carl Lerche
2016-07-23 09:49:59 -07:00
parent 1c2234f7fe
commit d650404bb8
8 changed files with 159 additions and 16 deletions
+4 -9
View File
@@ -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]
+111
View File
@@ -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]
}
}
+9 -5
View File
@@ -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()] }
}
}
+2
View File
@@ -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};
+1 -1
View File
@@ -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) {
+1 -1
View File
@@ -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) }
}
}
+1
View File
@@ -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;
+30
View File
@@ -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");
}
}