byte_slab/
managed_arc_slab.rs

1//! A convenience type for abstracting over slices or `SlabSliceArc`s.
2
3use core::{
4    ops::Deref,
5    fmt::Debug,
6    marker::PhantomData,
7};
8use serde::{Deserialize, Serialize};
9
10use crate::{
11    byte_slab::BSlab,
12    slab_arc::{SlabArc, RerooterKey},
13    slab_slice_arc::{SlabSliceArc, SlabStrArc}
14};
15
16/// A `ManagedArcSlab` may contain EITHER a borrowed `&[u8]` slice,
17/// or a `SlabSliceArc`. `ManagedArcSlab`s implement the `Deref` trait
18/// for access to the underlying data, and implement `serde`'s `Serialize`
19/// and `Deserialize` traits, to allow them to be serialized as a slice of
20/// bytes.
21#[derive(Clone)]
22pub enum ManagedArcSlab<'a, const N: usize, const SZ: usize> {
23    Borrowed(&'a [u8]),
24    Owned(SlabSliceArc<N, SZ>),
25}
26
27#[cfg(feature = "defmt")]
28impl<'a, const N: usize, const SZ: usize> defmt::Format for ManagedArcSlab<'a, N, SZ> {
29    fn format(&self, fmt: defmt::Formatter<'_>) {
30        self.deref().format(fmt)
31    }
32}
33
34impl<'a, const N: usize, const SZ: usize> Debug for ManagedArcSlab<'a, N, SZ> {
35    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
36        // TODO: Probably want a better debug impl than this
37        match self {
38            ManagedArcSlab::Borrowed(b) => b.fmt(f),
39            ManagedArcSlab::Owned(o) => o.deref().fmt(f),
40        }
41    }
42}
43
44impl<'a, const N: usize, const SZ: usize> Serialize for ManagedArcSlab<'a, N, SZ> {
45    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
46    where
47        S: serde::Serializer,
48    {
49        let data: &[u8] = self.deref();
50        data.serialize(serializer)
51    }
52}
53
54impl<'de: 'a, 'a, const N: usize, const SZ: usize> Deserialize<'de> for ManagedArcSlab<'a, N, SZ> {
55    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
56    where
57        D: serde::Deserializer<'de>,
58    {
59        struct ByteVisitor<'a, const N: usize, const SZ: usize> {
60            pd: PhantomData<&'a ()>,
61        }
62
63        impl<'d: 'ai, 'ai, const NI: usize, const SZI: usize> serde::de::Visitor<'d>
64            for ByteVisitor<'ai, NI, SZI>
65        {
66            type Value = ManagedArcSlab<'ai, NI, SZI>;
67
68            fn expecting(&self, formatter: &mut core::fmt::Formatter) -> core::fmt::Result {
69                write!(formatter, "a byte slice")
70            }
71
72            fn visit_borrowed_bytes<E>(self, v: &'d [u8]) -> Result<Self::Value, E>
73            where
74                E: serde::de::Error,
75            {
76                Ok(ManagedArcSlab::Borrowed(v))
77            }
78        }
79        deserializer.deserialize_bytes(ByteVisitor { pd: PhantomData })
80    }
81}
82
83impl<'a, const N: usize, const SZ: usize> Deref for ManagedArcSlab<'a, N, SZ> {
84    type Target = [u8];
85
86    fn deref(&self) -> &Self::Target {
87        match self {
88            ManagedArcSlab::Borrowed(data) => data,
89            ManagedArcSlab::Owned(ssa) => ssa.deref(),
90        }
91    }
92}
93
94impl<'a, const N: usize, const SZ: usize> ManagedArcSlab<'a, N, SZ> {
95    pub fn from_arc(arc: &SlabArc<N, SZ>) -> ManagedArcSlab<'static, N, SZ> {
96        ManagedArcSlab::Owned(arc.full_sub_slice_arc())
97    }
98
99    pub fn from_slice(sli: &'a [u8]) -> ManagedArcSlab<'a, N, SZ> {
100        ManagedArcSlab::Borrowed(sli)
101    }
102
103    pub fn from_slab_slice_arc(arc: &SlabSliceArc<N, SZ>) -> ManagedArcSlab<'static, N, SZ> {
104        ManagedArcSlab::Owned(arc.clone())
105    }
106
107    pub fn reroot_with_key(self, key: &RerooterKey) -> Option<ManagedArcSlab<'static, N, SZ>> {
108        match self {
109            ManagedArcSlab::Owned(e) => Some(ManagedArcSlab::Owned(e)),
110            ManagedArcSlab::Borrowed(b) => {
111                if b.is_empty() {
112                    // TODO: nuance
113                    return None;
114                }
115
116                // TODO: yolo ub
117                let start: usize = key.start as usize;
118                let end: usize = key.end as usize;
119                let b_start: usize = b.as_ptr() as usize;
120
121                if (start <= b_start) && (b_start < end) {
122                    let bslab: &'static BSlab<N, SZ> = unsafe { &*key.slab.cast::<BSlab<N, SZ>>() };
123
124                    // NOTE: We *don't* increase the refcount for this arc, as it technically aliases the
125                    // borrowed one that created the `key` here. Since we will forget this arc after
126                    // creating the ssa, don't update!
127                    let arc = SlabArc {
128                        slab: bslab,
129                        idx: key.idx,
130                    };
131
132                    let ssa = arc
133                        .sub_slice_arc(b_start - start, b.len())
134                        .ok()?;
135
136                    // Okay, now forget that arc ever happened
137                    core::mem::forget(arc);
138
139                    Some(ManagedArcSlab::Owned(ssa))
140                } else {
141                    None
142                }
143            }
144        }
145    }
146}
147
148/// A `ManagedArcStr` may contain EITHER a borrowed `&str` slice,
149/// or a `SlabStrArc`. `ManagedArcStr`s implement the `Deref` trait
150/// for access to the underlying data, and implement `serde`'s `Serialize`
151/// and `Deserialize` traits, to allow them to be serialized as a string slice.
152#[derive(Clone)]
153pub enum ManagedArcStr<'a, const N: usize, const SZ: usize> {
154    Borrowed(&'a str),
155    Owned(SlabStrArc<N, SZ>),
156}
157
158
159impl<'a, const N: usize, const SZ: usize> PartialEq<str> for ManagedArcStr<'a, N, SZ> {
160    fn eq(&self, other: &str) -> bool {
161        let stir: &str = self.deref();
162        stir.eq(other)
163    }
164}
165
166impl<'a, const N: usize, const SZ: usize> PartialEq for ManagedArcStr<'a, N, SZ> {
167    fn eq(&self, other: &Self) -> bool {
168        let stir_me: &str = self.deref();
169        let stir_ot: &str = other.deref();
170        stir_me.eq(stir_ot)
171    }
172}
173
174impl<'a, const N: usize, const SZ: usize> Eq for ManagedArcStr<'a, N, SZ> { }
175
176#[cfg(feature = "defmt")]
177impl<'a, const N: usize, const SZ: usize> defmt::Format for ManagedArcStr<'a, N, SZ> {
178    fn format(&self, fmt: defmt::Formatter<'_>) {
179        self.deref().format(fmt)
180    }
181}
182
183impl<'a, const N: usize, const SZ: usize> Debug for ManagedArcStr<'a, N, SZ> {
184    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
185        // TODO: Probably want a better debug impl than this
186        match self {
187            ManagedArcStr::Borrowed(b) => b.fmt(f),
188            ManagedArcStr::Owned(o) => o.deref().fmt(f),
189        }
190    }
191}
192
193impl<'a, const N: usize, const SZ: usize> Serialize for ManagedArcStr<'a, N, SZ> {
194    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
195    where
196        S: serde::Serializer,
197    {
198        let data: &str = self.deref();
199        data.serialize(serializer)
200    }
201}
202
203impl<'de: 'a, 'a, const N: usize, const SZ: usize> Deserialize<'de> for ManagedArcStr<'a, N, SZ> {
204    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
205    where
206        D: serde::Deserializer<'de>,
207    {
208        struct StrVisitor<'a, const N: usize, const SZ: usize> {
209            pd: PhantomData<&'a ()>,
210        }
211
212        impl<'d: 'ai, 'ai, const NI: usize, const SZI: usize> serde::de::Visitor<'d>
213            for StrVisitor<'ai, NI, SZI>
214        {
215            type Value = ManagedArcStr<'ai, NI, SZI>;
216
217            fn expecting(&self, formatter: &mut core::fmt::Formatter) -> core::fmt::Result {
218                write!(formatter, "a byte slice")
219            }
220
221            fn visit_borrowed_str<E>(self, v: &'d str) -> Result<Self::Value, E>
222            where
223                E: serde::de::Error,
224            {
225                Ok(ManagedArcStr::Borrowed(v))
226            }
227        }
228        deserializer.deserialize_str(StrVisitor { pd: PhantomData })
229    }
230}
231
232impl<'a, const N: usize, const SZ: usize> Deref for ManagedArcStr<'a, N, SZ> {
233    type Target = str;
234
235    fn deref(&self) -> &Self::Target {
236        match self {
237            ManagedArcStr::Borrowed(data) => data,
238            ManagedArcStr::Owned(ssa) => ssa.deref(),
239        }
240    }
241}
242
243impl<'a, const N: usize, const SZ: usize> ManagedArcStr<'a, N, SZ> {
244    pub fn from_slice(sli: &'a str) -> ManagedArcStr<'a, N, SZ> {
245        ManagedArcStr::Borrowed(sli)
246    }
247
248    pub fn from_slab_str_arc(arc: &SlabStrArc<N, SZ>) -> ManagedArcStr<'static, N, SZ> {
249        ManagedArcStr::Owned(arc.clone())
250    }
251
252    pub fn reroot_with_key(self, key: &RerooterKey) -> Option<ManagedArcStr<'static, N, SZ>> {
253        match self {
254            ManagedArcStr::Owned(e) => Some(ManagedArcStr::Owned(e)),
255            ManagedArcStr::Borrowed(b) => {
256                if b.is_empty() {
257                    // TODO: nuance
258                    return None;
259                }
260
261                // TODO: yolo ub
262                let start: usize = key.start as usize;
263                let end: usize = key.end as usize;
264                let b_start: usize = b.as_ptr() as usize;
265
266                if (start <= b_start) && (b_start < end) {
267                    let bslab: &'static BSlab<N, SZ> = unsafe { &*key.slab.cast::<BSlab<N, SZ>>() };
268
269                    // NOTE: We *don't* increase the refcount for this arc, as it technically aliases the
270                    // borrowed one that created the `key` here. Since we will forget this arc after
271                    // creating the ssa, don't update!
272                    let arc = SlabArc {
273                        slab: bslab,
274                        idx: key.idx,
275                    };
276
277                    let ssa = arc
278                        .sub_slice_arc(b_start - start, b.len())
279                        .ok()?
280                        .into_str_arc()
281                        .ok()?;
282
283                    // Okay, now forget that arc ever happened
284                    core::mem::forget(arc);
285
286                    Some(ManagedArcStr::Owned(ssa))
287                } else {
288                    None
289                }
290            }
291        }
292    }
293}
294
295/// A trait for restoring Managed types to a given SlabArc
296///
297/// For types that are not Managed, this trait should just return
298/// the type as-is.
299pub trait Reroot
300{
301    /// The return type. This is used to change the lifetimes, e.g. to convert
302    /// `Self: ManagedArcSlab<'a, _, _>` to `Retval: ManagedArcSlab<'static, _, _>`.
303    type Retval;
304
305    fn reroot(self, key: &RerooterKey) -> Result<Self::Retval, ()>;
306}
307
308impl<'a, const N: usize, const SZ: usize> Reroot for ManagedArcSlab<'a, N, SZ> {
309    type Retval = ManagedArcSlab<'static, N, SZ>;
310
311    fn reroot(self, key: &RerooterKey) -> Result<Self::Retval, ()>
312    {
313        self.reroot_with_key(key).ok_or(())
314    }
315}
316
317impl<'a, const N: usize, const SZ: usize> Reroot for ManagedArcStr<'a, N, SZ> {
318    type Retval = ManagedArcStr<'static, N, SZ>;
319
320    fn reroot(self, key: &RerooterKey) -> Result<Self::Retval, ()>
321    {
322        self.reroot_with_key(key).ok_or(())
323    }
324}
325
326macro_rules! reroot_nop {
327    (
328        [$($rrty:ty),+]
329    ) => {
330        $(
331            impl Reroot for $rrty {
332                type Retval = $rrty;
333
334                #[inline(always)]
335                fn reroot(self, _key: &RerooterKey) -> Result<Self::Retval, ()>
336                {
337                    Ok(self)
338                }
339            }
340        )+
341    };
342}
343
344reroot_nop!([u8, u16, u32, u64]);
345reroot_nop!([i8, i16, i32, i64]);
346reroot_nop!([bool, char, ()]);
347reroot_nop!([f32, f64]);
348
349impl<T, E> Reroot for Result<T, E>
350where
351    T: Reroot,
352    E: Reroot,
353{
354    type Retval = Result<T::Retval, E::Retval>;
355
356    fn reroot(self, key: &RerooterKey) -> Result<Result<T::Retval, E::Retval>, ()> {
357        match self {
358            Ok(t) => Ok(Ok(t.reroot(key)?)),
359            Err(e) => Ok(Err(e.reroot(key)?)),
360        }
361    }
362}
363
364impl<T> Reroot for Option<T>
365where
366    T: Reroot,
367{
368    type Retval = Option<T::Retval>;
369
370    #[inline]
371    fn reroot(self, key: &RerooterKey) -> Result<Self::Retval, ()> {
372        match self {
373            Some(t) => Ok(Some(t.reroot(key)?)),
374            None => Ok(None),
375        }
376    }
377}
378
379#[cfg(test)]
380mod test {
381    use crate::{BSlab, ManagedArcSlab, Reroot};
382    use std::ops::Deref;
383
384    #[test]
385    fn smoke() {
386        static SLAB: BSlab<4, 128> = BSlab::new();
387        SLAB.init().unwrap();
388
389        let mut sbox = SLAB.alloc_box().unwrap();
390
391        sbox[..4].copy_from_slice(&[1, 2, 3, 4]);
392
393        let arc_1 = sbox.into_arc();
394
395        let brw = ManagedArcSlab::<4, 128>::Borrowed(&arc_1[..4]);
396        let own: ManagedArcSlab<'static, 4, 128> = brw.reroot(&arc_1).unwrap();
397
398        match own {
399            ManagedArcSlab::Owned(ssa) => {
400                assert_eq!(&[1, 2, 3, 4], ssa.deref());
401            }
402            _ => panic!("Not owned!"),
403        }
404    }
405
406}