1 use std::alloc::{GlobalAlloc, Layout, System};
2 use std::{mem, ptr};
3
4 use bytes::{Buf, Bytes};
5
6 #[global_allocator]
7 static LEDGER: Ledger = Ledger;
8
9 struct Ledger;
10
11 const USIZE_SIZE: usize = mem::size_of::<usize>();
12
13 unsafe impl GlobalAlloc for Ledger {
alloc(&self, layout: Layout) -> *mut u814 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
15 if layout.align() == 1 && layout.size() > 0 {
16 // Allocate extra space to stash a record of
17 // how much space there was.
18 let orig_size = layout.size();
19 let size = orig_size + USIZE_SIZE;
20 let new_layout = match Layout::from_size_align(size, 1) {
21 Ok(layout) => layout,
22 Err(_err) => return ptr::null_mut(),
23 };
24 let ptr = System.alloc(new_layout);
25 if !ptr.is_null() {
26 (ptr as *mut usize).write(orig_size);
27 let ptr = ptr.offset(USIZE_SIZE as isize);
28 ptr
29 } else {
30 ptr
31 }
32 } else {
33 System.alloc(layout)
34 }
35 }
36
dealloc(&self, ptr: *mut u8, layout: Layout)37 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
38 if layout.align() == 1 && layout.size() > 0 {
39 let off_ptr = (ptr as *mut usize).offset(-1);
40 let orig_size = off_ptr.read();
41 if orig_size != layout.size() {
42 panic!(
43 "bad dealloc: alloc size was {}, dealloc size is {}",
44 orig_size,
45 layout.size()
46 );
47 }
48
49 let new_layout = match Layout::from_size_align(layout.size() + USIZE_SIZE, 1) {
50 Ok(layout) => layout,
51 Err(_err) => std::process::abort(),
52 };
53 System.dealloc(off_ptr as *mut u8, new_layout);
54 } else {
55 System.dealloc(ptr, layout);
56 }
57 }
58 }
59 #[test]
test_bytes_advance()60 fn test_bytes_advance() {
61 let mut bytes = Bytes::from(vec![10, 20, 30]);
62 bytes.advance(1);
63 drop(bytes);
64 }
65
66 #[test]
test_bytes_truncate()67 fn test_bytes_truncate() {
68 let mut bytes = Bytes::from(vec![10, 20, 30]);
69 bytes.truncate(2);
70 drop(bytes);
71 }
72
73 #[test]
test_bytes_truncate_and_advance()74 fn test_bytes_truncate_and_advance() {
75 let mut bytes = Bytes::from(vec![10, 20, 30]);
76 bytes.truncate(2);
77 bytes.advance(1);
78 drop(bytes);
79 }
80