diff --git a/src/util/arena.rs b/src/util/arena.rs new file mode 100644 index 0000000000000000000000000000000000000000..d9b3f9585e08ce4771a73d2cd8960a8555f0600e --- /dev/null +++ b/src/util/arena.rs @@ -0,0 +1,106 @@ +use std::alloc::{alloc, dealloc, Layout}; +use std::ptr::NonNull; +use std::slice; + +// Arena block size +const ARENA_BLOCK_SIZE: usize = 4096; + +pub struct Arena { + alloc_ptr: Option>, + alloc_bytes_remaining: usize, + blocks: Vec<(NonNull, Layout)>, + memory_usage: usize, +} + +impl Default for Arena { + fn default() -> Self { + Self { + alloc_ptr: None, + alloc_bytes_remaining: 0, + blocks: vec![], + memory_usage: 0, + } + } +} + +impl Arena { + + /// 申请一块内存 + /// + /// # Arguments + /// + /// * `bytes`: 申请内存大小(byte) + /// + /// returns: &mut [u8] + /// 内存的 byte 数组 + /// # Examples + /// + /// ``` + ///let arena = Arena::default(); + /// // 申请 12 字节大小的内存 + /// let buf = arena.allocate(12); + /// ``` + #[inline] + pub fn allocate(&mut self, bytes: usize) -> &mut [u8] { + self.allocate_align(bytes, 1) + } + + pub fn allocate_align(&mut self, bytes: usize, align: usize) -> &mut [u8] { + if bytes <= self.alloc_bytes_remaining { + self.alloc_bytes_remaining -= bytes; + let result = unsafe { slice::from_raw_parts_mut(self.alloc_ptr.unwrap().as_ptr(), bytes) }; + unsafe { + let new_ptr = self.alloc_ptr.unwrap().as_ptr().offset(bytes as isize); + self.alloc_ptr = Some(NonNull::new_unchecked(new_ptr)); + }; + return result; + } + return self.allocate_fallback(bytes, align); + } + + pub fn memory_usage(&self) -> usize { + self.memory_usage + } + + fn allocate_fallback(&mut self, bytes: usize, align: usize) -> &mut [u8] { + if bytes > ARENA_BLOCK_SIZE / 4 { + unsafe { + let layout = Layout::from_size_align_unchecked(bytes, align); + return self.allocate_new_block(layout); + } + } + unsafe { + self.alloc_bytes_remaining = ARENA_BLOCK_SIZE - bytes; + let layout = Layout::from_size_align_unchecked(ARENA_BLOCK_SIZE, align); + let new_block = self.allocate_new_block(layout); + unsafe { + let ptr = new_block.as_ptr() as *mut u8; + let result = slice::from_raw_parts_mut(ptr, bytes); + self.alloc_ptr = Some(NonNull::new_unchecked(ptr.offset(bytes as isize))); + result + } + } + } + + + /// 分配一块新的内存 + fn allocate_new_block(&mut self, layout: Layout) -> &mut [u8] { + unsafe { + let data = alloc(layout); + self.memory_usage += layout.size(); + self.blocks.push((NonNull::new_unchecked(data), layout)); + slice::from_raw_parts_mut(data, layout.size()) + } + } +} + +impl Drop for Arena { + /// 释放内存 + fn drop(&mut self) { + for (block, layout) in self.blocks.iter() { + unsafe { + dealloc(block.as_ptr(), *layout) + } + } + } +} \ No newline at end of file diff --git a/src/util/arena_test.rs b/src/util/arena_test.rs new file mode 100644 index 0000000000000000000000000000000000000000..dbc8cf12ac94c5b61572f5b0f410a2c09f9a3fa3 --- /dev/null +++ b/src/util/arena_test.rs @@ -0,0 +1,35 @@ +use crate::util::Arena; + +#[test] +fn test_memory_usage() { + let mut arena = Arena::default(); + let _buf0 = arena.allocate(12); + let _buf1 = arena.allocate(16); + assert_eq!(4096, arena.memory_usage()); + let _buf2 = arena.allocate(3900); + assert_eq!(4096, arena.memory_usage()); + let _buf3 = arena.allocate(1200); + assert_eq!(4096 + 1200, arena.memory_usage()); +} + +#[test] +fn test_allocate() { + let mut arena = Arena::default(); + for i in 0..=12 { + let byte_size = 1 << i; + let buf = arena.allocate(byte_size); + assert_eq!(byte_size, buf.len()); + } + assert_eq!(8192, arena.memory_usage()); +} + +#[test] +fn test_allocate_align() { + let mut arena = Arena::default(); + for i in 0..=12 { + let byte_size = 1 << i; + let buf = arena.allocate_align(byte_size, 8); + assert_eq!(byte_size, buf.len()); + } + assert_eq!(4096 * 2, arena.memory_usage()); +} \ No newline at end of file diff --git a/src/util/coding.rs b/src/util/coding.rs index 74f2492e3cef4de6067d84fabb54e2da55413728..11709506e4f9a5e44ee5547313e51f383cd1485e 100644 --- a/src/util/coding.rs +++ b/src/util/coding.rs @@ -1,47 +1,3 @@ -const B: u8 = 128; +pub struct Coding { -/// -/// -/// -pub trait Coding { - fn varint32(self, buf: &mut [u8]); -} - -impl Coding for u32 -{ - fn varint32(self, buf: &mut [u8]) { - // if self < 1 << 7 { - // buf[0] = self as u8; - // } else if self < 1 << 14 { - // buf[0] = self | B; - // println!("{:b}", buf[0]); - // buf[1] = buf[0] >> 7; - // println!("{:b}", buf[1]); - // } else if self < 1 << 21 { - // buf[0] = (self | B as u8) as u8; - // println!("{:b}", buf[0]); - // buf[1] = (buf[0] >> 7) | B; - // println!("{:b}", buf[1]); - // buf[2] = buf[1] >> 14; - // println!("{:b}", buf[2]); - // } else if self < 1 << 28 { - // buf[0] = (self | B as u8) as u8; - // println!("{:b}", buf[0]); - // buf[1] = (buf[0] >> 7) | B; - // println!("{:b}", buf[1]); - // buf[2] = (buf[1] >> 14) | B; - // println!("{:b}", buf[2]); - // buf[3] = buf[2] >> 21; - // } else { - // buf[0] = (self | B as u8) as u8; - // println!("{:b}", buf[0]); - // buf[1] = (buf[0] >> 7) | B; - // println!("{:b}", buf[1]); - // buf[2] = (buf[1] >> 14) | B; - // println!("{:b}", buf[2]); - // buf[3] = (buf[2] >> 21) | B; - // println!("{:b}", buf[3]); - // buf[4] = buf[3] >> 28; - // } - } } \ No newline at end of file diff --git a/src/util/coding_test.rs b/src/util/coding_test.rs index 3c4184f62558e28f92d57534635ed78c81c3fc5c..ed0356514f7193fe33373cd1c9afc4120ae4903f 100644 --- a/src/util/coding_test.rs +++ b/src/util/coding_test.rs @@ -1,20 +1,4 @@ mod test { use crate::util::coding::Coding; - #[test] - fn test_i32() { - let mut buf: [u8; 8] = [0, 0, 0, 0, 0, 0, 0, 0]; - let x = 127.varint32(&mut buf); - println!("{:?}", buf); - println!("{:?}", x); - println!("{:b}", buf[0]); - } - - #[test] - fn test_offset() { - let buf: [u8; 3] = [0, 1, 2]; - let mut offset: usize = 0; - offset += 1; - println!("{:?}", buf[offset]) - } } diff --git a/src/util/mod.rs b/src/util/mod.rs index fe7c884251c3c2b51d0c269ad6bcae918f4b8908..567d645f98bcb5353685d655513fbd6d8c88612d 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -2,4 +2,8 @@ pub mod slice; pub mod status; mod slice_test; pub mod coding; -mod coding_test; \ No newline at end of file +mod coding_test; +pub mod arena; +mod arena_test; + +pub use arena::Arena; \ No newline at end of file