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::ptr::NonNull;
  65 use core::sync::atomic::{AtomicPtr, Ordering};
  66 use core::{mem, ptr};
  67
  68 use crate::sys_common::util::dumb_print;
  69
  70 #[stable(feature = "alloc_module", since = "1.28.0")]
  71 #[doc(inline)]
  72 pub use alloc_crate::alloc::*;
  73
  74 /// The default memory allocator provided by the operating system.
  75 ///
  76 /// This is based on `malloc` on Unix platforms and `HeapAlloc` on Windows,
  77 /// plus related functions.
  78 ///
  79 /// This type implements the `GlobalAlloc` trait and Rust programs by default
  80 /// work as if they had this definition:
  81 ///
  82 /// ```rust
  83 /// use std::alloc::System;
  84 ///
  85 /// #[global_allocator]
  86 /// static A: System = System;
  87 ///
  88 /// fn main() {
  89 ///     let a = Box::new(4); // Allocates from the system allocator.
  90 ///     println!("{}", a);
  91 /// }
  92 /// ```
  93 ///
  94 /// You can also define your own wrapper around `System` if you'd like, such as
  95 /// keeping track of the number of all bytes allocated:
  96 ///
  97 /// ```rust
  98 /// use std::alloc::{System, GlobalAlloc, Layout};
  99 /// use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
 100 ///
 101 /// struct Counter;
 102 ///
 103 /// static ALLOCATED: AtomicUsize = AtomicUsize::new(0);
 104 ///
 105 /// unsafe impl GlobalAlloc for Counter {
 106 ///     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
 107 ///         let ret = System.alloc(layout);
 108 ///         if !ret.is_null() {
 109 ///             ALLOCATED.fetch_add(layout.size(), SeqCst);
 110 ///         }
 111 ///         return ret
 112 ///     }
 113 ///
 114 ///     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
 115 ///         System.dealloc(ptr, layout);
 116 ///         ALLOCATED.fetch_sub(layout.size(), SeqCst);
 117 ///     }
 118 /// }
 119 ///
 120 /// #[global_allocator]
 121 /// static A: Counter = Counter;
 122 ///
 123 /// fn main() {
 124 ///     println!("allocated bytes before main: {}", ALLOCATED.load(SeqCst));
 125 /// }
 126 /// ```
 127 ///
 128 /// It can also be used directly to allocate memory independently of whatever
 129 /// global allocator has been selected for a Rust program. For example if a Rust
 130 /// program opts in to using jemalloc as the global allocator, `System` will
 131 /// still allocate memory using `malloc` and `HeapAlloc`.
 132 #[stable(feature = "alloc_system_type", since = "1.28.0")]
 133 #[derive(Debug, Default, Copy, Clone)]
 134 pub struct System;
 135
 136 // The AllocRef impl checks the layout size to be non-zero and forwards to the GlobalAlloc impl,
 137 // which is in `std::sys::*::alloc`.
 138 #[unstable(feature = "allocator_api", issue = "32838")]
 139 unsafe impl AllocRef for System {
 140     #[inline]
 141     fn alloc(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
 142         if layout.size() == 0 {
 143             Ok((layout.dangling(), 0))
 144         } else {
 145             unsafe {
 146                 NonNull::new(GlobalAlloc::alloc(self, layout))
 147                     .ok_or(AllocErr)
 148                     .map(|p| (p, layout.size()))
 149             }
 150         }
 151     }
 152
 153     #[inline]
 154     fn alloc_zeroed(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
 155         if layout.size() == 0 {
 156             Ok((layout.dangling(), 0))
 157         } else {
 158             unsafe {
 159                 NonNull::new(GlobalAlloc::alloc_zeroed(self, layout))
 160                     .ok_or(AllocErr)
 161                     .map(|p| (p, layout.size()))
 162             }
 163         }
 164     }
 165
 166     #[inline]
 167     unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
 168         if layout.size() != 0 {
 169             GlobalAlloc::dealloc(self, ptr.as_ptr(), layout)
 170         }
 171     }
 172
 173     #[inline]
 174     unsafe fn realloc(
 175         &mut self,
 176         ptr: NonNull<u8>,
 177         layout: Layout,
 178         new_size: usize,
 179     ) -> Result<(NonNull<u8>, usize), AllocErr> {
 180         match (layout.size(), new_size) {
 181             (0, 0) => Ok((layout.dangling(), 0)),
 182             (0, _) => self.alloc(Layout::from_size_align_unchecked(new_size, layout.align())),
 183             (_, 0) => {
 184                 self.dealloc(ptr, layout);
 185                 Ok((layout.dangling(), 0))
 186             }
 187             (_, _) => NonNull::new(GlobalAlloc::realloc(self, ptr.as_ptr(), layout, new_size))
 188                 .ok_or(AllocErr)
 189                 .map(|p| (p, new_size)),
 190         }
 191     }
 192 }
 193
 194 static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
 195
 196 /// Registers a custom allocation error hook, replacing any that was previously registered.
 197 ///
 198 /// The allocation error hook is invoked when an infallible memory allocation fails, before
 199 /// the runtime aborts. The default hook prints a message to standard error,
 200 /// but this behavior can be customized with the [`set_alloc_error_hook`] and
 201 /// [`take_alloc_error_hook`] functions.
 202 ///
 203 /// The hook is provided with a `Layout` struct which contains information
 204 /// about the allocation that failed.
 205 ///
 206 /// The allocation error hook is a global resource.
 207 ///
 208 /// [`set_alloc_error_hook`]: fn.set_alloc_error_hook.html
 209 /// [`take_alloc_error_hook`]: fn.take_alloc_error_hook.html
 210 #[unstable(feature = "alloc_error_hook", issue = "51245")]
 211 pub fn set_alloc_error_hook(hook: fn(Layout)) {
 212     HOOK.store(hook as *mut (), Ordering::SeqCst);
 213 }
 214
 215 /// Unregisters the current allocation error hook, returning it.
 216 ///
 217 /// *See also the function [`set_alloc_error_hook`].*
 218 ///
 219 /// If no custom hook is registered, the default hook will be returned.
 220 ///
 221 /// [`set_alloc_error_hook`]: fn.set_alloc_error_hook.html
 222 #[unstable(feature = "alloc_error_hook", issue = "51245")]
 223 pub fn take_alloc_error_hook() -> fn(Layout) {
 224     let hook = HOOK.swap(ptr::null_mut(), Ordering::SeqCst);
 225     if hook.is_null() { default_alloc_error_hook } else { unsafe { mem::transmute(hook) } }
 226 }
 227
 228 fn default_alloc_error_hook(layout: Layout) {
 229     dumb_print(format_args!("memory allocation of {} bytes failed", layout.size()));
 230 }
 231
 232 #[cfg(not(test))]
 233 #[doc(hidden)]
 234 #[alloc_error_handler]
 235 #[unstable(feature = "alloc_internals", issue = "none")]
 236 pub fn rust_oom(layout: Layout) -> ! {
 237     let hook = HOOK.load(Ordering::SeqCst);
 238     let hook: fn(Layout) =
 239         if hook.is_null() { default_alloc_error_hook } else { unsafe { mem::transmute(hook) } };
 240     hook(layout);
 241     unsafe {
 242         crate::sys::abort_internal();
 243     }
 244 }
 245
 246 #[cfg(not(test))]
 247 #[doc(hidden)]
 248 #[allow(unused_attributes)]
 249 #[unstable(feature = "alloc_internals", issue = "none")]
 250 pub mod __default_lib_allocator {
 251     use super::{GlobalAlloc, Layout, System};
 252     // These magic symbol names are used as a fallback for implementing the
 253     // `__rust_alloc` etc symbols (see `src/liballoc/alloc.rs) when there is
 254     // no `#[global_allocator]` attribute.
 255
 256     // for symbol names src/librustc/middle/allocator.rs
 257     // for signatures src/librustc_allocator/lib.rs
 258
 259     // linkage directives are provided as part of the current compiler allocator
 260     // ABI
 261
 262     #[rustc_std_internal_symbol]
 263     pub unsafe extern "C" fn __rdl_alloc(size: usize, align: usize) -> *mut u8 {
 264         let layout = Layout::from_size_align_unchecked(size, align);
 265         System.alloc(layout)
 266     }
 267
 268     #[rustc_std_internal_symbol]
 269     pub unsafe extern "C" fn __rdl_dealloc(ptr: *mut u8, size: usize, align: usize) {
 270         System.dealloc(ptr, Layout::from_size_align_unchecked(size, align))
 271     }
 272
 273     #[rustc_std_internal_symbol]
 274     pub unsafe extern "C" fn __rdl_realloc(
 275         ptr: *mut u8,
 276         old_size: usize,
 277         align: usize,
 278         new_size: usize,
 279     ) -> *mut u8 {
 280         let old_layout = Layout::from_size_align_unchecked(old_size, align);
 281         System.realloc(ptr, old_layout, new_size)
 282     }
 283
 284     #[rustc_std_internal_symbol]
 285     pub unsafe extern "C" fn __rdl_alloc_zeroed(size: usize, align: usize) -> *mut u8 {
 286         let layout = Layout::from_size_align_unchecked(size, align);
 287         System.alloc_zeroed(layout)
 288     }
 289 }