diff --git a/src/alloc/mod.rs b/src/alloc/mod.rs index 891412f491e8a5199d222cf536f7e4dbad0aa556..820ede7aa285c7377f7f033ab51307e27b827218 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 0000000000000000000000000000000000000000..6528eda8694e33e50ebd5985252775e1382fa6a2 --- /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 661abf9118817c86f4f64a76776ed3b1e0e7af5d..0bfe6d68034ee5628ba48e345287476c7c10c5cd 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 1ce2c00577bf620f80b31c497517f4acde03f6df..4b1e3eb275583c31eaa9823f067540027feae082 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 891eb12e7afd4eccea56207147df0250614f3fa6..e799e2a235ec2ccac2f61090f7b6ae810f856859 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 044f5079c4d9d58bce931c9eca50dc90e10db15e..b36456ad1d513d265c931bb4cacf4c84d41b3a70 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 9fc3f427abcbd1b4385a3fde731fc35daa3b5b4a..d8172d653596698421b8178f5afcd75b7ba76c65 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 0000000000000000000000000000000000000000..227e9aeedaab300ff468b932b71d33aaf7bce121 --- /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"); + } +}