src/libstd/alloc.rs

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