]> git.lizzy.rs Git - rust.git/blob - src/libarena/lib.rs
Fix the new capacity measurement in arenas.
[rust.git] / src / libarena / lib.rs
1 //! The arena, a fast but limited type of allocator.
2 //!
3 //! Arenas are a type of allocator that destroy the objects within, all at
4 //! once, once the arena itself is destroyed. They do not support deallocation
5 //! of individual objects while the arena itself is still alive. The benefit
6 //! of an arena is very fast allocation; just a pointer bump.
7 //!
8 //! This crate implements `TypedArena`, a simple arena that can only hold
9 //! objects of a single type.
10
11 #![doc(
12     html_root_url = "https://doc.rust-lang.org/nightly/",
13     test(no_crate_inject, attr(deny(warnings)))
14 )]
15 #![feature(core_intrinsics)]
16 #![feature(dropck_eyepatch)]
17 #![feature(raw_vec_internals)]
18 #![cfg_attr(test, feature(test))]
19 #![allow(deprecated)]
20
21 extern crate alloc;
22
23 use rustc_data_structures::cold_path;
24 use smallvec::SmallVec;
25
26 use std::cell::{Cell, RefCell};
27 use std::cmp;
28 use std::intrinsics;
29 use std::marker::{PhantomData, Send};
30 use std::mem;
31 use std::ptr;
32 use std::slice;
33
34 use alloc::raw_vec::RawVec;
35
36 /// An arena that can hold objects of only one type.
37 pub struct TypedArena<T> {
38     /// A pointer to the next object to be allocated.
39     ptr: Cell<*mut T>,
40
41     /// A pointer to the end of the allocated area. When this pointer is
42     /// reached, a new chunk is allocated.
43     end: Cell<*mut T>,
44
45     /// A vector of arena chunks.
46     chunks: RefCell<Vec<TypedArenaChunk<T>>>,
47
48     /// Marker indicating that dropping the arena causes its owned
49     /// instances of `T` to be dropped.
50     _own: PhantomData<T>,
51 }
52
53 struct TypedArenaChunk<T> {
54     /// The raw storage for the arena chunk.
55     storage: RawVec<T>,
56     /// The number of valid entries in the chunk.
57     entries: usize,
58 }
59
60 impl<T> TypedArenaChunk<T> {
61     #[inline]
62     unsafe fn new(capacity: usize) -> TypedArenaChunk<T> {
63         TypedArenaChunk { storage: RawVec::with_capacity(capacity), entries: 0 }
64     }
65
66     /// Destroys this arena chunk.
67     #[inline]
68     unsafe fn destroy(&mut self, len: usize) {
69         // The branch on needs_drop() is an -O1 performance optimization.
70         // Without the branch, dropping TypedArena<u8> takes linear time.
71         if mem::needs_drop::<T>() {
72             let mut start = self.start();
73             // Destroy all allocated objects.
74             for _ in 0..len {
75                 ptr::drop_in_place(start);
76                 start = start.offset(1);
77             }
78         }
79     }
80
81     // Returns a pointer to the first allocated object.
82     #[inline]
83     fn start(&self) -> *mut T {
84         self.storage.ptr()
85     }
86
87     // Returns a pointer to the end of the allocated space.
88     #[inline]
89     fn end(&self) -> *mut T {
90         unsafe {
91             if mem::size_of::<T>() == 0 {
92                 // A pointer as large as possible for zero-sized elements.
93                 !0 as *mut T
94             } else {
95                 self.start().add(self.storage.capacity())
96             }
97         }
98     }
99 }
100
101 const PAGE: usize = 4096;
102
103 impl<T> Default for TypedArena<T> {
104     /// Creates a new `TypedArena`.
105     fn default() -> TypedArena<T> {
106         TypedArena {
107             // We set both `ptr` and `end` to 0 so that the first call to
108             // alloc() will trigger a grow().
109             ptr: Cell::new(ptr::null_mut()),
110             end: Cell::new(ptr::null_mut()),
111             chunks: RefCell::new(vec![]),
112             _own: PhantomData,
113         }
114     }
115 }
116
117 impl<T> TypedArena<T> {
118     /// Allocates an object in the `TypedArena`, returning a reference to it.
119     #[inline]
120     pub fn alloc(&self, object: T) -> &mut T {
121         if self.ptr == self.end {
122             self.grow(1)
123         }
124
125         unsafe {
126             if mem::size_of::<T>() == 0 {
127                 self.ptr.set(intrinsics::arith_offset(self.ptr.get() as *mut u8, 1) as *mut T);
128                 let ptr = mem::align_of::<T>() as *mut T;
129                 // Don't drop the object. This `write` is equivalent to `forget`.
130                 ptr::write(ptr, object);
131                 &mut *ptr
132             } else {
133                 let ptr = self.ptr.get();
134                 // Advance the pointer.
135                 self.ptr.set(self.ptr.get().offset(1));
136                 // Write into uninitialized memory.
137                 ptr::write(ptr, object);
138                 &mut *ptr
139             }
140         }
141     }
142
143     #[inline]
144     fn can_allocate(&self, len: usize) -> bool {
145         let available_capacity_bytes = self.end.get() as usize - self.ptr.get() as usize;
146         let at_least_bytes = len.checked_mul(mem::size_of::<T>()).unwrap();
147         available_capacity_bytes >= at_least_bytes
148     }
149
150     /// Ensures there's enough space in the current chunk to fit `len` objects.
151     #[inline]
152     fn ensure_capacity(&self, len: usize) {
153         if !self.can_allocate(len) {
154             self.grow(len);
155             debug_assert!(self.can_allocate(len));
156         }
157     }
158
159     #[inline]
160     unsafe fn alloc_raw_slice(&self, len: usize) -> *mut T {
161         assert!(mem::size_of::<T>() != 0);
162         assert!(len != 0);
163
164         self.ensure_capacity(len);
165
166         let start_ptr = self.ptr.get();
167         self.ptr.set(start_ptr.add(len));
168         start_ptr
169     }
170
171     /// Allocates a slice of objects that are copied into the `TypedArena`, returning a mutable
172     /// reference to it. Will panic if passed a zero-sized types.
173     ///
174     /// Panics:
175     ///
176     ///  - Zero-sized types
177     ///  - Zero-length slices
178     #[inline]
179     pub fn alloc_slice(&self, slice: &[T]) -> &mut [T]
180     where
181         T: Copy,
182     {
183         unsafe {
184             let len = slice.len();
185             let start_ptr = self.alloc_raw_slice(len);
186             slice.as_ptr().copy_to_nonoverlapping(start_ptr, len);
187             slice::from_raw_parts_mut(start_ptr, len)
188         }
189     }
190
191     #[inline]
192     pub fn alloc_from_iter<I: IntoIterator<Item = T>>(&self, iter: I) -> &mut [T] {
193         assert!(mem::size_of::<T>() != 0);
194         let mut vec: SmallVec<[_; 8]> = iter.into_iter().collect();
195         if vec.is_empty() {
196             return &mut [];
197         }
198         // Move the content to the arena by copying it and then forgetting
199         // the content of the SmallVec
200         unsafe {
201             let len = vec.len();
202             let start_ptr = self.alloc_raw_slice(len);
203             vec.as_ptr().copy_to_nonoverlapping(start_ptr, len);
204             vec.set_len(0);
205             slice::from_raw_parts_mut(start_ptr, len)
206         }
207     }
208
209     /// Grows the arena.
210     #[inline(never)]
211     #[cold]
212     fn grow(&self, n: usize) {
213         unsafe {
214             let mut chunks = self.chunks.borrow_mut();
215             let (chunk, mut new_capacity);
216             if let Some(last_chunk) = chunks.last_mut() {
217                 let used_bytes = self.ptr.get() as usize - last_chunk.start() as usize;
218                 let currently_used_cap = used_bytes / mem::size_of::<T>();
219                 last_chunk.entries = currently_used_cap;
220                 if last_chunk.storage.reserve_in_place(currently_used_cap, n) {
221                     self.end.set(last_chunk.end());
222                     return;
223                 } else {
224                     new_capacity = last_chunk.storage.capacity();
225                     loop {
226                         new_capacity = new_capacity.checked_mul(2).unwrap();
227                         if new_capacity >= n {
228                             break;
229                         }
230                     }
231                 }
232             } else {
233                 let elem_size = cmp::max(1, mem::size_of::<T>());
234                 new_capacity = cmp::max(n, PAGE / elem_size);
235             }
236             chunk = TypedArenaChunk::<T>::new(new_capacity);
237             self.ptr.set(chunk.start());
238             self.end.set(chunk.end());
239             chunks.push(chunk);
240         }
241     }
242
243     /// Clears the arena. Deallocates all but the longest chunk which may be reused.
244     pub fn clear(&mut self) {
245         unsafe {
246             // Clear the last chunk, which is partially filled.
247             let mut chunks_borrow = self.chunks.borrow_mut();
248             if let Some(mut last_chunk) = chunks_borrow.last_mut() {
249                 self.clear_last_chunk(&mut last_chunk);
250                 let len = chunks_borrow.len();
251                 // If `T` is ZST, code below has no effect.
252                 for mut chunk in chunks_borrow.drain(..len - 1) {
253                     chunk.destroy(chunk.entries);
254                 }
255             }
256         }
257     }
258
259     // Drops the contents of the last chunk. The last chunk is partially empty, unlike all other
260     // chunks.
261     fn clear_last_chunk(&self, last_chunk: &mut TypedArenaChunk<T>) {
262         // Determine how much was filled.
263         let start = last_chunk.start() as usize;
264         // We obtain the value of the pointer to the first uninitialized element.
265         let end = self.ptr.get() as usize;
266         // We then calculate the number of elements to be dropped in the last chunk,
267         // which is the filled area's length.
268         let diff = if mem::size_of::<T>() == 0 {
269             // `T` is ZST. It can't have a drop flag, so the value here doesn't matter. We get
270             // the number of zero-sized values in the last and only chunk, just out of caution.
271             // Recall that `end` was incremented for each allocated value.
272             end - start
273         } else {
274             (end - start) / mem::size_of::<T>()
275         };
276         // Pass that to the `destroy` method.
277         unsafe {
278             last_chunk.destroy(diff);
279         }
280         // Reset the chunk.
281         self.ptr.set(last_chunk.start());
282     }
283 }
284
285 unsafe impl<#[may_dangle] T> Drop for TypedArena<T> {
286     fn drop(&mut self) {
287         unsafe {
288             // Determine how much was filled.
289             let mut chunks_borrow = self.chunks.borrow_mut();
290             if let Some(mut last_chunk) = chunks_borrow.pop() {
291                 // Drop the contents of the last chunk.
292                 self.clear_last_chunk(&mut last_chunk);
293                 // The last chunk will be dropped. Destroy all other chunks.
294                 for chunk in chunks_borrow.iter_mut() {
295                     chunk.destroy(chunk.entries);
296                 }
297             }
298             // RawVec handles deallocation of `last_chunk` and `self.chunks`.
299         }
300     }
301 }
302
303 unsafe impl<T: Send> Send for TypedArena<T> {}
304
305 pub struct DroplessArena {
306     /// A pointer to the next object to be allocated.
307     ptr: Cell<*mut u8>,
308
309     /// A pointer to the end of the allocated area. When this pointer is
310     /// reached, a new chunk is allocated.
311     end: Cell<*mut u8>,
312
313     /// A vector of arena chunks.
314     chunks: RefCell<Vec<TypedArenaChunk<u8>>>,
315 }
316
317 unsafe impl Send for DroplessArena {}
318
319 impl Default for DroplessArena {
320     #[inline]
321     fn default() -> DroplessArena {
322         DroplessArena {
323             ptr: Cell::new(ptr::null_mut()),
324             end: Cell::new(ptr::null_mut()),
325             chunks: Default::default(),
326         }
327     }
328 }
329
330 impl DroplessArena {
331     #[inline]
332     fn align(&self, align: usize) {
333         let final_address = ((self.ptr.get() as usize) + align - 1) & !(align - 1);
334         self.ptr.set(final_address as *mut u8);
335         assert!(self.ptr <= self.end);
336     }
337
338     #[inline(never)]
339     #[cold]
340     fn grow(&self, needed_bytes: usize) {
341         unsafe {
342             let mut chunks = self.chunks.borrow_mut();
343             let (chunk, mut new_capacity);
344             if let Some(last_chunk) = chunks.last_mut() {
345                 let used_bytes = self.ptr.get() as usize - last_chunk.start() as usize;
346                 if last_chunk.storage.reserve_in_place(used_bytes, needed_bytes) {
347                     self.end.set(last_chunk.end());
348                     return;
349                 } else {
350                     new_capacity = last_chunk.storage.capacity();
351                     loop {
352                         new_capacity = new_capacity.checked_mul(2).unwrap();
353                         if new_capacity >= needed_bytes {
354                             break;
355                         }
356                     }
357                 }
358             } else {
359                 new_capacity = cmp::max(needed_bytes, PAGE);
360             }
361             chunk = TypedArenaChunk::<u8>::new(new_capacity);
362             self.ptr.set(chunk.start());
363             self.end.set(chunk.end());
364             chunks.push(chunk);
365         }
366     }
367
368     #[inline]
369     pub fn alloc_raw(&self, bytes: usize, align: usize) -> &mut [u8] {
370         unsafe {
371             assert!(bytes != 0);
372
373             self.align(align);
374
375             let future_end = intrinsics::arith_offset(self.ptr.get(), bytes as isize);
376             if (future_end as *mut u8) >= self.end.get() {
377                 self.grow(bytes);
378             }
379
380             let ptr = self.ptr.get();
381             // Set the pointer past ourselves
382             self.ptr.set(intrinsics::arith_offset(self.ptr.get(), bytes as isize) as *mut u8);
383             slice::from_raw_parts_mut(ptr, bytes)
384         }
385     }
386
387     #[inline]
388     pub fn alloc<T>(&self, object: T) -> &mut T {
389         assert!(!mem::needs_drop::<T>());
390
391         let mem = self.alloc_raw(mem::size_of::<T>(), mem::align_of::<T>()) as *mut _ as *mut T;
392
393         unsafe {
394             // Write into uninitialized memory.
395             ptr::write(mem, object);
396             &mut *mem
397         }
398     }
399
400     /// Allocates a slice of objects that are copied into the `DroplessArena`, returning a mutable
401     /// reference to it. Will panic if passed a zero-sized type.
402     ///
403     /// Panics:
404     ///
405     ///  - Zero-sized types
406     ///  - Zero-length slices
407     #[inline]
408     pub fn alloc_slice<T>(&self, slice: &[T]) -> &mut [T]
409     where
410         T: Copy,
411     {
412         assert!(!mem::needs_drop::<T>());
413         assert!(mem::size_of::<T>() != 0);
414         assert!(!slice.is_empty());
415
416         let mem = self.alloc_raw(slice.len() * mem::size_of::<T>(), mem::align_of::<T>()) as *mut _
417             as *mut T;
418
419         unsafe {
420             let arena_slice = slice::from_raw_parts_mut(mem, slice.len());
421             arena_slice.copy_from_slice(slice);
422             arena_slice
423         }
424     }
425
426     #[inline]
427     unsafe fn write_from_iter<T, I: Iterator<Item = T>>(
428         &self,
429         mut iter: I,
430         len: usize,
431         mem: *mut T,
432     ) -> &mut [T] {
433         let mut i = 0;
434         // Use a manual loop since LLVM manages to optimize it better for
435         // slice iterators
436         loop {
437             let value = iter.next();
438             if i >= len || value.is_none() {
439                 // We only return as many items as the iterator gave us, even
440                 // though it was supposed to give us `len`
441                 return slice::from_raw_parts_mut(mem, i);
442             }
443             ptr::write(mem.add(i), value.unwrap());
444             i += 1;
445         }
446     }
447
448     #[inline]
449     pub fn alloc_from_iter<T, I: IntoIterator<Item = T>>(&self, iter: I) -> &mut [T] {
450         let iter = iter.into_iter();
451         assert!(mem::size_of::<T>() != 0);
452         assert!(!mem::needs_drop::<T>());
453
454         let size_hint = iter.size_hint();
455
456         match size_hint {
457             (min, Some(max)) if min == max => {
458                 // We know the exact number of elements the iterator will produce here
459                 let len = min;
460
461                 if len == 0 {
462                     return &mut [];
463                 }
464                 let size = len.checked_mul(mem::size_of::<T>()).unwrap();
465                 let mem = self.alloc_raw(size, mem::align_of::<T>()) as *mut _ as *mut T;
466                 unsafe { self.write_from_iter(iter, len, mem) }
467             }
468             (_, _) => {
469                 cold_path(move || -> &mut [T] {
470                     let mut vec: SmallVec<[_; 8]> = iter.collect();
471                     if vec.is_empty() {
472                         return &mut [];
473                     }
474                     // Move the content to the arena by copying it and then forgetting
475                     // the content of the SmallVec
476                     unsafe {
477                         let len = vec.len();
478                         let start_ptr = self
479                             .alloc_raw(len * mem::size_of::<T>(), mem::align_of::<T>())
480                             as *mut _ as *mut T;
481                         vec.as_ptr().copy_to_nonoverlapping(start_ptr, len);
482                         vec.set_len(0);
483                         slice::from_raw_parts_mut(start_ptr, len)
484                     }
485                 })
486             }
487         }
488     }
489 }
490
491 /// Calls the destructor for an object when dropped.
492 struct DropType {
493     drop_fn: unsafe fn(*mut u8),
494     obj: *mut u8,
495 }
496
497 unsafe fn drop_for_type<T>(to_drop: *mut u8) {
498     std::ptr::drop_in_place(to_drop as *mut T)
499 }
500
501 impl Drop for DropType {
502     fn drop(&mut self) {
503         unsafe { (self.drop_fn)(self.obj) }
504     }
505 }
506
507 /// An arena which can be used to allocate any type.
508 /// Allocating in this arena is unsafe since the type system
509 /// doesn't know which types it contains. In order to
510 /// allocate safely, you must store a PhantomData<T>
511 /// alongside this arena for each type T you allocate.
512 #[derive(Default)]
513 pub struct DropArena {
514     /// A list of destructors to run when the arena drops.
515     /// Ordered so `destructors` gets dropped before the arena
516     /// since its destructor can reference memory in the arena.
517     destructors: RefCell<Vec<DropType>>,
518     arena: DroplessArena,
519 }
520
521 impl DropArena {
522     #[inline]
523     pub unsafe fn alloc<T>(&self, object: T) -> &mut T {
524         let mem =
525             self.arena.alloc_raw(mem::size_of::<T>(), mem::align_of::<T>()) as *mut _ as *mut T;
526         // Write into uninitialized memory.
527         ptr::write(mem, object);
528         let result = &mut *mem;
529         // Record the destructor after doing the allocation as that may panic
530         // and would cause `object`'s destuctor to run twice if it was recorded before
531         self.destructors
532             .borrow_mut()
533             .push(DropType { drop_fn: drop_for_type::<T>, obj: result as *mut T as *mut u8 });
534         result
535     }
536
537     #[inline]
538     pub unsafe fn alloc_from_iter<T, I: IntoIterator<Item = T>>(&self, iter: I) -> &mut [T] {
539         let mut vec: SmallVec<[_; 8]> = iter.into_iter().collect();
540         if vec.is_empty() {
541             return &mut [];
542         }
543         let len = vec.len();
544
545         let start_ptr = self
546             .arena
547             .alloc_raw(len.checked_mul(mem::size_of::<T>()).unwrap(), mem::align_of::<T>())
548             as *mut _ as *mut T;
549
550         let mut destructors = self.destructors.borrow_mut();
551         // Reserve space for the destructors so we can't panic while adding them
552         destructors.reserve(len);
553
554         // Move the content to the arena by copying it and then forgetting
555         // the content of the SmallVec
556         vec.as_ptr().copy_to_nonoverlapping(start_ptr, len);
557         mem::forget(vec.drain(..));
558
559         // Record the destructors after doing the allocation as that may panic
560         // and would cause `object`'s destuctor to run twice if it was recorded before
561         for i in 0..len {
562             destructors.push(DropType {
563                 drop_fn: drop_for_type::<T>,
564                 obj: start_ptr.offset(i as isize) as *mut u8,
565             });
566         }
567
568         slice::from_raw_parts_mut(start_ptr, len)
569     }
570 }
571
572 #[macro_export]
573 macro_rules! arena_for_type {
574     ([][$ty:ty]) => {
575         $crate::TypedArena<$ty>
576     };
577     ([few $(, $attrs:ident)*][$ty:ty]) => {
578         ::std::marker::PhantomData<$ty>
579     };
580     ([$ignore:ident $(, $attrs:ident)*]$args:tt) => {
581         $crate::arena_for_type!([$($attrs),*]$args)
582     };
583 }
584
585 #[macro_export]
586 macro_rules! which_arena_for_type {
587     ([][$arena:expr]) => {
588         ::std::option::Option::Some($arena)
589     };
590     ([few$(, $attrs:ident)*][$arena:expr]) => {
591         ::std::option::Option::None
592     };
593     ([$ignore:ident$(, $attrs:ident)*]$args:tt) => {
594         $crate::which_arena_for_type!([$($attrs),*]$args)
595     };
596 }
597
598 #[macro_export]
599 macro_rules! declare_arena {
600     ([], [$($a:tt $name:ident: $ty:ty,)*], $tcx:lifetime) => {
601         #[derive(Default)]
602         pub struct Arena<$tcx> {
603             pub dropless: $crate::DroplessArena,
604             drop: $crate::DropArena,
605             $($name: $crate::arena_for_type!($a[$ty]),)*
606         }
607
608         #[marker]
609         pub trait ArenaAllocatable {}
610
611         impl<T: Copy> ArenaAllocatable for T {}
612
613         unsafe trait ArenaField<'tcx>: Sized {
614             /// Returns a specific arena to allocate from.
615             /// If `None` is returned, the `DropArena` will be used.
616             fn arena<'a>(arena: &'a Arena<'tcx>) -> Option<&'a $crate::TypedArena<Self>>;
617         }
618
619         unsafe impl<'tcx, T> ArenaField<'tcx> for T {
620             #[inline]
621             default fn arena<'a>(_: &'a Arena<'tcx>) -> Option<&'a $crate::TypedArena<Self>> {
622                 panic!()
623             }
624         }
625
626         $(
627             #[allow(unused_lifetimes)]
628             impl<$tcx> ArenaAllocatable for $ty {}
629             unsafe impl<$tcx> ArenaField<$tcx> for $ty {
630                 #[inline]
631                 fn arena<'a>(_arena: &'a Arena<$tcx>) -> Option<&'a $crate::TypedArena<Self>> {
632                     $crate::which_arena_for_type!($a[&_arena.$name])
633                 }
634             }
635         )*
636
637         impl<'tcx> Arena<'tcx> {
638             #[inline]
639             pub fn alloc<T: ArenaAllocatable>(&self, value: T) -> &mut T {
640                 if !::std::mem::needs_drop::<T>() {
641                     return self.dropless.alloc(value);
642                 }
643                 match <T as ArenaField<'tcx>>::arena(self) {
644                     ::std::option::Option::Some(arena) => arena.alloc(value),
645                     ::std::option::Option::None => unsafe { self.drop.alloc(value) },
646                 }
647             }
648
649             #[inline]
650             pub fn alloc_slice<T: ::std::marker::Copy>(&self, value: &[T]) -> &mut [T] {
651                 if value.is_empty() {
652                     return &mut [];
653                 }
654                 self.dropless.alloc_slice(value)
655             }
656
657             pub fn alloc_from_iter<'a, T: ArenaAllocatable>(
658                 &'a self,
659                 iter: impl ::std::iter::IntoIterator<Item = T>,
660             ) -> &'a mut [T] {
661                 if !::std::mem::needs_drop::<T>() {
662                     return self.dropless.alloc_from_iter(iter);
663                 }
664                 match <T as ArenaField<'tcx>>::arena(self) {
665                     ::std::option::Option::Some(arena) => arena.alloc_from_iter(iter),
666                     ::std::option::Option::None => unsafe { self.drop.alloc_from_iter(iter) },
667                 }
668             }
669         }
670     }
671 }
672
673 #[cfg(test)]
674 mod tests;