fdtd-coremem/crates/spirv_backend/src/support.rs

use core::ops::Index;

use coremem_types::compound::Optional;
use coremem_types::vec::Vec3u;

/// This struct allows doing things like *(ptr + offset) = value.
/// Such code wouldn't ordinarily compile with the spirv target because of
/// unsupported pointer math. RuntimeArray exists to overcome this, however
/// it emits invalid code for non-primitive types. Hence, this hack.
// XXX: maximum bytes an array may occupy is 0x7fff_ffff
// We don't know the element size, so assume it to be <= 44 bytes
// pub const MAX_UNSIZED_ARRAY: usize = 0x2e7_ffff;
pub const MAX_UNSIZED_ARRAY: usize = 0x1ff_ffff;
pub struct UnsizedArray<T>(pub [T; MAX_UNSIZED_ARRAY]);

impl<T: Copy> UnsizedArray<T> {
    /// user is responsible for bounds-checking the index before calling this.
    /// unsafe because this can trigger OOB reads if used incorrectly.
    pub unsafe fn index(&self, index: usize) -> T {
        // *self.0.index(index)
        self.0[index]
        // *self.0.get_unchecked(index)
        // &self.0
        //asm! {
        //    "%result = OpAccessChain _ {arr} {index}",
        //    "OpReturnValue %result",
        //    "%unused = OpLabel",
        //    arr = in(reg) self,
        //    index = in(reg) index,
        //}
    }
}

impl<T> UnsizedArray<T> {
    /// user is responsible for bounds-checking the index before calling this.
    /// unsafe because this can trigger OOB writes if used incorrectly.
    pub unsafe fn write(&mut self, index: usize, value: T) {
        self.0[index] = value;
    }
}

impl<T> UnsizedArray<T> {
    /// user is responsible for bounds-checking the index before calling this.
    /// unsafe because this can trigger OOB access if used incorrectly.
    pub unsafe fn index_ref(&self, index: usize) -> &T {
        &self.0[index]
    }

    pub unsafe fn get_handle<'a>(&'a self, index: usize) -> ArrayHandle<'a, T> {
        ArrayHandle {
            base: self,
            index,
        }
    }

    pub unsafe fn get_handle_mut<'a>(&'a mut self, index: usize) -> ArrayHandleMut<'a, T> {
        ArrayHandleMut {
            base: self,
            index,
        }
    }
}

pub struct ArrayHandle<'a, T> {
    base: &'a UnsizedArray<T>,
    index: usize,
}

impl<'a, T: Copy> ArrayHandle<'a, T> {
    pub fn get(&self) -> T {
        unsafe {
            self.base.index(self.index)
        }
    }
    pub fn get_into(&self, out: &mut T) {
        core::clone::Clone::clone_from(out, self.get_ref());
    }
}
impl<'a, T> ArrayHandle<'a, T> {
    pub fn get_ref(&self) -> &'a T {
        unsafe {
            self.base.index_ref(self.index)
        }
    }
}


pub struct ArrayHandleMut<'a, T> {
    base: &'a mut UnsizedArray<T>,
    index: usize,
}

impl<'a, T: Copy> ArrayHandleMut<'a, T> {
    pub fn write(&mut self, value: T) {
        unsafe {
            self.base.write(self.index, value);
        }
    }
    pub fn get(&self) -> T {
        unsafe {
            self.base.index(self.index)
        }
    }
}

/// 3d dynamically-sized array backed by a borrowed buffer.
#[derive(Clone, Copy)]
pub struct Array3<'a, T> {
    data: &'a UnsizedArray<T>,
    dim: Vec3u,
}

fn index(loc: Vec3u, dim: Vec3u) -> usize {
    ((loc.z()*dim.y() + loc.y())*dim.x() + loc.x()) as usize
}

fn checked_index(idx: Vec3u, dim: Vec3u) -> Optional<usize> {
    if idx.x() < dim.x() && idx.y() < dim.y() && idx.z() < dim.z() {
        let flat_idx = index(idx, dim);
        Optional::some(flat_idx)
    } else {
        Optional::none()
    }
}

impl<'a, T> Array3<'a, T> {
    pub fn new(data: &'a UnsizedArray<T>, dim: Vec3u) -> Self {
        Self {
            data,
            dim,
        }
    }

    pub fn index(self, idx: Vec3u) -> Optional<usize> {
        checked_index(idx, self.dim)
    }

    pub fn into_handle(self, idx: Vec3u) -> ArrayHandle<'a, T> {
        let idx = checked_index(idx, self.dim).unwrap();
        unsafe {
            self.data.get_handle(idx)
        }
    }

    pub fn into_ref(self, idx: Vec3u) -> &'a T {
        let idx = checked_index(idx, self.dim).unwrap();
        unsafe {
            self.data.index_ref(idx)
        }
    }
}

impl<'a, T: Copy + Default> Array3<'a, T> {
    pub fn get(&self, idx: Vec3u) -> Optional<T> {
        let idx = checked_index(idx, self.dim);
        if idx.is_some() {
            Optional::some(unsafe {
                self.data.index(idx.unwrap())
            })
        } else {
            Optional::none()
        }
    }
}

impl<'a, T> Index<Vec3u> for Array3<'a, T> {
    type Output = T;
    fn index(&self, idx: Vec3u) -> &Self::Output {
        let idx = checked_index(idx, self.dim).unwrap();
        unsafe {
            self.data.index_ref(idx)
        }
    }
}

/// 3d dynamically-sized array backed by a mutably borrowed buffer.
pub struct Array3Mut<'a, T> {
    data: &'a mut UnsizedArray<T>,
    dim: Vec3u,
}

impl<'a, T> Array3Mut<'a, T> {
    pub fn new(data: &'a mut UnsizedArray<T>, dim: Vec3u) -> Self {
        Self {
            data,
            dim,
        }
    }

    // fn as_array3<'b>(&'b self) -> Array3<'b, T> {
    //     Array3::new(self.data, self.dim)
    // }

    pub fn index(&self, idx: Vec3u) -> Optional<usize> {
        if idx.x() < self.dim.x() && idx.y() < self.dim.y() && idx.z() < self.dim.z() {
            let flat_idx = index(idx, self.dim);
            Optional::some(flat_idx)
        } else {
            Optional::none()
        }
    }
}

impl<'a, T: Copy> Array3Mut<'a, T> {
    pub fn into_mut_handle(self, idx: Vec3u) -> ArrayHandleMut<'a, T> {
        let idx = self.index(idx).unwrap();
        unsafe {
            self.data.get_handle_mut(idx)
        }
    }
}


#[cfg(test)]
mod test {
    use super::*;
    #[test]
    fn test_index() {
        let dim = Vec3u::new(2, 3, 7);
        assert_eq!(index(Vec3u::new(0, 0, 0), dim), 0);
        assert_eq!(index(Vec3u::new(1, 0, 0), dim), 1);
        assert_eq!(index(Vec3u::new(0, 1, 0), dim), 2);
        assert_eq!(index(Vec3u::new(1, 1, 0), dim), 3);
        assert_eq!(index(Vec3u::new(0, 2, 0), dim), 4);
        assert_eq!(index(Vec3u::new(0, 0, 1), dim), 6);
        assert_eq!(index(Vec3u::new(1, 0, 1), dim), 7);
        assert_eq!(index(Vec3u::new(0, 1, 1), dim), 8);
        assert_eq!(index(Vec3u::new(1, 2, 1), dim), 11);
        assert_eq!(index(Vec3u::new(1, 2, 2), dim), 17);
    }
}