metalssh 0.0.1

use core::ops::Deref;
use core::ops::DerefMut;

use zeroize::Zeroize;
use zeroize::ZeroizeOnDrop;

/// Adaptable [`Vec`] wrapper that prefers stack allocation up to constant `N`,
/// falling back to heap if this threshold is crossed. If no heap is available,
/// only stack will be used.
///
/// Also zeroizes on drop.
#[derive(Zeroize, ZeroizeOnDrop)]
pub enum Vec<T: Zeroize, const N: usize> {
    Stack(heapless::Vec<T, N>),
    #[cfg(feature = "alloc")]
    Heap(alloc::vec::Vec<T>),
}

impl<T, const N: usize> Vec<T, N>
where
    T: Zeroize,
{
    pub fn new() -> Self {
        Self::Stack(heapless::Vec::new())
    }

    /// Create a [`Vec`] from a slice.
    /// Uses stack if it fits, otherwise falls back to heap (if available).
    #[must_use]
    pub fn from_slice(slice: &[T]) -> Self
    where
        T: Clone,
    {
        if let Ok(v) = heapless::Vec::from_slice(slice) {
            Self::Stack(v)
        } else {
            #[cfg(feature = "alloc")]
            {
                Self::Heap(slice.to_vec())
            }
            #[cfg(not(feature = "alloc"))]
            {
                panic!("Secret too large for stack and no heap available");
            }
        }
    }

    #[inline]
    #[must_use]
    pub fn len(&self) -> usize {
        match self {
            Self::Stack(v) => v.len(),
            #[cfg(feature = "alloc")]
            Self::Heap(v) => v.len(),
        }
    }

    #[inline]
    #[must_use]
    pub fn is_empty(&self) -> bool {
        match self {
            Self::Stack(v) => v.is_empty(),
            #[cfg(feature = "alloc")]
            Self::Heap(v) => v.is_empty(),
        }
    }

    #[inline]
    #[must_use]
    pub fn as_slice(&self) -> &[T] {
        match self {
            Self::Stack(v) => v.as_slice(),
            #[cfg(feature = "alloc")]
            Self::Heap(v) => v.as_slice(),
        }
    }

    #[inline]
    pub fn as_mut_slice(&mut self) -> &mut [T] {
        match self {
            Self::Stack(v) => v.as_mut_slice(),
            #[cfg(feature = "alloc")]
            Self::Heap(v) => v.as_mut_slice(),
        }
    }

    /// Clear all items from the buffer.
    pub fn clear(&mut self) {
        match self {
            Self::Stack(v) => v.clear(),
            #[cfg(feature = "alloc")]
            Self::Heap(v) => v.clear(),
        }
    }

    /// Push a single item to the buffer.
    /// Will convert from Stack to Heap if capacity is exceeded (when alloc
    /// feature is enabled).
    pub fn push(&mut self, item: T)
    where
        T: Clone,
    {
        match self {
            Self::Stack(v) => {
                if v.push(item.clone()).is_err() {
                    #[cfg(feature = "alloc")]
                    {
                        // Convert to heap
                        let mut heap_vec = v.to_vec();
                        heap_vec.push(item);
                        *self = Self::Heap(heap_vec);
                    }
                    #[cfg(not(feature = "alloc"))]
                    {
                        panic!("Stack buffer is full and no heap available");
                    }
                }
            }
            #[cfg(feature = "alloc")]
            Self::Heap(v) => v.push(item),
        }
    }

    /// Extend the buffer with items from a slice.
    /// Will convert from Stack to Heap if capacity is exceeded (when alloc
    /// feature is enabled).
    pub fn extend_from_slice(&mut self, slice: &[T])
    where
        T: Clone,
    {
        match self {
            Self::Stack(v) => {
                if v.extend_from_slice(slice).is_err() {
                    #[cfg(feature = "alloc")]
                    {
                        // Convert to heap
                        let mut heap_vec = v.to_vec();
                        heap_vec.extend_from_slice(slice);
                        *self = Self::Heap(heap_vec);
                    }
                    #[cfg(not(feature = "alloc"))]
                    {
                        panic!("Stack buffer is full and no heap available");
                    }
                }
            }
            #[cfg(feature = "alloc")]
            Self::Heap(v) => v.extend_from_slice(slice),
        }
    }
}

impl<T, const N: usize> Deref for Vec<T, N>
where
    T: Zeroize,
{
    type Target = [T];

    fn deref(&self) -> &Self::Target {
        self.as_slice()
    }
}

impl<T, const N: usize> DerefMut for Vec<T, N>
where
    T: Zeroize,
{
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.as_mut_slice()
    }
}

impl<T, const N: usize> AsRef<[T]> for Vec<T, N>
where
    T: Zeroize,
{
    fn as_ref(&self) -> &[T] {
        self
    }
}

impl<T, const N: usize> AsMut<[T]> for Vec<T, N>
where
    T: Zeroize,
{
    fn as_mut(&mut self) -> &mut [T] {
        self
    }
}