use std::alloc::{alloc, dealloc, Layout};
use std::ptr::NonNull;
const BLOCK_SIZE: usize = 128;
const BLOCK_ALIGN: usize = 8;
pub struct SimpleMemoryPool {
memory: NonNull<u8>,
total_blocks: usize,
free_list: Vec<usize>,
block_usage: Vec<bool>,
}
impl SimpleMemoryPool {
pub fn new(blocks: usize) -> Self {
let layout = Layout::from_size_align(blocks * BLOCK_SIZE, BLOCK_ALIGN)
.expect("无法创建 layout");
let memory = unsafe {
let ptr = alloc(layout);
NonNull::new(ptr).expect("内存分配失败")
};
let mut free_list = Vec::with_capacity(blocks);
for i in 0..blocks {
free_list.push(i);
}
SimpleMemoryPool {
memory,
total_blocks: blocks,
free_list,
block_usage: vec![false; blocks],
}
}
pub fn alloc_block(&mut self) -> Option<usize> {
let idx = self.free_list.pop()?;
self.block_usage[idx] = true;
Some(idx)
}
pub fn free_block(&mut self, idx: usize) {
if idx < self.total_blocks && self.block_usage[idx] {
self.block_usage[idx] = false;
self.free_list.push(idx);
}
}
pub fn block_ptr(&self, idx: usize) -> Option<*mut u8> {
if idx < self.total_blocks && self.block_usage[idx] {
let offset = idx * BLOCK_SIZE;
unsafe { Some(self.memory.as_ptr().add(offset)) }
} else {
None
}
}
pub unsafe fn write_str(&self, idx: usize, s: &str) {
if let Some(ptr) = self.block_ptr(idx) {
let bytes = s.as_bytes();
std::ptr::copy_nonoverlapping(bytes.as_ptr(), ptr, bytes.len());
*ptr.add(bytes.len()) = 0;
}
}
pub unsafe fn read_str(&self, idx: usize) -> Option<&str> {
self.block_ptr(idx).map(|ptr| {
let len = (0..BLOCK_SIZE)
.find(|&i| *ptr.add(i) == 0)
.unwrap_or(BLOCK_SIZE);
let slice = std::slice::from_raw_parts(ptr, len);
std::str::from_utf8_unchecked(slice)
})
}
}
impl Drop for SimpleMemoryPool {
fn drop(&mut self) {
let layout = Layout::from_size_align(self.total_blocks * BLOCK_SIZE, BLOCK_ALIGN)
.unwrap();
unsafe {
dealloc(self.memory.as_ptr(), layout);
}
}
}
unsafe impl Send for SimpleMemoryPool {}
unsafe impl Sync for SimpleMemoryPool {}