src/libstd/alloc.rs

   1 //! Memory allocation APIs
   2 //!
   3 //! In a given program, the standard library has one “global” memory allocator
   4 //! that is used for example by `Box<T>` and `Vec<T>`.
   5 //!
   6 //! Currently the default global allocator is unspecified. Libraries, however,
   7 //! like `cdylib`s and `staticlib`s are guaranteed to use the [`System`] by
   8 //! default.
   9 //!
  10 //! [`System`]: struct.System.html
  11 //!
  12 //! # The `#[global_allocator]` attribute
  13 //!
  14 //! This attribute allows configuring the choice of global allocator.
  15 //! You can use this to implement a completely custom global allocator
  16 //! to route all default allocation requests to a custom object.
  17 //!
  18 //! ```rust
  19 //! use std::alloc::{GlobalAlloc, System, Layout};
  20 //!
  21 //! struct MyAllocator;
  22 //!
  23 //! unsafe impl GlobalAlloc for MyAllocator {
  24 //!     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
  25 //!         System.alloc(layout)
  26 //!     }
  27 //!
  28 //!     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
  29 //!         System.dealloc(ptr, layout)
  30 //!     }
  31 //! }
  32 //!
  33 //! #[global_allocator]
  34 //! static GLOBAL: MyAllocator = MyAllocator;
  35 //!
  36 //! fn main() {
  37 //!     // This `Vec` will allocate memory through `GLOBAL` above
  38 //!     let mut v = Vec::new();
  39 //!     v.push(1);
  40 //! }
  41 //! ```
  42 //!
  43 //! The attribute is used on a `static` item whose type implements the
  44 //! [`GlobalAlloc`] trait. This type can be provided by an external library:
  45 //!
  46 //! [`GlobalAlloc`]: ../../core/alloc/trait.GlobalAlloc.html
  47 //!
  48 //! ```rust,ignore (demonstrates crates.io usage)
  49 //! extern crate jemallocator;
  50 //!
  51 //! use jemallocator::Jemalloc;
  52 //!
  53 //! #[global_allocator]
  54 //! static GLOBAL: Jemalloc = Jemalloc;
  55 //!
  56 //! fn main() {}
  57 //! ```
  58 //!
  59 //! The `#[global_allocator]` can only be used once in a crate
  60 //! or its recursive dependencies.
  61
  62 #![stable(feature = "alloc_module", since = "1.28.0")]
  63
  64 use core::intrinsics;
  65 use core::ptr::NonNull;
  66 use core::sync::atomic::{AtomicPtr, Ordering};
  67 use core::{mem, ptr};
  68
  69 use crate::sys_common::util::dumb_print;
  70
  71 #[stable(feature = "alloc_module", since = "1.28.0")]
  72 #[doc(inline)]
  73 pub use alloc_crate::alloc::*;
  74
  75 /// The default memory allocator provided by the operating system.
  76 ///
  77 /// This is based on `malloc` on Unix platforms and `HeapAlloc` on Windows,
  78 /// plus related functions.
  79 ///
  80 /// This type implements the `GlobalAlloc` trait and Rust programs by default
  81 /// work as if they had this definition:
  82 ///
  83 /// ```rust
  84 /// use std::alloc::System;
  85 ///
  86 /// #[global_allocator]
  87 /// static A: System = System;
  88 ///
  89 /// fn main() {
  90 ///     let a = Box::new(4); // Allocates from the system allocator.
  91 ///     println!("{}", a);
  92 /// }
  93 /// ```
  94 ///
  95 /// You can also define your own wrapper around `System` if you'd like, such as
  96 /// keeping track of the number of all bytes allocated:
  97 ///
  98 /// ```rust
  99 /// use std::alloc::{System, GlobalAlloc, Layout};
 100 /// use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
 101 ///
 102 /// struct Counter;
 103 ///
 104 /// static ALLOCATED: AtomicUsize = AtomicUsize::new(0);
 105 ///
 106 /// unsafe impl GlobalAlloc for Counter {
 107 ///     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
 108 ///         let ret = System.alloc(layout);
 109 ///         if !ret.is_null() {
 110 ///             ALLOCATED.fetch_add(layout.size(), SeqCst);
 111 ///         }
 112 ///         return ret
 113 ///     }
 114 ///
 115 ///     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
 116 ///         System.dealloc(ptr, layout);
 117 ///         ALLOCATED.fetch_sub(layout.size(), SeqCst);
 118 ///     }
 119 /// }
 120 ///
 121 /// #[global_allocator]
 122 /// static A: Counter = Counter;
 123 ///
 124 /// fn main() {
 125 ///     println!("allocated bytes before main: {}", ALLOCATED.load(SeqCst));
 126 /// }
 127 /// ```
 128 ///
 129 /// It can also be used directly to allocate memory independently of whatever
 130 /// global allocator has been selected for a Rust program. For example if a Rust
 131 /// program opts in to using jemalloc as the global allocator, `System` will
 132 /// still allocate memory using `malloc` and `HeapAlloc`.
 133 #[stable(feature = "alloc_system_type", since = "1.28.0")]
 134 #[derive(Debug, Default, Copy, Clone)]
 135 pub struct System;
 136
 137 #[unstable(feature = "allocator_api", issue = "32838")]
 138 unsafe impl AllocRef for System {
 139     #[inline]
 140     fn alloc(&mut self, layout: Layout, init: AllocInit) -> Result<MemoryBlock, AllocErr> {
 141         unsafe {
 142             if layout.size() == 0 {
 143                 Ok(MemoryBlock::new(layout.dangling(), layout))
 144             } else {
 145                 let raw_ptr = match init {
 146                     AllocInit::Uninitialized => GlobalAlloc::alloc(self, layout),
 147                     AllocInit::Zeroed => GlobalAlloc::alloc_zeroed(self, layout),
 148                 };
 149                 let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
 150                 Ok(MemoryBlock::new(ptr, layout))
 151             }
 152         }
 153     }
 154
 155     #[inline]
 156     unsafe fn dealloc(&mut self, memory: MemoryBlock) {
 157         if memory.size() != 0 {
 158             GlobalAlloc::dealloc(self, memory.ptr().as_ptr(), memory.layout())
 159         }
 160     }
 161
 162     #[inline]
 163     unsafe fn grow(
 164         &mut self,
 165         memory: &mut MemoryBlock,
 166         new_size: usize,
 167         placement: ReallocPlacement,
 168         init: AllocInit,
 169     ) -> Result<(), AllocErr> {
 170         let old_size = memory.size();
 171         debug_assert!(
 172             new_size >= old_size,
 173             "`new_size` must be greater than or equal to `memory.size()`"
 174         );
 175
 176         if old_size == new_size {
 177             return Ok(());
 178         }
 179
 180         let new_layout = Layout::from_size_align_unchecked(new_size, memory.align());
 181         match placement {
 182             ReallocPlacement::InPlace => return Err(AllocErr),
 183             ReallocPlacement::MayMove if memory.size() == 0 => {
 184                 *memory = self.alloc(new_layout, init)?
 185             }
 186             ReallocPlacement::MayMove => {
 187                 // `realloc` probably checks for `new_size > old_size` or something similar.
 188                 intrinsics::assume(new_size > old_size);
 189                 let ptr =
 190                     GlobalAlloc::realloc(self, memory.ptr().as_ptr(), memory.layout(), new_size);
 191                 *memory = MemoryBlock::new(NonNull::new(ptr).ok_or(AllocErr)?, new_layout);
 192                 memory.init_offset(init, old_size);
 193             }
 194         }
 195         Ok(())
 196     }
 197
 198     #[inline]
 199     unsafe fn shrink(
 200         &mut self,
 201         memory: &mut MemoryBlock,
 202         new_size: usize,
 203         placement: ReallocPlacement,
 204     ) -> Result<(), AllocErr> {
 205         let old_size = memory.size();
 206         debug_assert!(
 207             new_size <= old_size,
 208             "`new_size` must be smaller than or equal to `memory.size()`"
 209         );
 210
 211         if old_size == new_size {
 212             return Ok(());
 213         }
 214
 215         let new_layout = Layout::from_size_align_unchecked(new_size, memory.align());
 216         match placement {
 217             ReallocPlacement::InPlace => return Err(AllocErr),
 218             ReallocPlacement::MayMove if new_size == 0 => {
 219                 let new_memory = MemoryBlock::new(new_layout.dangling(), new_layout);
 220                 let old_memory = mem::replace(memory, new_memory);
 221                 self.dealloc(old_memory)
 222             }
 223             ReallocPlacement::MayMove => {
 224                 // `realloc` probably checks for `new_size < old_size` or something similar.
 225                 intrinsics::assume(new_size < old_size);
 226                 let ptr =
 227                     GlobalAlloc::realloc(self, memory.ptr().as_ptr(), memory.layout(), new_size);
 228                 *memory = MemoryBlock::new(NonNull::new(ptr).ok_or(AllocErr)?, new_layout);
 229             }
 230         }
 231         Ok(())
 232     }
 233 }
 234
 235 static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
 236
 237 /// Registers a custom allocation error hook, replacing any that was previously registered.
 238 ///
 239 /// The allocation error hook is invoked when an infallible memory allocation fails, before
 240 /// the runtime aborts. The default hook prints a message to standard error,
 241 /// but this behavior can be customized with the [`set_alloc_error_hook`] and
 242 /// [`take_alloc_error_hook`] functions.
 243 ///
 244 /// The hook is provided with a `Layout` struct which contains information
 245 /// about the allocation that failed.
 246 ///
 247 /// The allocation error hook is a global resource.
 248 ///
 249 /// [`set_alloc_error_hook`]: fn.set_alloc_error_hook.html
 250 /// [`take_alloc_error_hook`]: fn.take_alloc_error_hook.html
 251 #[unstable(feature = "alloc_error_hook", issue = "51245")]
 252 pub fn set_alloc_error_hook(hook: fn(Layout)) {
 253     HOOK.store(hook as *mut (), Ordering::SeqCst);
 254 }
 255
 256 /// Unregisters the current allocation error hook, returning it.
 257 ///
 258 /// *See also the function [`set_alloc_error_hook`].*
 259 ///
 260 /// If no custom hook is registered, the default hook will be returned.
 261 ///
 262 /// [`set_alloc_error_hook`]: fn.set_alloc_error_hook.html
 263 #[unstable(feature = "alloc_error_hook", issue = "51245")]
 264 pub fn take_alloc_error_hook() -> fn(Layout) {
 265     let hook = HOOK.swap(ptr::null_mut(), Ordering::SeqCst);
 266     if hook.is_null() { default_alloc_error_hook } else { unsafe { mem::transmute(hook) } }
 267 }
 268
 269 fn default_alloc_error_hook(layout: Layout) {
 270     dumb_print(format_args!("memory allocation of {} bytes failed", layout.size()));
 271 }
 272
 273 #[cfg(not(test))]
 274 #[doc(hidden)]
 275 #[alloc_error_handler]
 276 #[unstable(feature = "alloc_internals", issue = "none")]
 277 pub fn rust_oom(layout: Layout) -> ! {
 278     let hook = HOOK.load(Ordering::SeqCst);
 279     let hook: fn(Layout) =
 280         if hook.is_null() { default_alloc_error_hook } else { unsafe { mem::transmute(hook) } };
 281     hook(layout);
 282     unsafe { crate::sys::abort_internal() }
 283 }
 284
 285 #[cfg(not(test))]
 286 #[doc(hidden)]
 287 #[allow(unused_attributes)]
 288 #[unstable(feature = "alloc_internals", issue = "none")]
 289 pub mod __default_lib_allocator {
 290     use super::{GlobalAlloc, Layout, System};
 291     // These magic symbol names are used as a fallback for implementing the
 292     // `__rust_alloc` etc symbols (see `src/liballoc/alloc.rs) when there is
 293     // no `#[global_allocator]` attribute.
 294
 295     // for symbol names src/librustc/middle/allocator.rs
 296     // for signatures src/librustc_allocator/lib.rs
 297
 298     // linkage directives are provided as part of the current compiler allocator
 299     // ABI
 300
 301     #[rustc_std_internal_symbol]
 302     pub unsafe extern "C" fn __rdl_alloc(size: usize, align: usize) -> *mut u8 {
 303         let layout = Layout::from_size_align_unchecked(size, align);
 304         System.alloc(layout)
 305     }
 306
 307     #[rustc_std_internal_symbol]
 308     pub unsafe extern "C" fn __rdl_dealloc(ptr: *mut u8, size: usize, align: usize) {
 309         System.dealloc(ptr, Layout::from_size_align_unchecked(size, align))
 310     }
 311
 312     #[rustc_std_internal_symbol]
 313     pub unsafe extern "C" fn __rdl_realloc(
 314         ptr: *mut u8,
 315         old_size: usize,
 316         align: usize,
 317         new_size: usize,
 318     ) -> *mut u8 {
 319         let old_layout = Layout::from_size_align_unchecked(old_size, align);
 320         System.realloc(ptr, old_layout, new_size)
 321     }
 322
 323     #[rustc_std_internal_symbol]
 324     pub unsafe extern "C" fn __rdl_alloc_zeroed(size: usize, align: usize) -> *mut u8 {
 325         let layout = Layout::from_size_align_unchecked(size, align);
 326         System.alloc_zeroed(layout)
 327     }
 328 }