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