library/std/src/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 #![deny(unsafe_op_in_unsafe_fn)]
  63 #![stable(feature = "alloc_module", since = "1.28.0")]
  64
  65 use core::intrinsics;
  66 use core::ptr::NonNull;
  67 use core::sync::atomic::{AtomicPtr, Ordering};
  68 use core::{mem, ptr};
  69
  70 use crate::sys_common::util::dumb_print;
  71
  72 #[stable(feature = "alloc_module", since = "1.28.0")]
  73 #[doc(inline)]
  74 pub use alloc_crate::alloc::*;
  75
  76 /// The default memory allocator provided by the operating system.
  77 ///
  78 /// This is based on `malloc` on Unix platforms and `HeapAlloc` on Windows,
  79 /// plus related functions.
  80 ///
  81 /// This type implements the `GlobalAlloc` trait and Rust programs by default
  82 /// work as if they had this definition:
  83 ///
  84 /// ```rust
  85 /// use std::alloc::System;
  86 ///
  87 /// #[global_allocator]
  88 /// static A: System = System;
  89 ///
  90 /// fn main() {
  91 ///     let a = Box::new(4); // Allocates from the system allocator.
  92 ///     println!("{}", a);
  93 /// }
  94 /// ```
  95 ///
  96 /// You can also define your own wrapper around `System` if you'd like, such as
  97 /// keeping track of the number of all bytes allocated:
  98 ///
  99 /// ```rust
 100 /// use std::alloc::{System, GlobalAlloc, Layout};
 101 /// use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
 102 ///
 103 /// struct Counter;
 104 ///
 105 /// static ALLOCATED: AtomicUsize = AtomicUsize::new(0);
 106 ///
 107 /// unsafe impl GlobalAlloc for Counter {
 108 ///     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
 109 ///         let ret = System.alloc(layout);
 110 ///         if !ret.is_null() {
 111 ///             ALLOCATED.fetch_add(layout.size(), SeqCst);
 112 ///         }
 113 ///         return ret
 114 ///     }
 115 ///
 116 ///     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
 117 ///         System.dealloc(ptr, layout);
 118 ///         ALLOCATED.fetch_sub(layout.size(), SeqCst);
 119 ///     }
 120 /// }
 121 ///
 122 /// #[global_allocator]
 123 /// static A: Counter = Counter;
 124 ///
 125 /// fn main() {
 126 ///     println!("allocated bytes before main: {}", ALLOCATED.load(SeqCst));
 127 /// }
 128 /// ```
 129 ///
 130 /// It can also be used directly to allocate memory independently of whatever
 131 /// global allocator has been selected for a Rust program. For example if a Rust
 132 /// program opts in to using jemalloc as the global allocator, `System` will
 133 /// still allocate memory using `malloc` and `HeapAlloc`.
 134 #[stable(feature = "alloc_system_type", since = "1.28.0")]
 135 #[derive(Debug, Default, Copy, Clone)]
 136 pub struct System;
 137
 138 // The AllocRef impl checks the layout size to be non-zero and forwards to the GlobalAlloc impl,
 139 // which is in `std::sys::*::alloc`.
 140 #[unstable(feature = "allocator_api", issue = "32838")]
 141 unsafe impl AllocRef for System {
 142     #[inline]
 143     fn alloc(&mut self, layout: Layout) -> Result<NonNull<[u8]>, AllocErr> {
 144         let size = layout.size();
 145         let ptr = if size == 0 {
 146             layout.dangling()
 147         } else {
 148             // SAFETY: `layout` is non-zero in size,
 149             unsafe { NonNull::new(GlobalAlloc::alloc(&System, layout)).ok_or(AllocErr)? }
 150         };
 151         Ok(NonNull::slice_from_raw_parts(ptr, size))
 152     }
 153
 154     #[inline]
 155     fn alloc_zeroed(&mut self, layout: Layout) -> Result<NonNull<[u8]>, AllocErr> {
 156         let size = layout.size();
 157         let ptr = if size == 0 {
 158             layout.dangling()
 159         } else {
 160             // SAFETY: `layout` is non-zero in size,
 161             unsafe { NonNull::new(GlobalAlloc::alloc_zeroed(&System, layout)).ok_or(AllocErr)? }
 162         };
 163         Ok(NonNull::slice_from_raw_parts(ptr, size))
 164     }
 165
 166     #[inline]
 167     unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
 168         if layout.size() != 0 {
 169             // SAFETY: `layout` is non-zero in size,
 170             // other conditions must be upheld by the caller
 171             unsafe { GlobalAlloc::dealloc(&System, ptr.as_ptr(), layout) }
 172         }
 173     }
 174
 175     #[inline]
 176     unsafe fn grow(
 177         &mut self,
 178         ptr: NonNull<u8>,
 179         layout: Layout,
 180         new_size: usize,
 181     ) -> Result<NonNull<[u8]>, AllocErr> {
 182         debug_assert!(
 183             new_size >= layout.size(),
 184             "`new_size` must be greater than or equal to `layout.size()`"
 185         );
 186
 187         // SAFETY: `new_size` must be non-zero, which is checked in the match expression.
 188         // Other conditions must be upheld by the caller
 189         unsafe {
 190             match layout.size() {
 191                 old_size if old_size == new_size => {
 192                     Ok(NonNull::slice_from_raw_parts(ptr, new_size))
 193                 }
 194                 0 => self.alloc(Layout::from_size_align_unchecked(new_size, layout.align())),
 195                 old_size => {
 196                     // `realloc` probably checks for `new_size > size` or something similar.
 197                     intrinsics::assume(new_size > old_size);
 198                     let raw_ptr = GlobalAlloc::realloc(&System, ptr.as_ptr(), layout, new_size);
 199                     let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
 200                     Ok(NonNull::slice_from_raw_parts(ptr, new_size))
 201                 }
 202             }
 203         }
 204     }
 205
 206     #[inline]
 207     unsafe fn grow_zeroed(
 208         &mut self,
 209         ptr: NonNull<u8>,
 210         layout: Layout,
 211         new_size: usize,
 212     ) -> Result<NonNull<[u8]>, AllocErr> {
 213         debug_assert!(
 214             new_size >= layout.size(),
 215             "`new_size` must be greater than or equal to `layout.size()`"
 216         );
 217
 218         // SAFETY: `new_size` must be non-zero, which is checked in the match expression.
 219         // Other conditions must be upheld by the caller
 220         unsafe {
 221             match layout.size() {
 222                 old_size if old_size == new_size => {
 223                     Ok(NonNull::slice_from_raw_parts(ptr, new_size))
 224                 }
 225                 0 => self.alloc_zeroed(Layout::from_size_align_unchecked(new_size, layout.align())),
 226                 old_size => {
 227                     // `realloc` probably checks for `new_size > size` or something similar.
 228                     intrinsics::assume(new_size > old_size);
 229                     let raw_ptr = GlobalAlloc::realloc(&System, ptr.as_ptr(), layout, new_size);
 230                     raw_ptr.add(old_size).write_bytes(0, new_size - old_size);
 231                     let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
 232                     Ok(NonNull::slice_from_raw_parts(ptr, new_size))
 233                 }
 234             }
 235         }
 236     }
 237
 238     #[inline]
 239     unsafe fn shrink(
 240         &mut self,
 241         ptr: NonNull<u8>,
 242         layout: Layout,
 243         new_size: usize,
 244     ) -> Result<NonNull<[u8]>, AllocErr> {
 245         let old_size = layout.size();
 246         debug_assert!(
 247             new_size <= old_size,
 248             "`new_size` must be smaller than or equal to `layout.size()`"
 249         );
 250
 251         let ptr = if new_size == old_size {
 252             ptr
 253         } else if new_size == 0 {
 254             // SAFETY: `layout` is non-zero in size as `old_size` != `new_size`
 255             // Other conditions must be upheld by the caller
 256             unsafe {
 257                 self.dealloc(ptr, layout);
 258             }
 259             layout.dangling()
 260         } else {
 261             // SAFETY: new_size is not zero,
 262             // Other conditions must be upheld by the caller
 263             let raw_ptr = unsafe {
 264                 // `realloc` probably checks for `new_size < old_size` or something similar.
 265                 intrinsics::assume(new_size < old_size);
 266                 GlobalAlloc::realloc(&System, ptr.as_ptr(), layout, new_size)
 267             };
 268             NonNull::new(raw_ptr).ok_or(AllocErr)?
 269         };
 270
 271         Ok(NonNull::slice_from_raw_parts(ptr, new_size))
 272     }
 273 }
 274 static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
 275
 276 /// Registers a custom allocation error hook, replacing any that was previously registered.
 277 ///
 278 /// The allocation error hook is invoked when an infallible memory allocation fails, before
 279 /// the runtime aborts. The default hook prints a message to standard error,
 280 /// but this behavior can be customized with the [`set_alloc_error_hook`] and
 281 /// [`take_alloc_error_hook`] functions.
 282 ///
 283 /// The hook is provided with a `Layout` struct which contains information
 284 /// about the allocation that failed.
 285 ///
 286 /// The allocation error hook is a global resource.
 287 ///
 288 /// [`set_alloc_error_hook`]: fn.set_alloc_error_hook.html
 289 /// [`take_alloc_error_hook`]: fn.take_alloc_error_hook.html
 290 #[unstable(feature = "alloc_error_hook", issue = "51245")]
 291 pub fn set_alloc_error_hook(hook: fn(Layout)) {
 292     HOOK.store(hook as *mut (), Ordering::SeqCst);
 293 }
 294
 295 /// Unregisters the current allocation error hook, returning it.
 296 ///
 297 /// *See also the function [`set_alloc_error_hook`].*
 298 ///
 299 /// If no custom hook is registered, the default hook will be returned.
 300 ///
 301 /// [`set_alloc_error_hook`]: fn.set_alloc_error_hook.html
 302 #[unstable(feature = "alloc_error_hook", issue = "51245")]
 303 pub fn take_alloc_error_hook() -> fn(Layout) {
 304     let hook = HOOK.swap(ptr::null_mut(), Ordering::SeqCst);
 305     if hook.is_null() { default_alloc_error_hook } else { unsafe { mem::transmute(hook) } }
 306 }
 307
 308 fn default_alloc_error_hook(layout: Layout) {
 309     dumb_print(format_args!("memory allocation of {} bytes failed", layout.size()));
 310 }
 311
 312 #[cfg(not(test))]
 313 #[doc(hidden)]
 314 #[alloc_error_handler]
 315 #[unstable(feature = "alloc_internals", issue = "none")]
 316 pub fn rust_oom(layout: Layout) -> ! {
 317     let hook = HOOK.load(Ordering::SeqCst);
 318     let hook: fn(Layout) =
 319         if hook.is_null() { default_alloc_error_hook } else { unsafe { mem::transmute(hook) } };
 320     hook(layout);
 321     crate::process::abort()
 322 }
 323
 324 #[cfg(not(test))]
 325 #[doc(hidden)]
 326 #[allow(unused_attributes)]
 327 #[unstable(feature = "alloc_internals", issue = "none")]
 328 pub mod __default_lib_allocator {
 329     use super::{GlobalAlloc, Layout, System};
 330     // These magic symbol names are used as a fallback for implementing the
 331     // `__rust_alloc` etc symbols (see `src/liballoc/alloc.rs`) when there is
 332     // no `#[global_allocator]` attribute.
 333
 334     // for symbol names src/librustc_ast/expand/allocator.rs
 335     // for signatures src/librustc_allocator/lib.rs
 336
 337     // linkage directives are provided as part of the current compiler allocator
 338     // ABI
 339
 340     #[rustc_std_internal_symbol]
 341     pub unsafe extern "C" fn __rdl_alloc(size: usize, align: usize) -> *mut u8 {
 342         // SAFETY: see the guarantees expected by `Layout::from_size_align` and
 343         // `GlobalAlloc::alloc`.
 344         unsafe {
 345             let layout = Layout::from_size_align_unchecked(size, align);
 346             System.alloc(layout)
 347         }
 348     }
 349
 350     #[rustc_std_internal_symbol]
 351     pub unsafe extern "C" fn __rdl_dealloc(ptr: *mut u8, size: usize, align: usize) {
 352         // SAFETY: see the guarantees expected by `Layout::from_size_align` and
 353         // `GlobalAlloc::dealloc`.
 354         unsafe { System.dealloc(ptr, Layout::from_size_align_unchecked(size, align)) }
 355     }
 356
 357     #[rustc_std_internal_symbol]
 358     pub unsafe extern "C" fn __rdl_realloc(
 359         ptr: *mut u8,
 360         old_size: usize,
 361         align: usize,
 362         new_size: usize,
 363     ) -> *mut u8 {
 364         // SAFETY: see the guarantees expected by `Layout::from_size_align` and
 365         // `GlobalAlloc::realloc`.
 366         unsafe {
 367             let old_layout = Layout::from_size_align_unchecked(old_size, align);
 368             System.realloc(ptr, old_layout, new_size)
 369         }
 370     }
 371
 372     #[rustc_std_internal_symbol]
 373     pub unsafe extern "C" fn __rdl_alloc_zeroed(size: usize, align: usize) -> *mut u8 {
 374         // SAFETY: see the guarantees expected by `Layout::from_size_align` and
 375         // `GlobalAlloc::alloc_zeroed`.
 376         unsafe {
 377             let layout = Layout::from_size_align_unchecked(size, align);
 378             System.alloc_zeroed(layout)
 379         }
 380     }
 381 }