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

ab_contracts_io_type/
maybe_data.rs
