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