Code shared from the Rust Playground

playground.rs

// Creating a heapless vector type using MaybeUninit<T> safely with
// unintialized memory (useful in embedded systems where there's limited
// memory space for heap allocations).
#![no_std]

extern crate std;

use arrayvec::ArrayVec;

mod arrayvec {
    use core::mem::{ManuallyDrop, MaybeUninit};
    use core::ptr;

    #[derive(Debug)]
    pub struct ArrayVec<T, const N: usize> {
        values: [MaybeUninit<T>; N],
        len: usize,
    }

    impl<T, const N: usize> ArrayVec<T, N> {
        /// Creates a new empty ArrayVec
        pub fn new() -> Self {
            // [MaybeUninit<T>; N] is zero-initialized to uninit by default.
            // Meaning the array starts from a blank slate waiting to be
            // initialized by `write()`; filling uninit elements with
            // "garbage" bytes.
            ArrayVec {
                // values: unsafe { MaybeUninit::uninit().assume_init() },
                // Same as the commented code above but safer.
                values: [const { MaybeUninit::uninit() }; N],
                len: 0,
            }
        }

        /// Pushes a value if there's space, returning `Err(value)` if full.
        /// Safe: `.write()` takes ownership and marks the slot as
        /// initialized.
        pub fn try_push(&mut self, value: T) -> Result<(), T> {
            if self.len == N {
                return Err(value);
            }
            self.values[self.len].write(value);
            self.len += 1;
            Ok(())
        }

        /// Returns a reference to the element at `index` if within bounds
        /// and initialized.
        /// SAFETY: Unsafe internally: Assumes first `len` slots are init.
        pub fn get(&self, index: usize) -> Option<&T> {
            if index >= self.len {
                return None;
            }
            // Same as the commented code below
            unsafe { Some(&*self.values[index].as_ptr()) }
            // unsafe { Some(self.values[index].assume_init_ref()) }
        }

        /// Pops the last value if any, returning it (or `None` if empty).
        /// SAFETY: Safe: Uses `.assume_init_read()` to extract and mark
        /// as uninit.
        pub fn pop(&mut self) -> Option<T> {
            if self.len == 0 {
                return None;
            }
            self.len -= 1;
            Some(unsafe { self.values[self.len].assume_init_read() })
        }

        /// Returns the current length.
        pub fn len(&self) -> usize {
            self.len
        }

        /// Instead of `into_arr` lets return a slice using
        /// `slice::from_raw_parts()`.
        /// Returns a slice over init elements (& first `len` slots).
        /// SAFETY: Unsafe internally, but safe API: assumes invariant holds.
        pub fn as_slice(&self) -> &[T] {
            unsafe {
                core::slice::from_raw_parts(
                    self.values.as_ptr() as *const T, self.len
                )
            }
        }

        /// Returns a mutable slice over init elements.
        pub fn as_mut_slice(&mut self) -> &mut [T] {
            unsafe {
                core::slice::from_raw_parts_mut(
                    self.values.as_mut_ptr() as *mut T, self.len
                )
            }
        }

        /// Use slice iterator for immutable iteration
        pub fn iter(&self) -> core::slice::Iter<'_, T> {
            self.as_slice().iter()
        }

        /// Use slice iterator for mutable iteration
        pub fn iter_mut(&mut self) -> core::slice::IterMut<'_, T> {
            self.as_mut_slice().iter_mut()
        }
    }
    
    // Implement these only where `T` is `Copy`.
    /* Better this way because in some impls we dereference `T`;
    if `T` is not `Copy` we'll be creating a use-after-free or
    double-free which is UB.
    */
    impl<T, const N: usize> ArrayVec<T, N>
    where
        T: Copy,
    {
        /// Returns Self with a view of initialized and uninitialized
        /// accesses of memory.
        pub fn show_init<'a>(&self, other: &'a mut ArrayVec<Option<T>, N>)
        -> &'a [Option<T>]
        {
            let mut count = 0;
            
            for item in self {
                let _ = other.try_push(Some(*item)); /* deref `T` to copy
                primitive value.
                */
                count += 1;
            }
            while count < N {
                let _ = other.try_push(None);
                count += 1;
            }
            other.as_slice()
        }
    }

    // Implement Drop trait to safely deallocate initialized elements.
    impl<T, const N: usize> Drop for ArrayVec<T, N> {
        fn drop(&mut self) {
            // Explicity drop the first `len` initialized elements.
            for i in 0..self.len {
                unsafe {
                    self.values[i].assume_init_drop();
                }
            }
        }
    }

    // Build out iterator type for ArrayVec as ArrayVecIntoIter<T, N>
    // Consuming iterator (by-value): Moves out owned T.
    // But will provide iterators for fundamental types: & and &mut
    #[derive(Debug)]
    pub struct ArrayVecIntoIter<T, const N: usize> {
        values: [MaybeUninit<T>; N],
        len: usize,
        index: usize,
    }

    // Implement Iterator trait on ArrayVecIntoIter making it an iterator.
    impl<T, const N: usize> Iterator for ArrayVecIntoIter<T, N> {
        type Item = T;

        fn next(&mut self) -> Option<Self::Item> {
            if self.index >= self.len {
                return None;
            }
            let i = self.index;
            self.index += 1;
            // SAFETY: i < len, so slot is init
            Some(unsafe { self.values[i].assume_init_read() })
        }

        fn size_hint(&self) -> (usize, Option<usize>) {
            let remaining = self.len.saturating_sub(self.index);
            (remaining, Some(remaining))
        }
    }

    // Implement Drop trait to safely deallocate initialized elements
    // from ArrayVecIntoIter<T, N> MaybeUninit<T> values
    impl<T, const N: usize> Drop for ArrayVecIntoIter<T, N> {
        fn drop(&mut self) {
            // Drop remaining init elements (from index to len)
            // SAFETY: Invariant holds for those slots.
            for i in self.index..self.len {
                unsafe {
                    self.values[i].assume_init_drop();
                }
            }
            // Uninit slots auto-drop as MaybeUninit: meaning as "garbage"
            // bytes they are overwritten by the next write.
        }
    }

    // Finally, implement IntoIterator for ArrayVecIntoIter; returns
    // an iterator.
    impl<T, const N: usize> IntoIterator for ArrayVec<T, N> {
        type Item = T;
        type IntoIter = ArrayVecIntoIter<T, N>;

        fn into_iter(self) -> Self::IntoIter {
            // Wrap `self` inside ManuallyDrop so ArrayVec Drop cannot
            // be called on ArrayVec. Prevents double-free.
            // Drop will be handled by ArrayVecIntoIter.
            let this = ManuallyDrop::new(self);
            // SAFETY: Read (bitwise copy) into `values` and `len`,
            // since `self` is consumed by function call; `this` will not
            // be used again.
            let values = unsafe { ptr::read(&this.values) };
            let len = unsafe { ptr::read(&this.len) };
            
            ArrayVecIntoIter {
                values,
                len,
                index: 0,
            }
        }
    }

    // Implement IntoIterator for fundamental types of ArrayVec: & and &mut.
    // By reference: yields &T
    impl<'a, T, const N: usize> IntoIterator for &'a ArrayVec<T, N> {
        type Item = &'a T;
        type IntoIter = core::slice::Iter<'a, T>;

        fn into_iter(self) -> Self::IntoIter {
            self.as_slice().iter()
        }
    }

    // By mutable reference: yields &mut T
    impl<'a, T, const N: usize> IntoIterator for &'a mut ArrayVec<T, N> {
        type Item = &'a mut T;
        type IntoIter = core::slice::IterMut<'a, T>;

        fn into_iter(self) -> Self::IntoIter {
            self.as_mut_slice().iter_mut()
        }
    }
    
    // Extending functionality: Implementing FromIterator.
    impl<T, const N: usize> FromIterator<T> for ArrayVec<T, N> {
        fn from_iter<I>(iter: I) -> Self
        where
            I: IntoIterator<Item = T>
        {
            let mut arr_vec = Self::new();
            let iter = iter.into_iter();
            
            // Optimize: If hinted size > N, take only N to skip
            // excess early.
            let (lower, _) = iter.size_hint();
            if lower > N {
                for item in iter.take(N) {
                    let _ = arr_vec.try_push(item); /* Won't fail as
                    len < N.
                    */
                }
                return arr_vec;
            }
            
            // General case: Push until full (truncates excess).
            for item in iter {
                let _ = arr_vec.try_push(item); /* Err(item) dropped
                if full.
                */
            }
            arr_vec
        }
    }
    
    // Bonus: Extend for appending (up to "both" iter and ArrayVec capacity)
    impl<T, const N: usize> Extend<T> for ArrayVec<T, N> {
        fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
            for item in iter { // Stop iterating when `iter` is consumed
                if let Err(_) = self.try_push(item) { /* break if
                `self` is full
                */
                    break;
                }
            }
        }
    }
}

const CAP: usize = 5;

fn main() {
    {
        // A:
        let mut arr_vec = ArrayVec::<i32, CAP>::new();
        std::println!("{:?}", arr_vec); // View init ArrayVec
    
        let mut count;
    
        // Push a few elements
        for i in 0..(CAP - 2) {
            count = 1 + i as i32;
            arr_vec.try_push(count).unwrap()
        }
        std::println!("{:?}", arr_vec); // View pushed elements
    
        // Return element in initialized index;
        // if index is not initialized, return None.
        // let arr_els = ArrayVec::into_arr(&arr_vec);
        let arr_els = arr_vec.as_slice();
        std::println!("---\nInit values: {:?}", arr_els);
    
        // Pop from MaybeUninit ArrayVec; if uninit, return None.
        std::println!("---\nArrayVec `len` before pop: {}", arr_vec.len());
        std::println!("Popped value: {:?}", arr_vec.pop());
        std::println!("ArrayVec `len` after pop: {}", arr_vec.len());
    
        // TEST: Add more elements beyond `CAP` size for ArrayVec;
        // `try_push` should escape and return with Err.
        let arr_len = arr_vec.len();
        count = *arr_vec.get(arr_len - 1).unwrap();
        let mut arr_err_els: ArrayVec<Result<(), i32>, CAP> =
            ArrayVec::new();
    
        for _ in arr_len..(CAP * 2) {
            count += 1;
            if let Err(value) = arr_vec.try_push(count) {
                arr_err_els.try_push(Err(value)).unwrap();
            }
        }
        std::println!("---\nFilled ArrayVec: {:?}", arr_vec.as_slice());
        std::println!("Err beyond ArrayVec: {:?}", arr_err_els.as_slice());
    }
    
    {
        // B:
        // Iterating over ArrayVec through fundamental types: & and &mut

        // Simple iteration
        let mut arr_vec = ArrayVec::<u8, CAP>::new();
        let mut count;
        for i in 0..CAP {
            count = 1 + i as u8;
            arr_vec.try_push(count).unwrap(); // Init values on ArrayVec
        }
        std::println!("---");
        count = 0;
        for i in &arr_vec {
            std::println!("ArrayVec at index {}: {}", count, i);
            count += 1;
        }

        // Mutable iteration
        for value in &mut arr_vec {
            *value += 10;
        }
        std::println!("---\nMutated ArrayVec: {:?}", arr_vec.as_slice());
    }
    
    {
        // C:
        // Iterating over owned values of ArrayVec by calling `into_iter`
        let mut arr_vec = ArrayVec::<u8, CAP>::new();
        let mut count;
        for i in 0..(CAP - 1) {
            count = 1 + i as u8;
            arr_vec.try_push(count).unwrap();
        }
        let mut arr_iter = arr_vec.into_iter(); // move arr_vec
        // std::println!("{:?}", arr_vec); // should return move error
        /* Above code confirms `impl IntoIterator for ArrayVec` safety
        invariants.
        */
        std::println!("---\n{:?}", arr_iter); /* View created
        ArrayVecIntoIter. See `len` and `index` fields.
        */
        
        let mut arr_vec2 = ArrayVec::<Option<u8>, CAP>::new();
        loop { /* Loop through calls to `next()` to iterate through
        ArrayVecIntoInter.
        */
            match arr_iter.next() {
                Some(val) => arr_vec2.try_push(Some(val)).unwrap(),
                None => {
                    arr_vec2.try_push(None).unwrap();
                    break;
                }
            }
        }
        std::println!("{:?}", arr_iter); /* Review `len` and `index` fields.
        Notice index incr caused by calling `next()` on ArrayVecIntoIter.
        */
        std::println!("{:?}", arr_vec2.as_slice());
    }
    
    {
        // D:
        // Test FromIterator implementation with iterators on ArrayVec.
        
        // Using collect:
        let arr_vec = (-10..10).collect::<ArrayVec<i8, 15>>();
        std::println!("---\n{:?}", arr_vec.as_slice());
        
        // Using from_iter:
        let arr_vec = ArrayVec::<_, 5>::from_iter(-3..5 as i8); /* Using
        type inference for `T` in `ArrayVec<T, N>` helped by type cast
        on iterator.
        */
        std::println!("{:?}", arr_vec.as_slice());
    }
    
    {   
        // E:
        // Test Extend implementation with iterators on ArrayVec.
        
        /* Testing for two scenarios;
        
        1. `arr_vec1` has more CAP (array "capacity" - N) than `arr_vec2`
        has elements to fill it. Meaning, `arr_vec1` will retain spare CAP.
        
        2. `arr_vec1` has less CAP (array "capacity" - N) than `arr_vec2`
        has elements to fill it. Meaning, `arr_vec1` will max it's CAP
        without extending all of `arr_vec2`s elements.
        */
        
        // Scenario 1:
        type ArrayVecCap10<T> = ArrayVec<T, 10>;
        type ArrayVecCap5<T> = ArrayVec<T, 5>;
        
        let mut arr_vec1: ArrayVecCap10<u8> = ArrayVec::new();
        let mut count;
        for i in 0..3 { // Add elements to `arr_vec1`.
            count = 1 + i as u8;
            arr_vec1.try_push(count).unwrap();
        }
        
        let mut arr_vec2: ArrayVecCap5<u8> = ArrayVec::new();
        for i in 0..2 { // Add elements to `arr_vec2`.
            count = 1 + i as u8;
            arr_vec2.try_push(count).unwrap();
        }
        
        arr_vec1.extend(arr_vec2); // Extend moves `arr_vec2`.
        
        let mut empty_arr_vec: ArrayVecCap10<Option<u8>> = ArrayVec::new();
        std::println!("---\n{:?}", arr_vec1.show_init(&mut empty_arr_vec));
        drop(empty_arr_vec); /* Choosing to explicity drop (or free, or
        deallocate) the stack memory consumed by `empty_arr_vec` here
        as it's no longer in use.
        */
        
        // Scenario 2:
        let mut arr_vec1: ArrayVecCap5<i8> = ArrayVec::new();
        for i in -2..0 { arr_vec1.try_push(i as i8).unwrap(); }
        
        let mut arr_vec2: ArrayVecCap10<i8> = ArrayVec::new();
        for i in -5..0 { arr_vec2.try_push(i as i8).unwrap(); }
        
        arr_vec1.extend(arr_vec2); // Extend moves `arr_vec2`.
        
        let mut empty_arr_vec: ArrayVecCap5<Option<i8>> = ArrayVec::new();
        std::println!("{:?}", arr_vec1.show_init(&mut empty_arr_vec));
        /* Don't need to explicity drop `empty_arr_vec` here as before
        because this is the end of the scope, and the Rust Compiler (rustc)
        will call the drop method in ArrayVec's destructor `Drop`
        to drop `empty_arr_vec` implicity.
        */
    }
}