ic_stable_memory/collections/vec/
mod.rs

1use crate::collections::vec::iter::SVecIter;
2use crate::encoding::{AsFixedSizeBytes, Buffer};
3use crate::mem::allocator::EMPTY_PTR;
4use crate::mem::s_slice::SSlice;
5use crate::mem::StablePtr;
6use crate::primitive::s_ref::SRef;
7use crate::primitive::s_ref_mut::SRefMut;
8use crate::primitive::StableType;
9use crate::{allocate, deallocate, reallocate, OutOfMemory};
10use std::cmp::Ordering;
11use std::fmt::{Debug, Formatter};
12use std::marker::PhantomData;
13
14#[doc(hidden)]
15pub mod iter;
16
17const DEFAULT_CAPACITY: usize = 4;
18
19/// Stable analog of [Vec]
20///
21/// May reallocate on inserts. In this case will copy the underlying data to a new location.
22///
23/// This is a "finite" data structure, it can only holp up to [u32::MAX] / `T::SIZE` elements. Putting
24/// more elements inside will panic.
25///
26/// `T` has to implement both [StableType] and [AsFixedSizeBytes]. [SVec] itself implements these
27/// traits and can be nested inside other stable data structures.
28///
29/// When [SVec] is stable-dropped, its elements are also stable-dropped but in reverse order.
30pub struct SVec<T: StableType + AsFixedSizeBytes> {
31    ptr: u64,
32    len: usize,
33    cap: usize,
34    stable_drop_flag: bool,
35    _marker_t: PhantomData<T>,
36}
37
38impl<T: StableType + AsFixedSizeBytes> SVec<T> {
39    /// Creates a [SVec] of capacity equal to 4 elements.
40    ///
41    /// Does not allocate any heap or stable memory.
42    ///
43    /// # Example
44    /// ```rust
45    /// // won't allocate until you insert something in it
46    /// # use ic_stable_memory::collections::SVec;
47    /// # use ic_stable_memory::stable_memory_init;
48    /// # unsafe { ic_stable_memory::mem::clear(); }
49    /// # stable_memory_init();
50    /// let mut numbers = SVec::<u64>::new();
51    /// ```
52    #[inline]
53    pub fn new() -> Self {
54        Self {
55            len: 0,
56            cap: DEFAULT_CAPACITY,
57            ptr: EMPTY_PTR,
58            stable_drop_flag: true,
59            _marker_t: PhantomData::default(),
60        }
61    }
62
63    /// Creates a [SVec] of requested capacity.
64    ///
65    /// Does allocate stable memory, returning [OutOfMemory] if there is not enough of it.
66    /// If this function returns [Ok], you are guaranteed to have enough stable memory to store at
67    /// least `capacity` elements in it.
68    ///
69    /// # Example
70    /// ```rust
71    /// # use ic_stable_memory::collections::SVec;
72    /// # use ic_stable_memory::stable_memory_init;
73    /// # unsafe { ic_stable_memory::mem::clear(); }
74    /// # stable_memory_init();
75    /// let mut at_least_10_numbers = SVec::<u64>::new_with_capacity(10)
76    ///     .expect("Out of memory");
77    /// ```
78    #[inline]
79    pub fn new_with_capacity(capacity: usize) -> Result<Self, OutOfMemory> {
80        assert!(capacity <= Self::max_capacity());
81
82        Ok(Self {
83            len: 0,
84            cap: capacity,
85            ptr: unsafe { allocate((capacity * T::SIZE) as u64)?.as_ptr() },
86            stable_drop_flag: true,
87            _marker_t: PhantomData::default(),
88        })
89    }
90
91    /// Returns the capacity of this [SVec]
92    #[inline]
93    pub fn capacity(&self) -> usize {
94        self.cap
95    }
96
97    /// Returns the length of this [SVec]
98    #[inline]
99    pub fn len(&self) -> usize {
100        self.len
101    }
102
103    /// Returns [true] if the length of this [SVec] is `0`
104    #[inline]
105    pub fn is_empty(&self) -> bool {
106        self.len == 0
107    }
108
109    /// Returns the maximum possible capacity of this [SVec]
110    #[inline]
111    pub const fn max_capacity() -> usize {
112        u32::MAX as usize / T::SIZE
113    }
114
115    /// Inserts a new element at the end of this [SVec]
116    ///
117    /// Will try to reallocate if `capacity == length`. If the canister is out of stable memory,
118    /// will return [Err] with the element that was about to get inserted.
119    #[inline]
120    pub fn push(&mut self, mut element: T) -> Result<(), T> {
121        if self.maybe_reallocate().is_ok() {
122            let elem_ptr = SSlice::_offset(self.ptr, (self.len * T::SIZE) as u64);
123            unsafe { crate::mem::write_fixed(elem_ptr, &mut element) };
124
125            self.len += 1;
126
127            Ok(())
128        } else {
129            Err(element)
130        }
131    }
132
133    /// Removes the last element of the [SVec]
134    ///
135    /// If the [SVec] is empty, returns [None].
136    #[inline]
137    pub fn pop(&mut self) -> Option<T> {
138        if self.is_empty() {
139            return None;
140        }
141
142        let elem_ptr = self.get_element_ptr(self.len - 1)?;
143        self.len -= 1;
144
145        Some(unsafe { crate::mem::read_fixed_for_move(elem_ptr) })
146    }
147
148    /// Returns a [SRef] pointing to the element at requested index
149    ///
150    /// See also [SVec::get_mut].
151    ///
152    /// If out of bounds, returns [None]
153    ///
154    /// # Example
155    /// ```rust
156    /// # use ic_stable_memory::collections::SVec;
157    /// # use ic_stable_memory::stable_memory_init;
158    /// # unsafe { ic_stable_memory::mem::clear(); }
159    /// # stable_memory_init();
160    /// let mut vec = SVec::<u64>::new();
161    /// vec.push(20).expect("Out of memory");
162    ///
163    /// let elem = vec.get(0).unwrap();
164    ///
165    /// assert_eq!(*elem, 20);
166    /// ```
167    #[inline]
168    pub fn get(&self, idx: usize) -> Option<SRef<'_, T>> {
169        let ptr = self.get_element_ptr(idx)?;
170
171        unsafe { Some(SRef::new(ptr)) }
172    }
173
174    /// Returns [SRefMut] pointing to the element at requested index
175    ///
176    /// See also [SVec::get].
177    ///
178    /// If out of bounds, returns [None]
179    ///
180    /// # Example
181    /// ```rust
182    /// # use ic_stable_memory::collections::SVec;
183    /// # use ic_stable_memory::stable_memory_init;
184    /// # unsafe { ic_stable_memory::mem::clear(); }
185    /// # stable_memory_init();
186    /// let mut vec = SVec::<u64>::new();
187    /// vec.push(20).expect("Out of memory");
188    ///
189    /// let mut elem = vec.get_mut(0).unwrap();
190    /// *elem = 100;
191    /// ```
192    #[inline]
193    pub fn get_mut(&mut self, idx: usize) -> Option<SRefMut<'_, T>> {
194        let ptr = self.get_element_ptr(idx)?;
195
196        unsafe { Some(SRefMut::new(ptr)) }
197    }
198
199    /// Replaces an element at requested index with a provided value
200    ///
201    /// # Panics
202    /// Panics if out of bounds.
203    ///
204    /// # Example
205    /// ```rust
206    /// # use ic_stable_memory::collections::SVec;
207    /// # use ic_stable_memory::stable_memory_init;
208    /// # unsafe { ic_stable_memory::mem::clear(); }
209    /// # stable_memory_init();
210    /// let mut vec = SVec::<u64>::new();
211    /// vec.push(10).expect("Out of memory");
212    ///
213    /// let prev = vec.replace(0, 20);
214    ///
215    /// assert_eq!(prev, 10);
216    /// assert_eq!(*vec.get(0).unwrap(), 20);
217    /// ```
218    pub fn replace(&mut self, idx: usize, mut element: T) -> T {
219        assert!(idx < self.len(), "Out of bounds");
220
221        let elem_ptr = SSlice::_offset(self.ptr, (idx * T::SIZE) as u64);
222
223        let prev_element = unsafe { crate::mem::read_fixed_for_move(elem_ptr) };
224        unsafe { crate::mem::write_fixed(elem_ptr, &mut element) };
225
226        prev_element
227    }
228
229    /// Inserts a new element at the requested index, forward-shifting all elements after it
230    ///
231    /// Will try to reallocate, if `capacity == length`. If the canister is out of stable memory,
232    /// will return [Err] with the element that was about to get inserted.
233    ///
234    /// # Panics
235    /// Panics if out of bounds.
236    pub fn insert(&mut self, idx: usize, mut element: T) -> Result<(), T> {
237        if idx == self.len {
238            return self.push(element);
239        }
240
241        assert!(idx < self.len, "out of bounds");
242
243        if self.maybe_reallocate().is_ok() {
244            let elem_ptr = SSlice::_offset(self.ptr, (idx * T::SIZE) as u64);
245
246            // moving elements after idx one slot to the right
247            let mut buf = vec![0u8; (self.len - idx) * T::SIZE];
248            unsafe { crate::mem::read_bytes(elem_ptr, &mut buf) };
249            unsafe { crate::mem::write_bytes(elem_ptr + T::SIZE as u64, &buf) };
250
251            // writing the element
252            unsafe { crate::mem::write_fixed(elem_ptr, &mut element) };
253
254            self.len += 1;
255
256            Ok(())
257        } else {
258            Err(element)
259        }
260    }
261
262    /// Removes element at the requested index, back-shifting all elements after it
263    ///
264    /// # Panics
265    /// Panics if out of bounds.
266    pub fn remove(&mut self, idx: usize) -> T {
267        assert!(idx < self.len, "out of bounds");
268
269        if idx == self.len - 1 {
270            return unsafe { self.pop().unwrap_unchecked() };
271        }
272
273        let elem_ptr = SSlice::_offset(self.ptr, (idx * T::SIZE) as u64);
274        let elem = unsafe { crate::mem::read_fixed_for_move(elem_ptr) };
275
276        let mut buf = vec![0u8; (self.len - idx - 1) * T::SIZE];
277        unsafe { crate::mem::read_bytes(elem_ptr + T::SIZE as u64, &mut buf) };
278        unsafe { crate::mem::write_bytes(elem_ptr, &buf) };
279
280        self.len -= 1;
281
282        elem
283    }
284
285    /// Swaps elements at requested indices with each other
286    ///
287    /// # Panics
288    /// Panics if `idx1` == `idx2` or if any of indices are out of bounds.
289    pub fn swap(&mut self, idx1: usize, idx2: usize) {
290        assert!(
291            idx1 < self.len() && idx2 < self.len() && idx1 != idx2,
292            "invalid idx"
293        );
294
295        let ptr1 = SSlice::_offset(self.ptr, (idx1 * T::SIZE) as u64);
296        let ptr2 = SSlice::_offset(self.ptr, (idx2 * T::SIZE) as u64);
297
298        let mut buf_1 = T::Buf::new(T::SIZE);
299        let mut buf_2 = T::Buf::new(T::SIZE);
300
301        unsafe { crate::mem::read_bytes(ptr1, buf_1._deref_mut()) };
302        unsafe { crate::mem::read_bytes(ptr2, buf_2._deref_mut()) };
303
304        unsafe { crate::mem::write_bytes(ptr2, buf_1._deref()) };
305        unsafe { crate::mem::write_bytes(ptr1, buf_2._deref()) };
306    }
307
308    /// Clears the [SVec] from elements
309    ///
310    /// Does not reallocate or shrink the underlying memory block.
311    #[inline]
312    pub fn clear(&mut self) {
313        while self.pop().is_some() {}
314    }
315
316    /// Performs binary search on a sorted [SVec], using the provided lambda
317    ///
318    /// Works the same way as in [Vec].
319    ///
320    /// # Example
321    /// ```rust
322    /// # use ic_stable_memory::collections::SVec;
323    /// # use ic_stable_memory::stable_memory_init;
324    /// # unsafe { ic_stable_memory::mem::clear(); }
325    /// # stable_memory_init();
326    /// let mut vec = SVec::new_with_capacity(100).expect("Out of memory");
327    ///
328    /// for i in 0..100 {
329    ///     vec.push(i);
330    /// }
331    ///
332    /// let idx = vec.binary_search_by(|it| it.cmp(&10)).unwrap();
333    ///
334    /// assert_eq!(idx, 10);
335    /// ```
336    pub fn binary_search_by<FN>(&self, mut f: FN) -> Result<usize, usize>
337    where
338        FN: FnMut(&T) -> Ordering,
339    {
340        if self.is_empty() {
341            return Err(0);
342        }
343
344        let mut min = 0;
345        let mut max = self.len;
346        let mut mid = (max - min) / 2;
347
348        loop {
349            let elem_ptr = SSlice::_offset(self.ptr, (mid * T::SIZE) as u64);
350            let elem = unsafe { crate::mem::read_fixed_for_reference(elem_ptr) };
351
352            let res = f(&elem);
353
354            match res {
355                Ordering::Equal => return Ok(mid),
356                Ordering::Greater => {
357                    max = mid;
358                    let new_mid = (max - min) / 2 + min;
359
360                    if new_mid == mid {
361                        return Err(mid);
362                    }
363
364                    mid = new_mid;
365                    continue;
366                }
367                Ordering::Less => {
368                    min = mid;
369                    let new_mid = (max - min) / 2 + min;
370
371                    if new_mid == mid {
372                        return Err(mid + 1);
373                    }
374
375                    mid = new_mid;
376                    continue;
377                }
378            }
379        }
380    }
381
382    /// Returns an immutable iterator over this collection
383    ///
384    /// # Example
385    /// ```rust
386    /// # use ic_stable_memory::collections::SVec;
387    /// # use ic_stable_memory::stable_memory_init;
388    /// # unsafe { ic_stable_memory::mem::clear(); }
389    /// # stable_memory_init();
390    /// let mut vec = SVec::new_with_capacity(100).expect("Out of memory");
391    ///
392    /// for i in 0..100 {
393    ///     vec.push(i);
394    /// }
395    ///
396    /// for elem in vec.iter() {
397    ///     println!("{}", *elem); // will print '0, 1, 2, 3, 4, ...'
398    /// }
399    /// ```
400    #[inline]
401    pub fn iter(&self) -> SVecIter<T> {
402        SVecIter::new(self)
403    }
404
405    /// Prints byte representation of this collection
406    ///
407    /// Useful for tests
408    pub fn debug_print(&self) {
409        print!("SVec[");
410        for i in 0..self.len {
411            let mut b = T::Buf::new(T::SIZE);
412            unsafe {
413                crate::mem::read_bytes(
414                    SSlice::_offset(self.ptr, (i * T::SIZE) as u64),
415                    b._deref_mut(),
416                )
417            };
418
419            print!("{:?}", b._deref());
420
421            if i < self.len - 1 {
422                print!(", ");
423            }
424        }
425
426        println!("]");
427    }
428
429    fn maybe_reallocate(&mut self) -> Result<(), OutOfMemory> {
430        if self.ptr == EMPTY_PTR {
431            self.ptr = unsafe { allocate((self.capacity() * T::SIZE) as u64)?.as_ptr() };
432            return Ok(());
433        }
434
435        if self.len() == self.capacity() {
436            self.cap = self.cap.checked_mul(2).unwrap();
437            assert!(self.cap <= Self::max_capacity());
438
439            let slice = unsafe { SSlice::from_ptr(self.ptr).unwrap() };
440
441            self.ptr = unsafe { reallocate(slice, (self.cap * T::SIZE) as u64)?.as_ptr() };
442        }
443
444        Ok(())
445    }
446
447    pub(crate) fn get_element_ptr(&self, idx: usize) -> Option<StablePtr> {
448        if idx < self.len() {
449            Some(SSlice::_offset(self.ptr, (idx * T::SIZE) as u64))
450        } else {
451            None
452        }
453    }
454}
455
456impl<T: StableType + AsFixedSizeBytes> Default for SVec<T> {
457    #[inline]
458    fn default() -> Self {
459        Self::new()
460    }
461}
462
463impl<T: StableType + AsFixedSizeBytes + Debug> Debug for SVec<T> {
464    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
465        f.write_str("[")?;
466        for (idx, item) in self.iter().enumerate() {
467            item.fmt(f)?;
468
469            if idx < self.len - 1 {
470                f.write_str(", ")?;
471            }
472        }
473        f.write_str("]")
474    }
475}
476
477impl<T: StableType + AsFixedSizeBytes> AsFixedSizeBytes for SVec<T> {
478    const SIZE: usize = u64::SIZE + usize::SIZE * 2;
479    type Buf = [u8; u64::SIZE + usize::SIZE * 2];
480
481    fn as_fixed_size_bytes(&self, buf: &mut [u8]) {
482        self.ptr.as_fixed_size_bytes(&mut buf[0..u64::SIZE]);
483        self.len
484            .as_fixed_size_bytes(&mut buf[u64::SIZE..(u64::SIZE + usize::SIZE)]);
485        self.cap.as_fixed_size_bytes(
486            &mut buf[(u64::SIZE + usize::SIZE)..(u64::SIZE + usize::SIZE * 2)],
487        );
488    }
489
490    fn from_fixed_size_bytes(arr: &[u8]) -> Self {
491        let ptr = u64::from_fixed_size_bytes(&arr[0..u64::SIZE]);
492        let len = usize::from_fixed_size_bytes(&arr[u64::SIZE..(u64::SIZE + usize::SIZE)]);
493        let cap = usize::from_fixed_size_bytes(
494            &arr[(u64::SIZE + usize::SIZE)..(u64::SIZE + usize::SIZE * 2)],
495        );
496
497        Self {
498            ptr,
499            len,
500            cap,
501            stable_drop_flag: false,
502            _marker_t: PhantomData::default(),
503        }
504    }
505}
506
507impl<T: StableType + AsFixedSizeBytes> StableType for SVec<T> {
508    #[inline]
509    unsafe fn stable_drop_flag_off(&mut self) {
510        self.stable_drop_flag = false;
511    }
512
513    #[inline]
514    unsafe fn stable_drop_flag_on(&mut self) {
515        self.stable_drop_flag = true;
516    }
517
518    #[inline]
519    fn should_stable_drop(&self) -> bool {
520        self.stable_drop_flag
521    }
522
523    unsafe fn stable_drop(&mut self) {
524        if self.ptr != EMPTY_PTR {
525            self.clear();
526
527            let slice = SSlice::from_ptr(self.ptr).unwrap();
528
529            deallocate(slice);
530        }
531    }
532}
533
534impl<T: StableType + AsFixedSizeBytes> Drop for SVec<T> {
535    fn drop(&mut self) {
536        if self.should_stable_drop() {
537            unsafe {
538                self.stable_drop();
539            }
540        }
541    }
542}
543
544#[cfg(test)]
545mod tests {
546    use crate::collections::vec::{SVec, DEFAULT_CAPACITY};
547    use crate::encoding::{AsFixedSizeBytes, Buffer};
548    use crate::primitive::s_box::SBox;
549    use crate::primitive::StableType;
550    use crate::utils::mem_context::stable;
551    use crate::utils::test::generate_random_string;
552    use crate::utils::DebuglessUnwrap;
553    use crate::{
554        _debug_print_allocator, _debug_validate_allocator, deinit_allocator, get_allocated_size,
555        init_allocator, retrieve_custom_data, stable_memory_init, stable_memory_post_upgrade,
556        stable_memory_pre_upgrade, store_custom_data,
557    };
558    use rand::rngs::ThreadRng;
559    use rand::seq::SliceRandom;
560    use rand::{thread_rng, Rng};
561    use std::fmt::Debug;
562    use std::ops::Deref;
563
564    #[derive(Copy, Clone, Debug)]
565    struct Test {
566        a: usize,
567        b: bool,
568    }
569
570    impl AsFixedSizeBytes for Test {
571        const SIZE: usize = usize::SIZE + bool::SIZE;
572        type Buf = [u8; Self::SIZE];
573
574        fn as_fixed_size_bytes(&self, buf: &mut [u8]) {
575            self.a.as_fixed_size_bytes(&mut buf[0..usize::SIZE]);
576            self.b
577                .as_fixed_size_bytes(&mut buf[usize::SIZE..(usize::SIZE + bool::SIZE)]);
578        }
579
580        fn from_fixed_size_bytes(arr: &[u8]) -> Self {
581            let a = usize::from_fixed_size_bytes(&arr[0..usize::SIZE]);
582            let b = bool::from_fixed_size_bytes(&arr[usize::SIZE..(usize::SIZE + bool::SIZE)]);
583
584            Self { a, b }
585        }
586    }
587
588    impl StableType for Test {}
589
590    #[test]
591    fn create_destroy_work_fine() {
592        stable::clear();
593        stable_memory_init();
594
595        {
596            let mut stable_vec = SVec::new();
597            assert_eq!(stable_vec.capacity(), DEFAULT_CAPACITY);
598            assert_eq!(stable_vec.len(), 0);
599
600            stable_vec.push(10).unwrap();
601            assert_eq!(stable_vec.capacity(), DEFAULT_CAPACITY);
602            assert_eq!(stable_vec.len(), 1);
603        }
604
605        _debug_validate_allocator();
606        assert_eq!(get_allocated_size(), 0);
607    }
608
609    #[test]
610    fn push_pop_work_fine() {
611        stable::clear();
612        stable_memory_init();
613
614        {
615            let mut stable_vec = SVec::new();
616            let count = 10usize;
617
618            for i in 0..count {
619                let it = Test { a: i, b: true };
620
621                stable_vec.push(it).unwrap();
622            }
623
624            assert_eq!(stable_vec.len(), count, "Invalid len after push");
625
626            for i in 0..count {
627                let it = Test { a: i, b: false };
628
629                stable_vec.replace(i, it);
630            }
631
632            stable_vec.debug_print();
633            assert_eq!(stable_vec.len(), count, "Invalid len after push");
634
635            for i in 0..count {
636                println!("{} {}", i, stable_vec.len());
637
638                let it = stable_vec.pop().unwrap();
639
640                assert_eq!(stable_vec.len(), count - i - 1);
641                assert_eq!(it.a, count - 1 - i);
642                assert!(!it.b);
643            }
644
645            assert_eq!(stable_vec.len(), 0, "Invalid len after pop");
646
647            for i in 0..count {
648                let it = Test { a: i, b: true };
649
650                stable_vec.push(it).unwrap();
651            }
652        }
653
654        _debug_validate_allocator();
655        assert_eq!(get_allocated_size(), 0);
656    }
657
658    #[test]
659    fn basic_flow_works_fine() {
660        stable::clear();
661        stable_memory_init();
662
663        {
664            let mut v = SVec::default();
665            assert!(v.get(100).is_none());
666
667            v.push(10).unwrap();
668            v.push(20).unwrap();
669
670            assert_eq!(*v.get(0).unwrap(), 10);
671            assert_eq!(*v.get(1).unwrap(), 20);
672            assert_eq!(v.replace(0, 11), 10);
673        }
674
675        _debug_validate_allocator();
676        assert_eq!(get_allocated_size(), 0);
677    }
678
679    #[test]
680    fn insert_works_fine() {
681        stable::clear();
682        stable_memory_init();
683
684        {
685            let mut array = SVec::default();
686            let mut check = Vec::default();
687
688            for i in 0..30 {
689                array.insert(0, 29 - i).unwrap();
690                check.insert(0, 29 - i);
691            }
692
693            for i in 60..100 {
694                array.insert(array.len(), i).unwrap();
695                check.insert(check.len(), i);
696            }
697
698            for i in 30..60 {
699                array.insert(30 + (i - 30), i).unwrap();
700                check.insert(30 + (i - 30), i);
701            }
702
703            for i in 0..100 {
704                assert_eq!(*array.get(i).unwrap(), i);
705            }
706
707            let mut actual = Vec::new();
708            while let Some(elem) = array.pop() {
709                actual.insert(0, elem);
710            }
711
712            assert_eq!(actual, check);
713        }
714
715        _debug_validate_allocator();
716        assert_eq!(get_allocated_size(), 0);
717    }
718
719    #[test]
720    fn binary_search_work_fine() {
721        stable::clear();
722        stable_memory_init();
723
724        {
725            // 0..100 randomly shuffled
726            let initial = vec![
727                3, 24, 46, 92, 2, 21, 34, 95, 82, 22, 88, 32, 59, 13, 73, 51, 12, 83, 28, 17, 9,
728                23, 5, 63, 62, 38, 20, 40, 25, 98, 30, 43, 57, 86, 42, 4, 99, 33, 11, 74, 96, 94,
729                47, 31, 37, 71, 80, 70, 14, 67, 93, 56, 27, 39, 58, 41, 29, 84, 8, 0, 45, 54, 7,
730                26, 97, 6, 81, 65, 79, 10, 91, 68, 36, 60, 76, 75, 15, 87, 49, 35, 78, 64, 69, 52,
731                50, 61, 48, 53, 44, 19, 55, 72, 90, 77, 89, 16, 85, 66, 18, 1,
732            ];
733
734            let mut array = SVec::<i32>::default();
735            let mut check = Vec::<i32>::new();
736
737            for i in 0..100 {
738                match check.binary_search_by(|it| it.cmp(&initial[i])) {
739                    Err(idx) => check.insert(idx, initial[i]),
740                    _ => unreachable!(),
741                }
742
743                match array.binary_search_by(|it| it.cmp(&initial[i])) {
744                    Err(idx) => array.insert(idx, initial[i]).unwrap(),
745                    _ => unreachable!(),
746                }
747
748                _debug_print_allocator();
749            }
750
751            let mut actual = Vec::new();
752            while let Some(elem) = array.pop() {
753                actual.insert(0, elem);
754            }
755
756            assert_eq!(actual, check);
757        }
758
759        _debug_validate_allocator();
760        assert_eq!(get_allocated_size(), 0);
761    }
762
763    #[test]
764    fn remove_works_fine() {
765        stable::clear();
766        stable_memory_init();
767
768        {
769            let initial = vec![
770                3, 24, 46, 92, 2, 21, 34, 95, 82, 22, 88, 32, 59, 13, 73, 51, 12, 83, 28, 17, 9,
771                23, 5, 63, 62, 38, 20, 40, 25, 98, 30, 43, 57, 86, 42, 4, 99, 33, 11, 74, 96, 94,
772                47, 31, 37, 71, 80, 70, 14, 67, 93, 56, 27, 39, 58, 41, 29, 84, 8, 0, 45, 54, 7,
773                26, 97, 6, 81, 65, 79, 10, 91, 68, 36, 60, 76, 75, 15, 87, 49, 35, 78, 64, 69, 52,
774                50, 61, 48, 53, 44, 19, 55, 72, 90, 77, 89, 16, 85, 66, 18, 1,
775            ];
776
777            let mut array = SVec::new();
778            for i in 0..initial.len() {
779                array.push(initial[i]);
780            }
781
782            for i in 0..initial.len() {
783                assert_eq!(array.remove(0), initial[i]);
784            }
785        }
786
787        _debug_validate_allocator();
788        assert_eq!(get_allocated_size(), 0);
789    }
790
791    #[test]
792    fn serialization_works_fine() {
793        stable::clear();
794        stable_memory_init();
795
796        {
797            let mut vec = SVec::<u32>::new_with_capacity(10).debugless_unwrap();
798            vec.push(1);
799            vec.push(2);
800            vec.push(3);
801
802            let mut buf = <SVec<u32> as AsFixedSizeBytes>::Buf::new(SVec::<u32>::SIZE);
803            vec.as_fixed_size_bytes(buf._deref_mut());
804            let mut vec1 = SVec::<u32>::from_fixed_size_bytes(buf._deref());
805            unsafe { vec1.stable_drop_flag_off() };
806
807            assert_eq!(vec.ptr, vec1.ptr);
808            assert_eq!(vec.len, vec1.len);
809            assert_eq!(vec.cap, vec1.cap);
810
811            let ptr = vec.ptr;
812            let len = vec.len;
813            let cap = vec.cap;
814
815            let mut buf = <SVec<u32> as AsFixedSizeBytes>::Buf::new(SVec::<u32>::SIZE);
816            vec.as_fixed_size_bytes(buf._deref_mut());
817
818            let mut vec1 = SVec::<u32>::from_fixed_size_bytes(buf._deref());
819            unsafe { vec1.stable_drop_flag_off() };
820
821            assert_eq!(ptr, vec1.ptr);
822            assert_eq!(len, vec1.len);
823            assert_eq!(cap, vec1.cap);
824        }
825
826        _debug_print_allocator();
827        _debug_validate_allocator();
828        assert_eq!(get_allocated_size(), 0);
829    }
830
831    #[test]
832    fn iter_works_fine() {
833        stable::clear();
834        stable_memory_init();
835
836        {
837            let mut vec = SVec::new();
838            for i in 0..100 {
839                vec.push(i);
840            }
841
842            let mut c = 0;
843            vec.debug_print();
844
845            for (idx, mut i) in vec.iter().enumerate() {
846                c += 1;
847
848                assert_eq!(idx as i32, *i);
849            }
850
851            assert_eq!(c, 100);
852        }
853
854        _debug_validate_allocator();
855        assert_eq!(get_allocated_size(), 0);
856    }
857
858    #[test]
859    fn random_works_fine() {
860        stable::clear();
861        stable_memory_init();
862
863        {
864            let mut svec = SVec::new();
865
866            let mut rng = thread_rng();
867            let iterations = 10_000;
868
869            let mut example = Vec::new();
870            for i in 0..iterations {
871                example.push(i);
872            }
873            example.shuffle(&mut rng);
874
875            for i in 0..iterations {
876                svec.push(example[i]);
877            }
878
879            for i in 0..iterations {
880                svec.insert(rng.gen_range(0..svec.len()), example[i]);
881            }
882
883            for _ in 0..iterations {
884                svec.pop().unwrap();
885            }
886
887            for i in 0..iterations {
888                if svec.len() == 1 {
889                    svec.remove(0);
890                } else {
891                    let range = rng.gen_range(0..(svec.len() - 1));
892                    svec.remove(range);
893                }
894            }
895
896            assert!(svec.is_empty());
897        }
898
899        _debug_validate_allocator();
900        assert_eq!(get_allocated_size(), 0);
901    }
902
903    #[test]
904    fn sboxes_work_fine() {
905        stable::clear();
906        stable_memory_init();
907
908        {
909            let mut vec = SVec::new();
910
911            for _ in 0..100 {
912                let b = SBox::new(10).unwrap();
913
914                vec.push(b);
915            }
916
917            vec.debug_print();
918        }
919
920        _debug_validate_allocator();
921        assert_eq!(get_allocated_size(), 0);
922    }
923
924    #[derive(Debug)]
925    enum Action {
926        Push(usize),
927        Pop(usize),
928        Insert(usize, usize),
929        Remove(usize, usize),
930        Swap(usize, usize, usize),
931        Replace(usize, usize),
932        CanisterUpgrade,
933        GetMut(usize, usize),
934        Clear(usize),
935    }
936
937    struct Fuzzer {
938        vec: Option<SVec<SVec<SBox<String>>>>,
939        example: Vec<Vec<String>>,
940        rng: ThreadRng,
941        log: Vec<Action>,
942    }
943
944    impl Fuzzer {
945        fn new() -> Self {
946            let mut svec = SVec::default();
947            let mut vec = Vec::default();
948
949            for i in 0..10 {
950                svec.push(SVec::default());
951                vec.push(Vec::default());
952            }
953
954            Self {
955                vec: Some(svec),
956                example: vec,
957                rng: thread_rng(),
958                log: Vec::default(),
959            }
960        }
961
962        fn vec(&mut self) -> &mut SVec<SVec<SBox<String>>> {
963            self.vec.as_mut().unwrap()
964        }
965
966        fn next(&mut self) {
967            let action = self.rng.gen_range(0..1100);
968
969            match action {
970                // PUSH ~25%
971                0..=250 => {
972                    let str = generate_random_string(&mut self.rng);
973
974                    if let Ok(data) = SBox::new(str.clone()) {
975                        let outer_idx = self.rng.gen_range(0..10);
976
977                        if self.vec().get_mut(outer_idx).unwrap().push(data).is_err() {
978                            return;
979                        };
980
981                        self.example.get_mut(outer_idx).unwrap().push(str.clone());
982
983                        self.log.push(Action::Push(outer_idx));
984                    }
985                }
986                // INSERT ~25%
987                251..=500 => {
988                    let str = generate_random_string(&mut self.rng);
989
990                    if let Ok(data) = SBox::new(str.clone()) {
991                        let outer_idx = self.rng.gen_range(0..10);
992
993                        let len = self.vec().get(outer_idx).unwrap().len();
994                        let idx = if len == 0 {
995                            0
996                        } else {
997                            self.rng.gen_range(0..len + 1)
998                        };
999
1000                        if self
1001                            .vec()
1002                            .get_mut(outer_idx)
1003                            .unwrap()
1004                            .insert(idx, data)
1005                            .is_err()
1006                        {
1007                            return;
1008                        }
1009
1010                        self.example
1011                            .get_mut(outer_idx)
1012                            .unwrap()
1013                            .insert(idx, str.clone());
1014
1015                        self.log.push(Action::Insert(outer_idx, idx));
1016                    }
1017                }
1018                // POP ~10%
1019                501..=600 => {
1020                    let outer_idx = self.rng.gen_range(0..10);
1021
1022                    self.vec().get_mut(outer_idx).unwrap().pop();
1023                    self.example.get_mut(outer_idx).unwrap().pop();
1024
1025                    self.log.push(Action::Pop(outer_idx));
1026                }
1027                // REMOVE ~10%
1028                601..=700 => {
1029                    let outer_idx = self.rng.gen_range(0..10);
1030
1031                    let len = self.vec().get(outer_idx).unwrap().len();
1032
1033                    if len == 0 {
1034                        return self.next();
1035                    }
1036
1037                    let idx = if len == 1 {
1038                        0
1039                    } else {
1040                        self.rng.gen_range(0..len)
1041                    };
1042
1043                    self.vec().get_mut(outer_idx).unwrap().remove(idx);
1044                    self.example.get_mut(outer_idx).unwrap().remove(idx);
1045
1046                    self.log.push(Action::Remove(outer_idx, idx));
1047                }
1048                // SWAP ~10%
1049                701..=800 => {
1050                    let outer_idx = self.rng.gen_range(0..10);
1051
1052                    let len = self.vec().get(outer_idx).unwrap().len();
1053
1054                    if len < 2 {
1055                        return self.next();
1056                    }
1057
1058                    let mut idx1 = self.rng.gen_range(0..len);
1059                    let mut idx2 = self.rng.gen_range(0..len);
1060
1061                    if idx1 == idx2 {
1062                        if idx2 < len - 1 {
1063                            idx2 += 1;
1064                        } else {
1065                            idx2 -= 1;
1066                        }
1067                    }
1068
1069                    self.vec().get_mut(outer_idx).unwrap().swap(idx1, idx2);
1070                    self.example.get_mut(outer_idx).unwrap().swap(idx1, idx2);
1071
1072                    self.log.push(Action::Swap(outer_idx, idx1, idx2));
1073                }
1074                // REPLACE ~10%
1075                801..=900 => {
1076                    let outer_idx = self.rng.gen_range(0..10);
1077
1078                    let len = self.vec().get(outer_idx).unwrap().len();
1079
1080                    if len == 0 {
1081                        return self.next();
1082                    }
1083
1084                    let idx = self.rng.gen_range(0..len);
1085                    let str = generate_random_string(&mut self.rng);
1086
1087                    if let Ok(data) = SBox::new(str.clone()) {
1088                        self.vec().get_mut(outer_idx).unwrap().replace(idx, data);
1089
1090                        std::mem::replace(
1091                            self.example
1092                                .get_mut(outer_idx)
1093                                .unwrap()
1094                                .get_mut(idx)
1095                                .unwrap(),
1096                            str.clone(),
1097                        );
1098
1099                        self.log.push(Action::Replace(outer_idx, idx));
1100                    }
1101                }
1102                // GET MUT ~10%
1103                901..=1000 => {
1104                    let outer_idx = self.rng.gen_range(0..10);
1105
1106                    let len = self.vec().get_mut(outer_idx).unwrap().len();
1107
1108                    if len == 0 {
1109                        return self.next();
1110                    }
1111
1112                    let idx = self.rng.gen_range(0..len);
1113                    let str = generate_random_string(&mut self.rng);
1114
1115                    if self
1116                        .vec()
1117                        .get_mut(outer_idx)
1118                        .unwrap()
1119                        .get_mut(idx)
1120                        .unwrap()
1121                        .with(|s: &mut String| {
1122                            *s = str.clone();
1123                        })
1124                        .is_err()
1125                    {
1126                        return;
1127                    }
1128
1129                    *self
1130                        .example
1131                        .get_mut(outer_idx)
1132                        .unwrap()
1133                        .get_mut(idx)
1134                        .unwrap() = str;
1135
1136                    self.log.push(Action::GetMut(outer_idx, idx));
1137                }
1138                // CLEAR
1139                1001..=1010 => {
1140                    let outer_idx = self.rng.gen_range(0..10);
1141
1142                    self.vec().get_mut(outer_idx).unwrap().clear();
1143                    self.example.get_mut(outer_idx).unwrap().clear();
1144
1145                    self.log.push(Action::Clear(outer_idx));
1146                }
1147                // CANISTER UPGRADE
1148                _ => match SBox::new(self.vec.take().unwrap()) {
1149                    Ok(data) => {
1150                        store_custom_data(1, data);
1151
1152                        if stable_memory_pre_upgrade().is_ok() {
1153                            stable_memory_post_upgrade();
1154                        }
1155
1156                        self.vec = retrieve_custom_data::<SVec<SVec<SBox<String>>>>(1)
1157                            .map(|it| it.into_inner());
1158                        self.log.push(Action::CanisterUpgrade);
1159                    }
1160                    Err(vec) => {
1161                        self.vec = Some(vec);
1162                    }
1163                },
1164            }
1165
1166            _debug_validate_allocator();
1167            assert_eq!(self.vec().len(), self.example.len());
1168
1169            for i in 0..self.vec().len() {
1170                let svec = self.vec.as_ref().unwrap().get(i).unwrap();
1171                let example = self.example.get(i).unwrap();
1172
1173                assert_eq!(svec.len(), example.len());
1174
1175                for j in 0..svec.len() {
1176                    assert_eq!(
1177                        svec.get(j).unwrap().clone(),
1178                        example.get(j).unwrap().clone()
1179                    );
1180                }
1181            }
1182        }
1183    }
1184
1185    #[test]
1186    fn fuzzer_works_fine() {
1187        stable::clear();
1188        init_allocator(0);
1189
1190        {
1191            let mut fuzzer = Fuzzer::new();
1192
1193            for _ in 0..10_000 {
1194                fuzzer.next();
1195            }
1196        }
1197
1198        assert_eq!(get_allocated_size(), 0);
1199    }
1200
1201    #[test]
1202    fn fuzzer_works_fine_limited_memory() {
1203        stable::clear();
1204        init_allocator(10);
1205
1206        {
1207            let mut fuzzer = Fuzzer::new();
1208
1209            for _ in 0..10_000 {
1210                fuzzer.next();
1211            }
1212        }
1213
1214        assert_eq!(get_allocated_size(), 0);
1215    }
1216}
ic_stable_memory/collections/vec/mod.rs

ic_stable_memory/collections/vec/
mod.rs
