1 #![feature(allocator_api)]
2 #![cfg(feature = "allocator_api")]
3 use bumpalo::Bump;
4
5 use std::alloc::{AllocError, Allocator, Layout};
6 use std::ptr::NonNull;
7 use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
8
9 #[derive(Debug)]
10 struct AllocatorDebug {
11 bump: Bump,
12 grows: AtomicUsize,
13 shrinks: AtomicUsize,
14 allocs: AtomicUsize,
15 deallocs: AtomicUsize,
16 }
17
18 impl AllocatorDebug {
new(bump: Bump) -> AllocatorDebug19 fn new(bump: Bump) -> AllocatorDebug {
20 AllocatorDebug {
21 bump,
22 grows: AtomicUsize::new(0),
23 shrinks: AtomicUsize::new(0),
24 allocs: AtomicUsize::new(0),
25 deallocs: AtomicUsize::new(0),
26 }
27 }
28 }
29
30 unsafe impl Allocator for AllocatorDebug {
allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError>31 fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
32 self.allocs.fetch_add(1, Relaxed);
33 let ref bump = self.bump;
34 bump.allocate(layout)
35 }
36
deallocate(&self, ptr: NonNull<u8>, layout: Layout)37 unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
38 self.deallocs.fetch_add(1, Relaxed);
39 let ref bump = self.bump;
40 bump.deallocate(ptr, layout)
41 }
42
shrink( &self, ptr: NonNull<u8>, old_layout: Layout, new_layout: Layout, ) -> Result<NonNull<[u8]>, AllocError>43 unsafe fn shrink(
44 &self,
45 ptr: NonNull<u8>,
46 old_layout: Layout,
47 new_layout: Layout,
48 ) -> Result<NonNull<[u8]>, AllocError> {
49 self.shrinks.fetch_add(1, Relaxed);
50 let ref bump = self.bump;
51 bump.shrink(ptr, old_layout, new_layout)
52 }
53
grow( &self, ptr: NonNull<u8>, old_layout: Layout, new_layout: Layout, ) -> Result<NonNull<[u8]>, AllocError>54 unsafe fn grow(
55 &self,
56 ptr: NonNull<u8>,
57 old_layout: Layout,
58 new_layout: Layout,
59 ) -> Result<NonNull<[u8]>, AllocError> {
60 self.grows.fetch_add(1, Relaxed);
61 let ref bump = self.bump;
62 bump.grow(ptr, old_layout, new_layout)
63 }
64 }
65
66 #[test]
allocator_api_push_a_bunch_of_items()67 fn allocator_api_push_a_bunch_of_items() {
68 let b = AllocatorDebug::new(Bump::new());
69 let mut v = Vec::with_capacity_in(1024, &b);
70 assert_eq!(b.allocs.load(Relaxed), 1);
71
72 for x in 0..1024 {
73 v.push(x);
74 }
75
76 // Ensure we trigger a grow
77 assert_eq!(b.grows.load(Relaxed), 0);
78 for x in 1024..2048 {
79 v.push(x);
80 }
81 assert_ne!(b.grows.load(Relaxed), 0);
82
83 // Ensure we trigger a shrink
84 v.truncate(1024);
85 v.shrink_to_fit();
86 assert_eq!(b.shrinks.load(Relaxed), 1);
87
88 // Ensure we trigger a deallocation
89 assert_eq!(b.deallocs.load(Relaxed), 0);
90 drop(v);
91 assert_eq!(b.deallocs.load(Relaxed), 1);
92 }
93
94 #[test]
allocator_grow_zeroed()95 fn allocator_grow_zeroed() {
96 // Create a new bump arena.
97 let ref bump = Bump::new();
98
99 // Make an initial allocation.
100 let first_layout = Layout::from_size_align(4, 4).expect("create a layout");
101 let mut p = bump
102 .allocate_zeroed(first_layout)
103 .expect("allocate a first chunk");
104 let allocated = bump.allocated_bytes();
105 unsafe { p.as_mut().fill(42) };
106 let p = p.cast();
107
108 // Grow the last allocation. This should just reserve a few more bytes
109 // within the current chunk, not allocate a whole new memory block within a
110 // new chunk.
111 let second_layout = Layout::from_size_align(8, 4).expect("create a expanded layout");
112 let p = unsafe { bump.grow_zeroed(p, first_layout, second_layout) }
113 .expect("should grow_zeroed okay");
114 assert!(bump.allocated_bytes() <= allocated * 2);
115 assert_eq!(unsafe { p.as_ref() }, [42, 42, 42, 42, 0, 0, 0, 0]);
116 }
117