xref: /linux/rust/pin-init/src/alloc.rs (revision 0bd2f269ae892ce7283fee8fcfe2c6c971d871bc)
1 // SPDX-License-Identifier: Apache-2.0 OR MIT
2 
3 #[cfg(all(feature = "alloc", not(feature = "std")))]
4 use alloc::{boxed::Box, sync::Arc};
5 #[cfg(feature = "alloc")]
6 use core::alloc::AllocError;
7 use core::{mem::MaybeUninit, pin::Pin};
8 #[cfg(feature = "std")]
9 use std::sync::Arc;
10 
11 #[cfg(not(feature = "alloc"))]
12 type AllocError = core::convert::Infallible;
13 
14 use crate::{
15     init_from_closure, pin_init_from_closure, InPlaceWrite, Init, PinInit, ZeroableOption,
16 };
17 
18 pub 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>).
22 unsafe impl<T> ZeroableOption for Box<T> {}
23 
24 /// Smart pointer that can initialize memory in-place.
25 pub 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.
try_pin_init<E>(init: impl PinInit<T, E>) -> Result<Pin<Self>, E> where E: From<AllocError>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.
pin_init(init: impl PinInit<T>) -> Result<Pin<Self>, AllocError>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`.
try_init<E>(init: impl Init<T, E>) -> Result<Self, E> where E: From<AllocError>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`.
init(init: impl Init<T>) -> Result<Self, AllocError>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")]
68 macro_rules! try_new_uninit {
69     ($type:ident) => {
70         $type::try_new_uninit()?
71     };
72 }
73 #[cfg(all(feature = "std", not(feature = "alloc")))]
74 macro_rules! try_new_uninit {
75     ($type:ident) => {
76         $type::new_uninit()
77     };
78 }
79 
80 impl<T> InPlaceInit<T> for Box<T> {
81     #[inline]
try_pin_init<E>(init: impl PinInit<T, E>) -> Result<Pin<Self>, E> where E: From<AllocError>,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]
try_init<E>(init: impl Init<T, E>) -> Result<Self, E> where E: From<AllocError>,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 
98 impl<T> InPlaceInit<T> for Arc<T> {
99     #[inline]
try_pin_init<E>(init: impl PinInit<T, E>) -> Result<Pin<Self>, E> where E: From<AllocError>,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]
try_init<E>(init: impl Init<T, E>) -> Result<Self, E> where E: From<AllocError>,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 
136 impl<T> InPlaceWrite<T> for Box<MaybeUninit<T>> {
137     type Initialized = Box<T>;
138 
write_init<E>(mut self, init: impl Init<T, E>) -> Result<Self::Initialized, E>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 
write_pin_init<E>(mut self, init: impl PinInit<T, E>) -> Result<Pin<Self::Initialized>, E>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 }
157