ab_contracts_io_type/
maybe_data.rs

1use crate::trivial_type::TrivialType;
2use crate::{DerefWrapper, IoType, IoTypeOptional};
3use core::mem::MaybeUninit;
4use core::ops::{Deref, DerefMut};
5use core::ptr::NonNull;
6
7/// Wrapper type for `Data` that may or may not be filled with contents.
8///
9/// This is somewhat similar to [`VariableBytes`](crate::variable_bytes::VariableBytes), but instead
10/// of variable size data structure allows to either have it or not have the contents or not have
11/// it, which is a simpler and more convenient API that is also sufficient in many cases.
12#[derive(Debug)]
13pub struct MaybeData<Data>
14where
15    Data: TrivialType,
16{
17    data: NonNull<Data>,
18    size: NonNull<u32>,
19    capacity: u32,
20}
21
22unsafe impl<Data> IoType for MaybeData<Data>
23where
24    Data: TrivialType,
25{
26    const METADATA: &[u8] = Data::METADATA;
27
28    type PointerType = Data;
29
30    #[inline(always)]
31    fn size(&self) -> u32 {
32        // SAFETY: guaranteed to be initialized by constructors
33        unsafe { self.size.read() }
34    }
35
36    #[inline(always)]
37    unsafe fn size_ptr(&self) -> impl Deref<Target = NonNull<u32>> {
38        DerefWrapper(self.size)
39    }
40
41    #[inline(always)]
42    unsafe fn size_mut_ptr(&mut self) -> impl DerefMut<Target = *mut u32> {
43        DerefWrapper(self.size.as_ptr())
44    }
45
46    #[inline(always)]
47    fn capacity(&self) -> u32 {
48        self.size()
49    }
50
51    #[inline(always)]
52    unsafe fn capacity_ptr(&self) -> impl Deref<Target = NonNull<u32>> {
53        DerefWrapper(NonNull::from_ref(&self.capacity))
54    }
55
56    #[inline(always)]
57    #[track_caller]
58    unsafe fn set_size(&mut self, size: u32) {
59        debug_assert!(
60            size == 0 || size == self.size(),
61            "`set_size` called with invalid input {size} (self size {})",
62            self.size()
63        );
64
65        // SAFETY: guaranteed to be initialized by constructors
66        unsafe {
67            self.size.write(size);
68        }
69    }
70
71    #[inline(always)]
72    #[track_caller]
73    unsafe fn from_ptr<'a>(
74        data: &'a NonNull<Self::PointerType>,
75        size: &'a u32,
76        capacity: u32,
77    ) -> impl Deref<Target = Self> + 'a {
78        debug_assert!(data.is_aligned(), "Misaligned pointer");
79        debug_assert!(
80            *size == 0 || *size == capacity,
81            "Invalid size {size} for capacity {capacity}"
82        );
83        debug_assert!(
84            capacity >= Data::SIZE,
85            "Invalid capacity {capacity} for size {}",
86            Data::SIZE
87        );
88
89        let size = NonNull::from_ref(size);
90
91        DerefWrapper(MaybeData {
92            data: *data,
93            size,
94            capacity,
95        })
96    }
97
98    #[inline(always)]
99    #[track_caller]
100    unsafe fn from_mut_ptr<'a>(
101        data: &'a mut NonNull<Self::PointerType>,
102        size: &'a mut *mut u32,
103        capacity: u32,
104    ) -> impl DerefMut<Target = Self> + 'a {
105        debug_assert!(!size.is_null(), "`null` pointer for non-`TrivialType` size");
106        // SAFETY: Must be guaranteed by the caller + debug check above
107        let size = unsafe { NonNull::new_unchecked(*size) };
108        debug_assert!(data.is_aligned(), "Misaligned pointer");
109        // SAFETY: Must be guaranteed by the caller
110        {
111            let size = unsafe { size.read() };
112            debug_assert!(
113                size == 0 || size == capacity,
114                "Invalid size {size} for capacity {capacity}"
115            );
116        }
117        debug_assert!(
118            capacity >= Data::SIZE,
119            "Invalid capacity {capacity} for size {}",
120            Data::SIZE
121        );
122
123        DerefWrapper(MaybeData {
124            data: *data,
125            size,
126            capacity,
127        })
128    }
129
130    #[inline(always)]
131    unsafe fn as_ptr(&self) -> impl Deref<Target = NonNull<Self::PointerType>> {
132        &self.data
133    }
134
135    #[inline(always)]
136    unsafe fn as_mut_ptr(&mut self) -> impl DerefMut<Target = NonNull<Self::PointerType>> {
137        &mut self.data
138    }
139}
140
141impl<Data> IoTypeOptional for MaybeData<Data> where Data: TrivialType {}
142
143impl<Data> MaybeData<Data>
144where
145    Data: TrivialType,
146{
147    /// Create a new shared instance from provided data reference.
148    //
149    // `impl Deref` is used to tie lifetime of returned value to inputs, but still treat it as a
150    // shared reference for most practical purposes.
151    pub const fn from_buffer(data: Option<&'_ Data>) -> impl Deref<Target = Self> + '_ {
152        let (data, size) = if let Some(data) = data {
153            (NonNull::from_ref(data), &Data::SIZE)
154        } else {
155            (NonNull::dangling(), &0)
156        };
157
158        DerefWrapper(Self {
159            data,
160            size: NonNull::from_ref(size),
161            capacity: Data::SIZE,
162        })
163    }
164
165    /// Create a new exclusive instance from provided data reference.
166    ///
167    /// `size` can be either `0` or `Data::SIZE`, indicating that value is missing or present
168    /// accordingly.
169    ///
170    /// # Panics
171    /// Panics if `size != 0 && size != Data::SIZE`
172    //
173    // `impl DerefMut` is used to tie lifetime of returned value to inputs, but still treat it as an
174    // exclusive reference for most practical purposes.
175    #[track_caller]
176    pub fn from_buffer_mut<'a>(
177        buffer: &'a mut Data,
178        size: &'a mut u32,
179    ) -> impl DerefMut<Target = Self> + 'a {
180        debug_assert!(
181            *size == 0 || *size == Data::SIZE,
182            "Invalid size {size} (self size {})",
183            Data::SIZE
184        );
185
186        DerefWrapper(Self {
187            data: NonNull::from_mut(buffer),
188            size: NonNull::from_ref(size),
189            capacity: Data::SIZE,
190        })
191    }
192
193    /// Create a new shared instance from provided memory buffer.
194    ///
195    /// `size` must be `0`.
196    ///
197    /// # Panics
198    /// Panics if `size != 0`
199    //
200    // `impl Deref` is used to tie lifetime of returned value to inputs, but still treat it as a
201    // shared reference for most practical purposes.
202    // TODO: Change `usize` to `u32` once stabilized `generic_const_exprs` feature allows us to do
203    //  `CAPACITY as usize`
204    #[track_caller]
205    pub fn from_uninit<'a>(
206        uninit: &'a mut MaybeUninit<Data>,
207        size: &'a mut u32,
208    ) -> impl DerefMut<Target = Self> + 'a {
209        debug_assert_eq!(*size, 0, "Invalid size");
210
211        DerefWrapper(Self {
212            data: NonNull::from_mut(uninit).cast::<Data>(),
213            size: NonNull::from_mut(size),
214            capacity: Data::SIZE,
215        })
216    }
217
218    /// Try to get access to initialized `Data`, returns `None` if not initialized
219    #[inline(always)]
220    pub const fn get(&self) -> Option<&Data> {
221        // SAFETY: guaranteed to be initialized by constructors
222        if unsafe { self.size.read() } == self.capacity {
223            // SAFETY: initialized
224            Some(unsafe { self.data.as_ref() })
225        } else {
226            None
227        }
228    }
229
230    /// Try to get exclusive access to initialized `Data`, returns `None` if not initialized
231    #[inline(always)]
232    pub fn get_mut(&mut self) -> Option<&mut Data> {
233        // SAFETY: guaranteed to be initialized by constructors
234        if unsafe { self.size.read() } == self.capacity {
235            // SAFETY: initialized
236            Some(unsafe { self.data.as_mut() })
237        } else {
238            None
239        }
240    }
241
242    /// Initialize by inserting `Data` by value or replace existing value and return reference to it
243    #[inline(always)]
244    pub fn replace(&mut self, data: Data) -> &mut Data {
245        // SAFETY: guaranteed to be initialized by constructors
246        unsafe {
247            self.size.write(self.capacity);
248        }
249        // SAFETY: constructor guarantees that memory is aligned
250        unsafe {
251            self.data.write(data);
252            self.data.as_mut()
253        }
254    }
255
256    /// Remove `Data` inside and turn instance back into uninitialized
257    #[inline(always)]
258    pub fn remove(&mut self) {
259        // SAFETY: guaranteed to be initialized by constructors
260        unsafe {
261            self.size.write(0);
262        }
263    }
264
265    /// Get exclusive access to initialized `Data`, running provided initialization function if
266    /// necessary
267    #[inline(always)]
268    pub fn get_mut_or_init_with<'a, Init>(&'a mut self, init: Init) -> &'a mut Data
269    where
270        Init: FnOnce(NonNull<Data>) -> &'a mut Data,
271    {
272        // SAFETY: guaranteed to be initialized by constructors
273        if unsafe { self.size.read() } == self.capacity {
274            // SAFETY: initialized
275            unsafe { self.data.as_mut() }
276        } else {
277            let data = init(self.data);
278            // SAFETY: guaranteed to be initialized by constructors
279            unsafe {
280                self.size.write(self.capacity);
281            }
282            data
283        }
284    }
285
286    /// Assume value is initialized
287    ///
288    /// # Safety
289    /// Caller must ensure `Data` is actually properly initialized
290    #[inline(always)]
291    pub unsafe fn assume_init(&mut self) -> &mut Data {
292        // SAFETY: guaranteed to be initialized by constructors
293        unsafe {
294            self.size.write(self.capacity);
295        }
296        // SAFETY: guaranteed to be initialized by caller, the rest of guarantees are provided by
297        // constructors
298        unsafe { self.data.as_mut() }
299    }
300}
ab_contracts_io_type/maybe_data.rs

ab_contracts_io_type/
maybe_data.rs
