alloc.rs

  1// SPDX-License-Identifier: Apache-2.0 OR MIT
  2
  3#[cfg(all(feature = "alloc", not(feature = "std")))]
  4use alloc::{boxed::Box, sync::Arc};
  5#[cfg(feature = "alloc")]
  6use core::alloc::AllocError;
  7use core::{mem::MaybeUninit, pin::Pin};
  8#[cfg(feature = "std")]
  9use std::sync::Arc;
 10
 11#[cfg(not(feature = "alloc"))]
 12type AllocError = core::convert::Infallible;
 13
 14use crate::{
 15    init_from_closure, pin_init_from_closure, InPlaceWrite, Init, PinInit, ZeroableOption,
 16};
 17
 18pub extern crate alloc;
 19
 20// SAFETY: All zeros is equivalent to `None` (option layout optimization guarantee:
 21// <https://doc.rust-lang.org/stable/std/option/index.html#representation>).
 22unsafe impl<T> ZeroableOption for Box<T> {}
 23
 24/// Smart pointer that can initialize memory in-place.
 25pub trait InPlaceInit<T>: Sized {
 26    /// Use the given pin-initializer to pin-initialize a `T` inside of a new smart pointer of this
 27    /// type.
 28    ///
 29    /// If `T: !Unpin` it will not be able to move afterwards.
 30    fn try_pin_init<E>(init: impl PinInit<T, E>) -> Result<Pin<Self>, E>
 31    where
 32        E: From<AllocError>;
 33
 34    /// Use the given pin-initializer to pin-initialize a `T` inside of a new smart pointer of this
 35    /// type.
 36    ///
 37    /// If `T: !Unpin` it will not be able to move afterwards.
 38    fn pin_init(init: impl PinInit<T>) -> Result<Pin<Self>, AllocError> {
 39        // SAFETY: We delegate to `init` and only change the error type.
 40        let init = unsafe {
 41            pin_init_from_closure(|slot| match init.__pinned_init(slot) {
 42                Ok(()) => Ok(()),
 43                Err(i) => match i {},
 44            })
 45        };
 46        Self::try_pin_init(init)
 47    }
 48
 49    /// Use the given initializer to in-place initialize a `T`.
 50    fn try_init<E>(init: impl Init<T, E>) -> Result<Self, E>
 51    where
 52        E: From<AllocError>;
 53
 54    /// Use the given initializer to in-place initialize a `T`.
 55    fn init(init: impl Init<T>) -> Result<Self, AllocError> {
 56        // SAFETY: We delegate to `init` and only change the error type.
 57        let init = unsafe {
 58            init_from_closure(|slot| match init.__init(slot) {
 59                Ok(()) => Ok(()),
 60                Err(i) => match i {},
 61            })
 62        };
 63        Self::try_init(init)
 64    }
 65}
 66
 67#[cfg(feature = "alloc")]
 68macro_rules! try_new_uninit {
 69    ($type:ident) => {
 70        $type::try_new_uninit()?
 71    };
 72}
 73#[cfg(all(feature = "std", not(feature = "alloc")))]
 74macro_rules! try_new_uninit {
 75    ($type:ident) => {
 76        $type::new_uninit()
 77    };
 78}
 79
 80impl<T> InPlaceInit<T> for Box<T> {
 81    #[inline]
 82    fn try_pin_init<E>(init: impl PinInit<T, E>) -> Result<Pin<Self>, E>
 83    where
 84        E: From<AllocError>,
 85    {
 86        try_new_uninit!(Box).write_pin_init(init)
 87    }
 88
 89    #[inline]
 90    fn try_init<E>(init: impl Init<T, E>) -> Result<Self, E>
 91    where
 92        E: From<AllocError>,
 93    {
 94        try_new_uninit!(Box).write_init(init)
 95    }
 96}
 97
 98impl<T> InPlaceInit<T> for Arc<T> {
 99    #[inline]
100    fn try_pin_init<E>(init: impl PinInit<T, E>) -> Result<Pin<Self>, E>
101    where
102        E: From<AllocError>,
103    {
104        let mut this = try_new_uninit!(Arc);
105        let Some(slot) = Arc::get_mut(&mut this) else {
106            // SAFETY: the Arc has just been created and has no external references
107            unsafe { core::hint::unreachable_unchecked() }
108        };
109        let slot = slot.as_mut_ptr();
110        // SAFETY: When init errors/panics, slot will get deallocated but not dropped,
111        // slot is valid and will not be moved, because we pin it later.
112        unsafe { init.__pinned_init(slot)? };
113        // SAFETY: All fields have been initialized and this is the only `Arc` to that data.
114        Ok(unsafe { Pin::new_unchecked(this.assume_init()) })
115    }
116
117    #[inline]
118    fn try_init<E>(init: impl Init<T, E>) -> Result<Self, E>
119    where
120        E: From<AllocError>,
121    {
122        let mut this = try_new_uninit!(Arc);
123        let Some(slot) = Arc::get_mut(&mut this) else {
124            // SAFETY: the Arc has just been created and has no external references
125            unsafe { core::hint::unreachable_unchecked() }
126        };
127        let slot = slot.as_mut_ptr();
128        // SAFETY: When init errors/panics, slot will get deallocated but not dropped,
129        // slot is valid.
130        unsafe { init.__init(slot)? };
131        // SAFETY: All fields have been initialized.
132        Ok(unsafe { this.assume_init() })
133    }
134}
135
136impl<T> InPlaceWrite<T> for Box<MaybeUninit<T>> {
137    type Initialized = Box<T>;
138
139    fn write_init<E>(mut self, init: impl Init<T, E>) -> Result<Self::Initialized, E> {
140        let slot = self.as_mut_ptr();
141        // SAFETY: When init errors/panics, slot will get deallocated but not dropped,
142        // slot is valid.
143        unsafe { init.__init(slot)? };
144        // SAFETY: All fields have been initialized.
145        Ok(unsafe { self.assume_init() })
146    }
147
148    fn write_pin_init<E>(mut self, init: impl PinInit<T, E>) -> Result<Pin<Self::Initialized>, E> {
149        let slot = self.as_mut_ptr();
150        // SAFETY: When init errors/panics, slot will get deallocated but not dropped,
151        // slot is valid and will not be moved, because we pin it later.
152        unsafe { init.__pinned_init(slot)? };
153        // SAFETY: All fields have been initialized.
154        Ok(unsafe { self.assume_init() }.into())
155    }
156}