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