80 lines
2.2 KiB
Rust
80 lines
2.2 KiB
Rust
use std::alloc::{GlobalAlloc, Layout, System};
|
|
use std::{mem, ptr};
|
|
|
|
use bytes::{Buf, Bytes};
|
|
|
|
#[global_allocator]
|
|
static LEDGER: Ledger = Ledger;
|
|
|
|
struct Ledger;
|
|
|
|
const USIZE_SIZE: usize = mem::size_of::<usize>();
|
|
|
|
unsafe impl GlobalAlloc for Ledger {
|
|
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
|
|
if layout.align() == 1 && layout.size() > 0 {
|
|
// Allocate extra space to stash a record of
|
|
// how much space there was.
|
|
let orig_size = layout.size();
|
|
let size = orig_size + USIZE_SIZE;
|
|
let new_layout = match Layout::from_size_align(size, 1) {
|
|
Ok(layout) => layout,
|
|
Err(_err) => return ptr::null_mut(),
|
|
};
|
|
let ptr = System.alloc(new_layout);
|
|
if !ptr.is_null() {
|
|
(ptr as *mut usize).write(orig_size);
|
|
let ptr = ptr.offset(USIZE_SIZE as isize);
|
|
ptr
|
|
} else {
|
|
ptr
|
|
}
|
|
} else {
|
|
System.alloc(layout)
|
|
}
|
|
}
|
|
|
|
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
|
|
if layout.align() == 1 && layout.size() > 0 {
|
|
let off_ptr = (ptr as *mut usize).offset(-1);
|
|
let orig_size = off_ptr.read();
|
|
if orig_size != layout.size() {
|
|
panic!(
|
|
"bad dealloc: alloc size was {}, dealloc size is {}",
|
|
orig_size,
|
|
layout.size()
|
|
);
|
|
}
|
|
|
|
let new_layout = match Layout::from_size_align(layout.size() + USIZE_SIZE, 1) {
|
|
Ok(layout) => layout,
|
|
Err(_err) => std::process::abort(),
|
|
};
|
|
System.dealloc(off_ptr as *mut u8, new_layout);
|
|
} else {
|
|
System.dealloc(ptr, layout);
|
|
}
|
|
}
|
|
}
|
|
#[test]
|
|
fn test_bytes_advance() {
|
|
let mut bytes = Bytes::from(vec![10, 20, 30]);
|
|
bytes.advance(1);
|
|
drop(bytes);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bytes_truncate() {
|
|
let mut bytes = Bytes::from(vec![10, 20, 30]);
|
|
bytes.truncate(2);
|
|
drop(bytes);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bytes_truncate_and_advance() {
|
|
let mut bytes = Bytes::from(vec![10, 20, 30]);
|
|
bytes.truncate(2);
|
|
bytes.advance(1);
|
|
drop(bytes);
|
|
}
|