ex3_ic_stable_structures/
storable.rs

1use std::borrow::{Borrow, Cow};
2use std::cmp::{Ordering, Reverse};
3use std::convert::{TryFrom, TryInto};
4use std::fmt;
5
6#[cfg(test)]
7mod tests;
8
9/// A trait with convenience methods for storing an element into a stable structure.
10pub trait Storable {
11    /// Converts an element into bytes.
12    ///
13    /// NOTE: `Cow` is used here to avoid unnecessary cloning.
14    fn to_bytes(&self) -> Cow<[u8]>;
15
16    /// Converts bytes into an element.
17    fn from_bytes(bytes: Cow<[u8]>) -> Self;
18
19    /// The size bounds of the type.
20    const BOUND: Bound;
21
22    /// Like `to_bytes`, but includes additional checks to ensure the element's serialized bytes
23    /// are within the element's bounds.
24    fn to_bytes_checked(&self) -> Cow<[u8]> {
25        let bytes = self.to_bytes();
26        if let Bound::Bounded {
27            max_size,
28            is_fixed_size,
29        } = Self::BOUND
30        {
31            if is_fixed_size {
32                assert_eq!(
33                    bytes.len(),
34                    max_size as usize,
35                    "expected a fixed-size element with length {} bytes, but found {} bytes",
36                    max_size,
37                    bytes.len()
38                );
39            } else {
40                assert!(
41                    bytes.len() <= max_size as usize,
42                    "expected an element with length <= {} bytes, but found {} bytes",
43                    max_size,
44                    bytes.len()
45                );
46            }
47        }
48        bytes
49    }
50}
51
52/// States whether the type's size is bounded or unbounded.
53pub enum Bound {
54    /// The type has no size bounds.
55    Unbounded,
56
57    /// The type has size bounds.
58    Bounded {
59        /// The maximum size, in bytes, of the type when serialized.
60        max_size: u32,
61
62        /// True if all the values of this type have fixed-width encoding.
63        /// Some data structures, such as stable vector, can take
64        /// advantage of fixed size to avoid storing an explicit entry
65        /// size.
66        ///
67        /// Examples: little-/big-endian encoding of u16/u32/u64, tuples
68        /// and arrays of fixed-size types.
69        is_fixed_size: bool,
70    },
71}
72
73impl Bound {
74    /// Returns the maximum size of the type if bounded, panics if unbounded.
75    pub const fn max_size(&self) -> u32 {
76        if let Bound::Bounded { max_size, .. } = self {
77            *max_size
78        } else {
79            panic!("Cannot get max size of unbounded type.");
80        }
81    }
82
83    /// Returns true if the type is fixed in size, false otherwise.
84    pub const fn is_fixed_size(&self) -> bool {
85        if let Bound::Bounded { is_fixed_size, .. } = self {
86            *is_fixed_size
87        } else {
88            false
89        }
90    }
91}
92
93/// Variable-size, but limited in capacity byte array.
94#[derive(Eq, Copy, Clone)]
95pub struct Blob<const N: usize> {
96    storage: [u8; N],
97    size: u32,
98}
99
100impl<const N: usize> Blob<N> {
101    /// Returns the contents of this array as a byte slice.
102    pub fn as_slice(&self) -> &[u8] {
103        &self.storage[0..self.len()]
104    }
105
106    /// Returns true if the array is empty.
107    pub fn is_empty(&self) -> bool {
108        self.len() == 0
109    }
110
111    /// Returns the actual length of this array.
112    pub fn len(&self) -> usize {
113        self.size as usize
114    }
115}
116
117impl<const N: usize> Default for Blob<N> {
118    fn default() -> Self {
119        Self {
120            storage: [0; N],
121            size: 0,
122        }
123    }
124}
125
126#[derive(Debug)]
127pub struct TryFromSliceError;
128
129impl<const N: usize> TryFrom<&[u8]> for Blob<N> {
130    type Error = TryFromSliceError;
131
132    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
133        if value.len() > N {
134            return Err(TryFromSliceError);
135        }
136        let mut result = Self::default();
137        result.storage[0..value.len()].copy_from_slice(value);
138        result.size = value.len() as u32;
139        Ok(result)
140    }
141}
142
143impl<const N: usize> AsRef<[u8]> for Blob<N> {
144    fn as_ref(&self) -> &[u8] {
145        self.as_slice()
146    }
147}
148
149impl<const N: usize> PartialEq for Blob<N> {
150    fn eq(&self, other: &Self) -> bool {
151        self.as_slice().eq(other.as_slice())
152    }
153}
154
155impl<const N: usize> PartialOrd for Blob<N> {
156    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
157        self.as_slice().partial_cmp(other.as_slice())
158    }
159}
160
161impl<const N: usize> Ord for Blob<N> {
162    fn cmp(&self, other: &Self) -> Ordering {
163        self.as_slice().cmp(other.as_slice())
164    }
165}
166
167impl<const N: usize> fmt::Debug for Blob<N> {
168    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
169        self.as_slice().fmt(fmt)
170    }
171}
172
173impl<const N: usize> Storable for Blob<N> {
174    fn to_bytes(&self) -> Cow<[u8]> {
175        Cow::Borrowed(self.as_slice())
176    }
177
178    fn from_bytes(bytes: Cow<[u8]>) -> Self {
179        Self::try_from(bytes.borrow()).unwrap()
180    }
181
182    const BOUND: Bound = Bound::Bounded {
183        max_size: N as u32,
184        is_fixed_size: false,
185    };
186}
187
188// NOTE: Below are a few implementations of `Storable` for common types.
189// Some of these implementations use `unwrap`, as opposed to returning a `Result`
190// with a possible error. The reason behind this decision is that these
191// `unwrap`s would be triggered in one of the following cases:
192//
193// 1) The implementation of `Storable` has a bug.
194// 2) The data being stored in the stable structure is corrupt.
195//
196// Both of these errors are irrecoverable at runtime, and given the additional
197// complexity of exposing these errors in the API of stable structures, an `unwrap`
198// in case of a detected error is preferable and safer.
199
200impl Storable for () {
201    fn to_bytes(&self) -> Cow<[u8]> {
202        Cow::Borrowed(&[])
203    }
204
205    fn from_bytes(bytes: Cow<[u8]>) -> Self {
206        assert!(bytes.is_empty());
207    }
208
209    const BOUND: Bound = Bound::Bounded {
210        max_size: 0,
211        // A `()` should in theory be fixed in size, but this flag was initially
212        // set incorrectly and it cannot be fixed to maintain backward-compatibility.
213        is_fixed_size: false,
214    };
215}
216
217impl Storable for Vec<u8> {
218    fn to_bytes(&self) -> Cow<[u8]> {
219        Cow::Borrowed(self)
220    }
221
222    fn from_bytes(bytes: Cow<[u8]>) -> Self {
223        bytes.to_vec()
224    }
225
226    const BOUND: Bound = Bound::Unbounded;
227}
228
229impl Storable for String {
230    fn to_bytes(&self) -> Cow<[u8]> {
231        Cow::Borrowed(self.as_bytes())
232    }
233
234    fn from_bytes(bytes: Cow<[u8]>) -> Self {
235        String::from_utf8(bytes.to_vec()).unwrap()
236    }
237
238    const BOUND: Bound = Bound::Unbounded;
239}
240
241impl Storable for u128 {
242    fn to_bytes(&self) -> Cow<[u8]> {
243        Cow::Owned(self.to_be_bytes().to_vec())
244    }
245
246    fn from_bytes(bytes: Cow<[u8]>) -> Self {
247        Self::from_be_bytes(bytes.as_ref().try_into().unwrap())
248    }
249
250    const BOUND: Bound = Bound::Bounded {
251        max_size: 16,
252        is_fixed_size: true,
253    };
254}
255
256impl Storable for u64 {
257    fn to_bytes(&self) -> Cow<[u8]> {
258        Cow::Owned(self.to_be_bytes().to_vec())
259    }
260
261    fn from_bytes(bytes: Cow<[u8]>) -> Self {
262        Self::from_be_bytes(bytes.as_ref().try_into().unwrap())
263    }
264
265    const BOUND: Bound = Bound::Bounded {
266        max_size: 8,
267        is_fixed_size: true,
268    };
269}
270
271impl Storable for f64 {
272    fn to_bytes(&self) -> Cow<[u8]> {
273        Cow::Owned(self.to_be_bytes().to_vec())
274    }
275
276    fn from_bytes(bytes: Cow<[u8]>) -> Self {
277        Self::from_be_bytes(bytes.as_ref().try_into().unwrap())
278    }
279
280    const BOUND: Bound = Bound::Bounded {
281        max_size: 8,
282        is_fixed_size: true,
283    };
284}
285
286impl Storable for u32 {
287    fn to_bytes(&self) -> Cow<[u8]> {
288        Cow::Owned(self.to_be_bytes().to_vec())
289    }
290
291    fn from_bytes(bytes: Cow<[u8]>) -> Self {
292        Self::from_be_bytes(bytes.as_ref().try_into().unwrap())
293    }
294
295    const BOUND: Bound = Bound::Bounded {
296        max_size: 4,
297        is_fixed_size: true,
298    };
299}
300
301impl Storable for f32 {
302    fn to_bytes(&self) -> Cow<[u8]> {
303        Cow::Owned(self.to_be_bytes().to_vec())
304    }
305
306    fn from_bytes(bytes: Cow<[u8]>) -> Self {
307        Self::from_be_bytes(bytes.as_ref().try_into().unwrap())
308    }
309
310    const BOUND: Bound = Bound::Bounded {
311        max_size: 4,
312        is_fixed_size: true,
313    };
314}
315
316impl Storable for u16 {
317    fn to_bytes(&self) -> Cow<[u8]> {
318        Cow::Owned(self.to_be_bytes().to_vec())
319    }
320
321    fn from_bytes(bytes: Cow<[u8]>) -> Self {
322        Self::from_be_bytes(bytes.as_ref().try_into().unwrap())
323    }
324
325    const BOUND: Bound = Bound::Bounded {
326        max_size: 2,
327        is_fixed_size: true,
328    };
329}
330
331impl Storable for u8 {
332    fn to_bytes(&self) -> Cow<[u8]> {
333        Cow::Owned(self.to_be_bytes().to_vec())
334    }
335
336    fn from_bytes(bytes: Cow<[u8]>) -> Self {
337        Self::from_be_bytes(bytes.as_ref().try_into().unwrap())
338    }
339
340    const BOUND: Bound = Bound::Bounded {
341        max_size: 1,
342        is_fixed_size: true,
343    };
344}
345
346impl<const N: usize> Storable for [u8; N] {
347    fn to_bytes(&self) -> Cow<[u8]> {
348        Cow::Borrowed(&self[..])
349    }
350
351    fn from_bytes(bytes: Cow<[u8]>) -> Self {
352        assert_eq!(bytes.len(), N);
353        let mut arr = [0; N];
354        arr[0..N].copy_from_slice(&bytes);
355        arr
356    }
357
358    const BOUND: Bound = Bound::Bounded {
359        max_size: N as u32,
360        is_fixed_size: true,
361    };
362}
363
364impl<T: Storable> Storable for Reverse<T> {
365    fn to_bytes(&self) -> Cow<[u8]> {
366        self.0.to_bytes()
367    }
368
369    fn from_bytes(bytes: Cow<[u8]>) -> Self {
370        Self(T::from_bytes(bytes))
371    }
372
373    const BOUND: Bound = T::BOUND;
374}
375
376impl<A, B> Storable for (A, B)
377where
378    A: Storable,
379    B: Storable,
380{
381    fn to_bytes(&self) -> Cow<[u8]> {
382        match Self::BOUND {
383            Bound::Bounded { max_size, .. } => {
384                let mut bytes = vec![0; max_size as usize];
385                let a_bytes = self.0.to_bytes();
386                let b_bytes = self.1.to_bytes();
387
388                let a_bounds = bounds::<A>();
389                let b_bounds = bounds::<B>();
390
391                let a_max_size = a_bounds.max_size as usize;
392                let b_max_size = b_bounds.max_size as usize;
393
394                debug_assert!(a_bytes.len() <= a_max_size);
395                debug_assert!(b_bytes.len() <= b_max_size);
396
397                bytes[0..a_bytes.len()].copy_from_slice(a_bytes.borrow());
398                bytes[a_max_size..a_max_size + b_bytes.len()].copy_from_slice(b_bytes.borrow());
399
400                let a_size_len = bytes_to_store_size(&a_bounds) as usize;
401                let b_size_len = bytes_to_store_size(&b_bounds) as usize;
402
403                let sizes_offset: usize = a_max_size + b_max_size;
404
405                encode_size(
406                    &mut bytes[sizes_offset..sizes_offset + a_size_len],
407                    a_bytes.len(),
408                    &a_bounds,
409                );
410                encode_size(
411                    &mut bytes[sizes_offset + a_size_len..sizes_offset + a_size_len + b_size_len],
412                    b_bytes.len(),
413                    &b_bounds,
414                );
415
416                Cow::Owned(bytes)
417            }
418            _ => todo!("Serializing tuples with unbounded types is not yet supported."),
419        }
420    }
421
422    fn from_bytes(bytes: Cow<[u8]>) -> Self {
423        match Self::BOUND {
424            Bound::Bounded { max_size, .. } => {
425                assert_eq!(bytes.len(), max_size as usize);
426
427                let a_bounds = bounds::<A>();
428                let b_bounds = bounds::<B>();
429                let a_max_size = a_bounds.max_size as usize;
430                let b_max_size = b_bounds.max_size as usize;
431                let sizes_offset = a_max_size + b_max_size;
432
433                let a_size_len = bytes_to_store_size(&a_bounds) as usize;
434                let b_size_len = bytes_to_store_size(&b_bounds) as usize;
435                let a_len = decode_size(&bytes[sizes_offset..sizes_offset + a_size_len], &a_bounds);
436                let b_len = decode_size(
437                    &bytes[sizes_offset + a_size_len..sizes_offset + a_size_len + b_size_len],
438                    &b_bounds,
439                );
440
441                let a = A::from_bytes(Cow::Borrowed(&bytes[0..a_len]));
442                let b = B::from_bytes(Cow::Borrowed(&bytes[a_max_size..a_max_size + b_len]));
443
444                (a, b)
445            }
446            _ => todo!("Deserializing tuples with unbounded types is not yet supported."),
447        }
448    }
449
450    const BOUND: Bound = {
451        match (A::BOUND, B::BOUND) {
452            (Bound::Bounded { .. }, Bound::Bounded { .. }) => {
453                let a_bounds = bounds::<A>();
454                let b_bounds = bounds::<B>();
455
456                let max_size = a_bounds.max_size
457                    + b_bounds.max_size
458                    + bytes_to_store_size(&a_bounds)
459                    + bytes_to_store_size(&b_bounds);
460
461                let is_fixed_size = a_bounds.is_fixed_size && b_bounds.is_fixed_size;
462
463                Bound::Bounded {
464                    max_size,
465                    is_fixed_size,
466                }
467            }
468            _ => Bound::Unbounded,
469        }
470    };
471}
472
473impl<T: Storable> Storable for Option<T> {
474    fn to_bytes(&self) -> Cow<[u8]> {
475        match self {
476            Some(t) => {
477                let mut bytes = t.to_bytes().into_owned();
478                bytes.push(1);
479                Cow::Owned(bytes)
480            }
481            None => Cow::Borrowed(&[0]),
482        }
483    }
484
485    fn from_bytes(bytes: Cow<[u8]>) -> Self {
486        match bytes.split_last() {
487            Some((last, rest)) => match last {
488                0 => {
489                    assert!(rest.is_empty(), "Invalid Option encoding: unexpected prefix before the None marker: {rest:?}");
490                    None
491                }
492                1 => Some(T::from_bytes(Cow::Borrowed(rest))),
493                _ => panic!("Invalid Option encoding: unexpected variant marker {last}"),
494            },
495            None => panic!("Invalid Option encoding: expected at least one byte"),
496        }
497    }
498
499    const BOUND: Bound = {
500        match T::BOUND {
501            Bound::Bounded {
502                max_size,
503                is_fixed_size,
504            } => Bound::Bounded {
505                max_size: max_size + 1,
506                is_fixed_size,
507            },
508            Bound::Unbounded => Bound::Unbounded,
509        }
510    };
511}
512
513pub(crate) struct Bounds {
514    pub max_size: u32,
515    pub is_fixed_size: bool,
516}
517
518/// Returns the bounds of the given type, panics if unbounded.
519pub(crate) const fn bounds<A: Storable>() -> Bounds {
520    if let Bound::Bounded {
521        max_size,
522        is_fixed_size,
523    } = A::BOUND
524    {
525        Bounds {
526            max_size,
527            is_fixed_size,
528        }
529    } else {
530        panic!("Cannot get bounds of unbounded type.");
531    }
532}
533
534fn decode_size(src: &[u8], bounds: &Bounds) -> usize {
535    if bounds.is_fixed_size {
536        bounds.max_size as usize
537    } else if bounds.max_size <= u8::MAX as u32 {
538        src[0] as usize
539    } else if bounds.max_size <= u16::MAX as u32 {
540        u16::from_be_bytes([src[0], src[1]]) as usize
541    } else {
542        u32::from_be_bytes([src[0], src[1], src[2], src[3]]) as usize
543    }
544}
545
546fn encode_size(dst: &mut [u8], n: usize, bounds: &Bounds) {
547    if bounds.is_fixed_size {
548        return;
549    }
550
551    if bounds.max_size <= u8::MAX as u32 {
552        dst[0] = n as u8;
553    } else if bounds.max_size <= u16::MAX as u32 {
554        dst[0..2].copy_from_slice(&(n as u16).to_be_bytes());
555    } else {
556        dst[0..4].copy_from_slice(&(n as u32).to_be_bytes());
557    }
558}
559
560pub(crate) const fn bytes_to_store_size(bounds: &Bounds) -> u32 {
561    if bounds.is_fixed_size {
562        0
563    } else if bounds.max_size <= u8::MAX as u32 {
564        1
565    } else if bounds.max_size <= u16::MAX as u32 {
566        2
567    } else {
568        4
569    }
570}
ex3_ic_stable_structures/storable.rs

ex3_ic_stable_structures/
storable.rs
