ublk: don't return 0 in case of any failure
[linux-block.git] / rust / alloc / alloc.rs
1 // SPDX-License-Identifier: Apache-2.0 OR MIT
2
3 //! Memory allocation APIs
4
5 #![stable(feature = "alloc_module", since = "1.28.0")]
6
7 #[cfg(not(test))]
8 use core::intrinsics;
9 use core::intrinsics::{min_align_of_val, size_of_val};
10
11 use core::ptr::Unique;
12 #[cfg(not(test))]
13 use core::ptr::{self, NonNull};
14
15 #[stable(feature = "alloc_module", since = "1.28.0")]
16 #[doc(inline)]
17 pub use core::alloc::*;
18
19 use core::marker::Destruct;
20
21 #[cfg(test)]
22 mod tests;
23
24 extern "Rust" {
25     // These are the magic symbols to call the global allocator.  rustc generates
26     // them to call `__rg_alloc` etc. if there is a `#[global_allocator]` attribute
27     // (the code expanding that attribute macro generates those functions), or to call
28     // the default implementations in libstd (`__rdl_alloc` etc. in `library/std/src/alloc.rs`)
29     // otherwise.
30     // The rustc fork of LLVM also special-cases these function names to be able to optimize them
31     // like `malloc`, `realloc`, and `free`, respectively.
32     #[rustc_allocator]
33     #[rustc_allocator_nounwind]
34     fn __rust_alloc(size: usize, align: usize) -> *mut u8;
35     #[rustc_allocator_nounwind]
36     fn __rust_dealloc(ptr: *mut u8, size: usize, align: usize);
37     #[rustc_allocator_nounwind]
38     fn __rust_realloc(ptr: *mut u8, old_size: usize, align: usize, new_size: usize) -> *mut u8;
39     #[rustc_allocator_nounwind]
40     fn __rust_alloc_zeroed(size: usize, align: usize) -> *mut u8;
41 }
42
43 /// The global memory allocator.
44 ///
45 /// This type implements the [`Allocator`] trait by forwarding calls
46 /// to the allocator registered with the `#[global_allocator]` attribute
47 /// if there is one, or the `std` crate’s default.
48 ///
49 /// Note: while this type is unstable, the functionality it provides can be
50 /// accessed through the [free functions in `alloc`](self#functions).
51 #[unstable(feature = "allocator_api", issue = "32838")]
52 #[derive(Copy, Clone, Default, Debug)]
53 #[cfg(not(test))]
54 pub struct Global;
55
56 #[cfg(test)]
57 pub use std::alloc::Global;
58
59 /// Allocate memory with the global allocator.
60 ///
61 /// This function forwards calls to the [`GlobalAlloc::alloc`] method
62 /// of the allocator registered with the `#[global_allocator]` attribute
63 /// if there is one, or the `std` crate’s default.
64 ///
65 /// This function is expected to be deprecated in favor of the `alloc` method
66 /// of the [`Global`] type when it and the [`Allocator`] trait become stable.
67 ///
68 /// # Safety
69 ///
70 /// See [`GlobalAlloc::alloc`].
71 ///
72 /// # Examples
73 ///
74 /// ```
75 /// use std::alloc::{alloc, dealloc, Layout};
76 ///
77 /// unsafe {
78 ///     let layout = Layout::new::<u16>();
79 ///     let ptr = alloc(layout);
80 ///
81 ///     *(ptr as *mut u16) = 42;
82 ///     assert_eq!(*(ptr as *mut u16), 42);
83 ///
84 ///     dealloc(ptr, layout);
85 /// }
86 /// ```
87 #[stable(feature = "global_alloc", since = "1.28.0")]
88 #[must_use = "losing the pointer will leak memory"]
89 #[inline]
90 pub unsafe fn alloc(layout: Layout) -> *mut u8 {
91     unsafe { __rust_alloc(layout.size(), layout.align()) }
92 }
93
94 /// Deallocate memory with the global allocator.
95 ///
96 /// This function forwards calls to the [`GlobalAlloc::dealloc`] method
97 /// of the allocator registered with the `#[global_allocator]` attribute
98 /// if there is one, or the `std` crate’s default.
99 ///
100 /// This function is expected to be deprecated in favor of the `dealloc` method
101 /// of the [`Global`] type when it and the [`Allocator`] trait become stable.
102 ///
103 /// # Safety
104 ///
105 /// See [`GlobalAlloc::dealloc`].
106 #[stable(feature = "global_alloc", since = "1.28.0")]
107 #[inline]
108 pub unsafe fn dealloc(ptr: *mut u8, layout: Layout) {
109     unsafe { __rust_dealloc(ptr, layout.size(), layout.align()) }
110 }
111
112 /// Reallocate memory with the global allocator.
113 ///
114 /// This function forwards calls to the [`GlobalAlloc::realloc`] method
115 /// of the allocator registered with the `#[global_allocator]` attribute
116 /// if there is one, or the `std` crate’s default.
117 ///
118 /// This function is expected to be deprecated in favor of the `realloc` method
119 /// of the [`Global`] type when it and the [`Allocator`] trait become stable.
120 ///
121 /// # Safety
122 ///
123 /// See [`GlobalAlloc::realloc`].
124 #[stable(feature = "global_alloc", since = "1.28.0")]
125 #[must_use = "losing the pointer will leak memory"]
126 #[inline]
127 pub unsafe fn realloc(ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
128     unsafe { __rust_realloc(ptr, layout.size(), layout.align(), new_size) }
129 }
130
131 /// Allocate zero-initialized memory with the global allocator.
132 ///
133 /// This function forwards calls to the [`GlobalAlloc::alloc_zeroed`] method
134 /// of the allocator registered with the `#[global_allocator]` attribute
135 /// if there is one, or the `std` crate’s default.
136 ///
137 /// This function is expected to be deprecated in favor of the `alloc_zeroed` method
138 /// of the [`Global`] type when it and the [`Allocator`] trait become stable.
139 ///
140 /// # Safety
141 ///
142 /// See [`GlobalAlloc::alloc_zeroed`].
143 ///
144 /// # Examples
145 ///
146 /// ```
147 /// use std::alloc::{alloc_zeroed, dealloc, Layout};
148 ///
149 /// unsafe {
150 ///     let layout = Layout::new::<u16>();
151 ///     let ptr = alloc_zeroed(layout);
152 ///
153 ///     assert_eq!(*(ptr as *mut u16), 0);
154 ///
155 ///     dealloc(ptr, layout);
156 /// }
157 /// ```
158 #[stable(feature = "global_alloc", since = "1.28.0")]
159 #[must_use = "losing the pointer will leak memory"]
160 #[inline]
161 pub unsafe fn alloc_zeroed(layout: Layout) -> *mut u8 {
162     unsafe { __rust_alloc_zeroed(layout.size(), layout.align()) }
163 }
164
165 #[cfg(not(test))]
166 impl Global {
167     #[inline]
168     fn alloc_impl(&self, layout: Layout, zeroed: bool) -> Result<NonNull<[u8]>, AllocError> {
169         match layout.size() {
170             0 => Ok(NonNull::slice_from_raw_parts(layout.dangling(), 0)),
171             // SAFETY: `layout` is non-zero in size,
172             size => unsafe {
173                 let raw_ptr = if zeroed { alloc_zeroed(layout) } else { alloc(layout) };
174                 let ptr = NonNull::new(raw_ptr).ok_or(AllocError)?;
175                 Ok(NonNull::slice_from_raw_parts(ptr, size))
176             },
177         }
178     }
179
180     // SAFETY: Same as `Allocator::grow`
181     #[inline]
182     unsafe fn grow_impl(
183         &self,
184         ptr: NonNull<u8>,
185         old_layout: Layout,
186         new_layout: Layout,
187         zeroed: bool,
188     ) -> Result<NonNull<[u8]>, AllocError> {
189         debug_assert!(
190             new_layout.size() >= old_layout.size(),
191             "`new_layout.size()` must be greater than or equal to `old_layout.size()`"
192         );
193
194         match old_layout.size() {
195             0 => self.alloc_impl(new_layout, zeroed),
196
197             // SAFETY: `new_size` is non-zero as `old_size` is greater than or equal to `new_size`
198             // as required by safety conditions. Other conditions must be upheld by the caller
199             old_size if old_layout.align() == new_layout.align() => unsafe {
200                 let new_size = new_layout.size();
201
202                 // `realloc` probably checks for `new_size >= old_layout.size()` or something similar.
203                 intrinsics::assume(new_size >= old_layout.size());
204
205                 let raw_ptr = realloc(ptr.as_ptr(), old_layout, new_size);
206                 let ptr = NonNull::new(raw_ptr).ok_or(AllocError)?;
207                 if zeroed {
208                     raw_ptr.add(old_size).write_bytes(0, new_size - old_size);
209                 }
210                 Ok(NonNull::slice_from_raw_parts(ptr, new_size))
211             },
212
213             // SAFETY: because `new_layout.size()` must be greater than or equal to `old_size`,
214             // both the old and new memory allocation are valid for reads and writes for `old_size`
215             // bytes. Also, because the old allocation wasn't yet deallocated, it cannot overlap
216             // `new_ptr`. Thus, the call to `copy_nonoverlapping` is safe. The safety contract
217             // for `dealloc` must be upheld by the caller.
218             old_size => unsafe {
219                 let new_ptr = self.alloc_impl(new_layout, zeroed)?;
220                 ptr::copy_nonoverlapping(ptr.as_ptr(), new_ptr.as_mut_ptr(), old_size);
221                 self.deallocate(ptr, old_layout);
222                 Ok(new_ptr)
223             },
224         }
225     }
226 }
227
228 #[unstable(feature = "allocator_api", issue = "32838")]
229 #[cfg(not(test))]
230 unsafe impl Allocator for Global {
231     #[inline]
232     fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
233         self.alloc_impl(layout, false)
234     }
235
236     #[inline]
237     fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
238         self.alloc_impl(layout, true)
239     }
240
241     #[inline]
242     unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
243         if layout.size() != 0 {
244             // SAFETY: `layout` is non-zero in size,
245             // other conditions must be upheld by the caller
246             unsafe { dealloc(ptr.as_ptr(), layout) }
247         }
248     }
249
250     #[inline]
251     unsafe fn grow(
252         &self,
253         ptr: NonNull<u8>,
254         old_layout: Layout,
255         new_layout: Layout,
256     ) -> Result<NonNull<[u8]>, AllocError> {
257         // SAFETY: all conditions must be upheld by the caller
258         unsafe { self.grow_impl(ptr, old_layout, new_layout, false) }
259     }
260
261     #[inline]
262     unsafe fn grow_zeroed(
263         &self,
264         ptr: NonNull<u8>,
265         old_layout: Layout,
266         new_layout: Layout,
267     ) -> Result<NonNull<[u8]>, AllocError> {
268         // SAFETY: all conditions must be upheld by the caller
269         unsafe { self.grow_impl(ptr, old_layout, new_layout, true) }
270     }
271
272     #[inline]
273     unsafe fn shrink(
274         &self,
275         ptr: NonNull<u8>,
276         old_layout: Layout,
277         new_layout: Layout,
278     ) -> Result<NonNull<[u8]>, AllocError> {
279         debug_assert!(
280             new_layout.size() <= old_layout.size(),
281             "`new_layout.size()` must be smaller than or equal to `old_layout.size()`"
282         );
283
284         match new_layout.size() {
285             // SAFETY: conditions must be upheld by the caller
286             0 => unsafe {
287                 self.deallocate(ptr, old_layout);
288                 Ok(NonNull::slice_from_raw_parts(new_layout.dangling(), 0))
289             },
290
291             // SAFETY: `new_size` is non-zero. Other conditions must be upheld by the caller
292             new_size if old_layout.align() == new_layout.align() => unsafe {
293                 // `realloc` probably checks for `new_size <= old_layout.size()` or something similar.
294                 intrinsics::assume(new_size <= old_layout.size());
295
296                 let raw_ptr = realloc(ptr.as_ptr(), old_layout, new_size);
297                 let ptr = NonNull::new(raw_ptr).ok_or(AllocError)?;
298                 Ok(NonNull::slice_from_raw_parts(ptr, new_size))
299             },
300
301             // SAFETY: because `new_size` must be smaller than or equal to `old_layout.size()`,
302             // both the old and new memory allocation are valid for reads and writes for `new_size`
303             // bytes. Also, because the old allocation wasn't yet deallocated, it cannot overlap
304             // `new_ptr`. Thus, the call to `copy_nonoverlapping` is safe. The safety contract
305             // for `dealloc` must be upheld by the caller.
306             new_size => unsafe {
307                 let new_ptr = self.allocate(new_layout)?;
308                 ptr::copy_nonoverlapping(ptr.as_ptr(), new_ptr.as_mut_ptr(), new_size);
309                 self.deallocate(ptr, old_layout);
310                 Ok(new_ptr)
311             },
312         }
313     }
314 }
315
316 /// The allocator for unique pointers.
317 #[cfg(all(not(no_global_oom_handling), not(test)))]
318 #[lang = "exchange_malloc"]
319 #[inline]
320 unsafe fn exchange_malloc(size: usize, align: usize) -> *mut u8 {
321     let layout = unsafe { Layout::from_size_align_unchecked(size, align) };
322     match Global.allocate(layout) {
323         Ok(ptr) => ptr.as_mut_ptr(),
324         Err(_) => handle_alloc_error(layout),
325     }
326 }
327
328 #[cfg_attr(not(test), lang = "box_free")]
329 #[inline]
330 #[rustc_const_unstable(feature = "const_box", issue = "92521")]
331 // This signature has to be the same as `Box`, otherwise an ICE will happen.
332 // When an additional parameter to `Box` is added (like `A: Allocator`), this has to be added here as
333 // well.
334 // For example if `Box` is changed to  `struct Box<T: ?Sized, A: Allocator>(Unique<T>, A)`,
335 // this function has to be changed to `fn box_free<T: ?Sized, A: Allocator>(Unique<T>, A)` as well.
336 pub(crate) const unsafe fn box_free<T: ?Sized, A: ~const Allocator + ~const Destruct>(
337     ptr: Unique<T>,
338     alloc: A,
339 ) {
340     unsafe {
341         let size = size_of_val(ptr.as_ref());
342         let align = min_align_of_val(ptr.as_ref());
343         let layout = Layout::from_size_align_unchecked(size, align);
344         alloc.deallocate(From::from(ptr.cast()), layout)
345     }
346 }
347
348 // # Allocation error handler
349
350 #[cfg(not(no_global_oom_handling))]
351 extern "Rust" {
352     // This is the magic symbol to call the global alloc error handler.  rustc generates
353     // it to call `__rg_oom` if there is a `#[alloc_error_handler]`, or to call the
354     // default implementations below (`__rdl_oom`) otherwise.
355     fn __rust_alloc_error_handler(size: usize, align: usize) -> !;
356 }
357
358 /// Abort on memory allocation error or failure.
359 ///
360 /// Callers of memory allocation APIs wishing to abort computation
361 /// in response to an allocation error are encouraged to call this function,
362 /// rather than directly invoking `panic!` or similar.
363 ///
364 /// The default behavior of this function is to print a message to standard error
365 /// and abort the process.
366 /// It can be replaced with [`set_alloc_error_hook`] and [`take_alloc_error_hook`].
367 ///
368 /// [`set_alloc_error_hook`]: ../../std/alloc/fn.set_alloc_error_hook.html
369 /// [`take_alloc_error_hook`]: ../../std/alloc/fn.take_alloc_error_hook.html
370 #[stable(feature = "global_alloc", since = "1.28.0")]
371 #[rustc_const_unstable(feature = "const_alloc_error", issue = "92523")]
372 #[cfg(all(not(no_global_oom_handling), not(test)))]
373 #[cold]
374 pub const fn handle_alloc_error(layout: Layout) -> ! {
375     const fn ct_error(_: Layout) -> ! {
376         panic!("allocation failed");
377     }
378
379     fn rt_error(layout: Layout) -> ! {
380         unsafe {
381             __rust_alloc_error_handler(layout.size(), layout.align());
382         }
383     }
384
385     unsafe { core::intrinsics::const_eval_select((layout,), ct_error, rt_error) }
386 }
387
388 // For alloc test `std::alloc::handle_alloc_error` can be used directly.
389 #[cfg(all(not(no_global_oom_handling), test))]
390 pub use std::alloc::handle_alloc_error;
391
392 #[cfg(all(not(no_global_oom_handling), not(test)))]
393 #[doc(hidden)]
394 #[allow(unused_attributes)]
395 #[unstable(feature = "alloc_internals", issue = "none")]
396 pub mod __alloc_error_handler {
397     use crate::alloc::Layout;
398
399     // called via generated `__rust_alloc_error_handler`
400
401     // if there is no `#[alloc_error_handler]`
402     #[rustc_std_internal_symbol]
403     pub unsafe extern "C-unwind" fn __rdl_oom(size: usize, _align: usize) -> ! {
404         panic!("memory allocation of {size} bytes failed")
405     }
406
407     // if there is an `#[alloc_error_handler]`
408     #[rustc_std_internal_symbol]
409     pub unsafe extern "C-unwind" fn __rg_oom(size: usize, align: usize) -> ! {
410         let layout = unsafe { Layout::from_size_align_unchecked(size, align) };
411         extern "Rust" {
412             #[lang = "oom"]
413             fn oom_impl(layout: Layout) -> !;
414         }
415         unsafe { oom_impl(layout) }
416     }
417 }
418
419 /// Specialize clones into pre-allocated, uninitialized memory.
420 /// Used by `Box::clone` and `Rc`/`Arc::make_mut`.
421 pub(crate) trait WriteCloneIntoRaw: Sized {
422     unsafe fn write_clone_into_raw(&self, target: *mut Self);
423 }
424
425 impl<T: Clone> WriteCloneIntoRaw for T {
426     #[inline]
427     default unsafe fn write_clone_into_raw(&self, target: *mut Self) {
428         // Having allocated *first* may allow the optimizer to create
429         // the cloned value in-place, skipping the local and move.
430         unsafe { target.write(self.clone()) };
431     }
432 }
433
434 impl<T: Copy> WriteCloneIntoRaw for T {
435     #[inline]
436     unsafe fn write_clone_into_raw(&self, target: *mut Self) {
437         // We can always copy in-place, without ever involving a local value.
438         unsafe { target.copy_from_nonoverlapping(self, 1) };
439     }
440 }