sav 0.2.0

use std::{
    alloc::{self, Layout},
    fmt,
    marker::PhantomData,
    ops::Index,
    ptr::{NonNull, read, write},
};

/** a static length, copy-on-write vector type with element type X */
#[derive(Debug)]
pub struct A<'a, X> {
    /** ptr to the buffer of elements */
    pub ptr: NonNull<X>,
    /** vector length */
    pub L: u32,
    /** do we own this buffer? */
    pub own: bool,
    /** is this buffer reversed? */
    pub Rev: bool,
    _mark: PhantomData<&'a X>,
}

/** initialize an A out of a list of expressions */
#[macro_export]
macro_rules! mkA {
    [$($x:expr),* $(,)*] => {{
        let a = $crate::A::A::empty(${count($x)});
        $(unsafe { core::ptr::write(a.ptr_add(${index()}), $x); })*
        a
    }};
}

impl<'a, X> A<'a, X> {
    /** allocate an empty buffer of X. not really useful by itself. */
    pub fn empty(len: usize) -> A<'a, X> {
        /* the layout is just a array of width len */
        let lay = Layout::array::<X>(len).expect("invalid array layout");
        assert!(lay.size() <= isize::MAX as usize, "allocation too large");

        /* perform the allocation */
        let ptr = unsafe { alloc::alloc(lay) };
        let ptr = match NonNull::new(ptr as *mut X) {
            Some(x) => x,
            None => alloc::handle_alloc_error(lay),
        };

        A {
            ptr: ptr as NonNull<X>,
            L: len.try_into().expect("len too big to fit in u32"),
            own: true,
            Rev: false,
            _mark: PhantomData,
        }
    }

    /** make a non-owned A from a ptr with length L */
    pub fn from_ptr(ptr: NonNull<X>, L: u32, Rev: bool) -> A<'a, X> {
        A {
            ptr,
            L,
            own: false,
            Rev,
            _mark: PhantomData,
        }
    }

    /** add i to the buffer ptr */
    #[inline]
    pub unsafe fn ptr_add(&self, i: usize) -> *mut X {
        unsafe { self.ptr.as_ptr().add(i) }
    }

    /** add i to the buffer ptr and make it NonNull */
    #[inline]
    pub unsafe fn ptr_add_NN(&self, i: usize) -> NonNull<X> {
        let ptr = unsafe { self.ptr_add(i) };
        NonNull::new(ptr).expect("ptr_add_NN(): ptr is null!!")
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.L as usize
    }

    /** return the vector with the reverse flag set */
    #[inline]
    pub fn rev<'b>(&'b self) -> A<'a, X>
    where
        'b: 'a,
    {
        A {
            ptr: self.ptr,
            L: self.L,
            own: false,
            Rev: !self.Rev,
            _mark: PhantomData,
        }
    }

    /** reverse in place */
    #[inline]
    pub fn rev_mut(&mut self) {
        self.Rev = !self.Rev;
    }

    pub fn mkref<'b>(&'b self) -> A<'a, X>
    where
        'b: 'a,
    {
        A {
            ptr: self.ptr,
            L: self.L,
            own: false,
            Rev: self.Rev,
            _mark: PhantomData,
        }
    }

    /** set a value in the buffer. THIS WRevITE DOES NOT COPY.
     *
     * this function is only safe to use if you know that the buffer is
     * owned (ie you just allocated it with `empty`). otherwise, the change
     * will be propagated to all `A`s that have ever been derived from
     * this same buffer. */
    #[inline]
    pub unsafe fn set_nocpy(&mut self, i: usize, x: X) {
        if i < self.len() {
            unsafe {
                write(self.ptr_add(i), x);
            }
        }
    }

    /** get a value from the buffer without checking the boundary. */
    #[inline]
    pub unsafe fn at_nochk(&self, i: usize) -> X {
        unsafe { read(self.ptr_add(i)) }
    }

    /** drop elements from the beginning of a vector */
    pub fn drop<'b>(&'b self, n: u32) -> A<'a, X>
    where
        'b: 'a,
    {
        let ptr = unsafe { self.ptr_add_NN(n as usize) };
        A::from_ptr(ptr, self.L - n, self.Rev)
    }

    /** cut an A into a matrix by columns (vertically). */
    pub fn cutC<'b>(&'b self, w: u32) -> A<'a, A<'a, X>>
    where
        'b: 'a,
    {
        let rev = self.Rev;
        let L = self.L;
        /* the index we'll use for each row */
        let mut i = 0;
        /* return a vector with L/w rows */
        let mut r = A::<A<X>>::empty(L.div_ceil(w) as usize);
        /* make sure we maintain reversal */
        if rev {
            r.rev_mut();
        }

        while (i * w) + w <= L {
            /* point to the first element in the row */
            let P = unsafe { self.ptr_add_NN((i * w) as usize) };

            /* make the row */
            let R = A::from_ptr(P, w, rev);

            /* set the row in r */
            unsafe {
                write(r.ptr_add(i as usize), R);
            }

            i += 1;
        }

        r
    }
}

impl<'a, X> Clone for A<'a, X> {
    /** clone() on A is actually a copy/ref */
    #[inline]
    fn clone(&self) -> Self {
        A {
            ptr: self.ptr,
            L: self.L,
            own: false,
            Rev: self.Rev,
            _mark: PhantomData,
        }
    }
}

impl<'a, X> Drop for A<'a, X> {
    fn drop(&mut self) {
        /* we only dealloc if X is not a zst.
         * TODO: is this the correct behavior? stolen from rustnomicon. */
        let Z = size_of::<X>();
        if self.own && Z != 0 {
            unsafe {
                alloc::dealloc(
                    self.ptr.as_ptr() as *mut u8,
                    Layout::array::<X>(self.len())
                        .expect("invalid array layout"),
                )
            }
        }
    }
}

impl<'a, X> Index<usize> for A<'a, X> {
    type Output = X;

    #[inline]
    fn index(&self, idx: usize) -> &'a Self::Output {
        let L = self.len();
        if idx < L {
            /* offset based on the reverse flag */
            let off = if self.Rev { L - 1 - idx } else { idx };
            unsafe { &*self.ptr_add(off) }
        } else {
            panic!("{idx} is out of bounds for A of len {}", self.L);
        }
    }
}

impl<'a, X> Into<A<'a, X>> for Vec<X> {
    fn into(self) -> A<'a, X> {
        let L = self.len();
        let mut r = A::<X>::empty(L);
        for (i, x) in self.into_iter().enumerate() {
            unsafe {
                r.set_nocpy(i, x);
            }
        }
        r
    }
}

impl<'a, X> PartialEq for A<'a, X>
where
    X: PartialEq,
{
    fn eq(&self, y: &Self) -> bool {
        if self.L != y.L {
            return false;
        }

        for i in 0..self.len() {
            if self[i] != y[i] {
                return false;
            }
        }

        return true;
    }
}

impl<'a, X> fmt::Display for A<'a, X>
where
    X: fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let L = self.len();
        if !self.own {
            write!(f, "&")?;
        }
        if L == 0 {
            return write!(f, "()");
        }

        write!(f, "(")?;

        for i in 0..L - 1 {
            write!(f, "{};", self[i])?;
        }
        write!(f, "{})", self[L - 1])
    }
}

impl<'a, X> A<'a, X>
where
    X: Copy + Into<usize>,
    usize: Into<X>,
{
    pub fn iota(x: X) -> A<'a, X> {
        let mut r = A::empty(x.into());
        for i in 0..x.into() {
            unsafe {
                r.set_nocpy(i, i.into());
            }
        }
        r
    }
}

#[cfg(test)]
pub mod test {
    use crate::A;
    use std::mem::size_of;

    /* size */
    #[test]
    fn Z() {
        assert!(size_of::<A::A<u8>>() * 8 == 128);
    }

    /* read and write */
    #[test]
    fn rw() {
        let mut a = A::A::<u8>::empty(3);

        unsafe {
            a.set_nocpy(1, 2);
            assert_eq!(a.at_nochk(1), 2);
        }
    }

    /* mkA constructor */
    #[test]
    fn mk() {
        let a: A::A<u8> = mkA![1, 2, 3];
        let f = |i: usize, e: u8| assert_eq!(a[i], e);
        f(0, 1);
        f(1, 2);
        f(2, 3);
    }

    #[test]
    fn eq() {
        let a = mkA![1, 2, 3];
        let b = mkA![1, 2, 3];
        let c = mkA![1];
        let d = mkA![2, 3, 4];
        assert!(a == b);
        assert!(a == a);
        assert!(b == b);
        assert!(a != c);
        assert!(a != d);
    }

    #[test]
    fn rev() {
        let a = mkA![1, 2, 3];
        let mut b = mkA![3, 2, 1];
        assert!(a.rev() == b);
        b.rev_mut();
        assert!(a == b);
    }

    #[test]
    fn into() {
        let v = vec![1, 2, 3];
        let a = mkA![1, 2, 3];
        let b: A::A<i32> = v.into();
        let c = mkA![3, 2, 1];
        assert_eq!(b, a);
        assert_eq!(b.rev(), c);
    }

    #[test]
    fn cut() {
        let x = mkA![1, 2, 3, 4, 5, 6, 7, 8, 9];
        let x = x.rev();
        let x = x.cutC(3);

        let a = mkA![1, 2, 3];
        let b = mkA![4, 5, 6];
        let c = mkA![7, 8, 9];
        let y = mkA![a.rev(), b.rev(), c.rev()];
        let y = y.rev();

        println!("x: {x}, y: {y}");
        assert_eq!(x, y);
    }

    #[test]
    fn drop() {
        let a = mkA![1, 2, 3, 4];
        let a = a.drop(2);
        let b = mkA![3, 4];
        println!("a: {a}, b: {b}");
        assert_eq!(a, b);
    }
}