boreholeio 0.1.0

mod array_proxy;
mod typed_array;
mod utils;

use crate::schema::array_uri_reference::ArrayUriReference;
use ndarray::ArrayViewD;
use std::{
    io::{Read, Write},
    path::{Path, PathBuf},
    sync::RwLock,
};

pub(crate) use utils::{convert, convert_vec};

pub use array_proxy::{ArrayProxy, ProxyCreator, ProxyFunctor, Statistics};
pub use typed_array::TypedArray;
pub use utils::Error;

#[derive(Debug)]
pub struct StridedArrayFile<'a> {
    path: PathBuf,
    file_mmap: RwLock<Option<memmap2::Mmap>>,
    arrays: RwLock<Option<Vec<ArrayProxy<'a>>>>,
}

impl<'a> StridedArrayFile<'a> {
    pub fn new(path: PathBuf) -> Self {
        Self {
            path,
            file_mmap: RwLock::new(None),
            arrays: RwLock::new(None),
        }
    }

    pub fn is_open(&self) -> bool {
        self.file_mmap.read().unwrap().is_some() && self.arrays.read().unwrap().is_some()
    }

    pub fn write(file_path: &Path, arrays: &[ArrayProxy<'a>]) -> Result<(), std::io::Error> {
        Self::validate_system_endianness()?;
        let mut f = std::fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(file_path)?;
        Header {
            n_arrays: convert(arrays.len()),
        }
        .write(&mut f)?;
        let header_size = 10
            + arrays
                .iter()
                .map(|arr| 13 + 16 * arr.dimensions())
                .sum::<usize>();
        struct ArrayDataParams {
            start: usize,
            size: usize,
        }
        let mut array_data_params = Vec::<ArrayDataParams>::with_capacity(arrays.len());
        let mut data_start = header_size;
        for a in arrays {
            f.write_all(&convert::<usize, u32>(a.dimensions()).to_le_bytes())?;
            let shape = a
                .shape()
                .iter()
                .map(|s| convert::<usize, u64>(*s))
                .collect::<Vec<u64>>();
            write_vec(&shape, &mut f)?;
            let elem_size = a.element_size();
            data_start = data_start.div_ceil(elem_size) * elem_size;
            let strides_in_bytes = convert_vec::<isize, i64>(&a.strides(true));
            write_vec(&strides_in_bytes, &mut f)?;
            f.write_all(&a.element_type().to_le_bytes())?;
            // TODO take negative strides (span/data_start) into account
            f.write_all(&convert::<usize, u64>(data_start).to_le_bytes())?;
            let data_size = a.data_size();
            array_data_params.push(ArrayDataParams {
                start: data_start,
                size: data_size,
            });
            data_start += data_size;
        }
        let mut cursor = header_size;
        for (arr, params) in std::iter::zip(arrays, array_data_params) {
            if params.start > cursor {
                f.write_all(&vec![0u8; params.start - cursor][..])?;
            }
            f.write_all(unsafe {
                // Note that the from_raw_parts method requires the memory to be nicely
                // aligned, which is guaranteed in this particular case
                std::slice::from_raw_parts(arr.data(), params.size)
            })?;
            cursor = params.start + params.size;
        }
        Ok(())
    }

    pub fn path(&self) -> &Path {
        &self.path
    }

    pub fn count(&self) -> Result<usize, Error> {
        if !self.is_open() {
            self.open()?
        }
        Ok(self.arrays.read().unwrap().as_ref().unwrap().len())
    }

    pub fn try_get(&self, i: usize) -> Result<ArrayProxy<'a>, Error> {
        if i >= self.count()? {
            Err(format!(
                "Array index {i} is greater than file size {}",
                self.count()?
            ))
        } else {
            Ok(self.arrays.read().unwrap().as_ref().unwrap()[i].clone())
        }
    }

    pub fn try_get_as<T: 'static>(&self, i: usize) -> Result<ArrayViewD<'a, T>, Error> {
        self.try_get(i)?.try_as::<T>()
    }

    pub fn get(&self, i: usize) -> ArrayProxy<'a> {
        self.try_get(i).unwrap()
    }

    pub fn get_as<T: 'static>(&self, i: usize) -> ArrayViewD<'a, T> {
        self.try_get_as(i).unwrap()
    }

    fn validate_system_endianness() -> Result<(), std::io::Error> {
        // https://doc.rust-lang.org/reference/conditional-compilation.html#target_endian
        if !cfg!(target_endian = "little") {
            return Err(std::io::Error::new(
                std::io::ErrorKind::Unsupported,
                "The OS is not little endian",
            ));
        }
        Ok(())
    }

    fn open(&self) -> Result<(), Error> {
        if let Err(err) = self.open_unchecked() {
            self.arrays.write().unwrap().take();
            self.file_mmap.write().unwrap().take();
            Err(err)
        } else {
            Ok(())
        }
    }

    /// This method does not clean up `self` if an error occurs midway. One
    /// must always call `self.open` instead.
    fn open_unchecked(&self) -> Result<(), Error> {
        Self::validate_system_endianness().map_err(|e| e.to_string())?;
        if !self.path.is_file() {
            return Err(format!(
                "Invalid path to a Strided Array File: {:?}",
                self.path
            ));
        }
        let mut f = std::fs::OpenOptions::new()
            .read(true)
            .create(false)
            .open(&self.path)
            .map_err(|e| e.to_string())?;
        let mmap = unsafe { memmap2::Mmap::map(&f).unwrap() };
        let n_arrays = Header::read(&mut f).map_err(|e| e.to_string())?.n_arrays;
        let mut offset = Header::N_BYTES;
        let mut arrays = Vec::<ArrayProxy>::new();
        for _ in 0..n_arrays {
            let n_dims = convert(u32::from_le_bytes(
                mmap[offset..offset + 4].try_into().unwrap(),
            ));
            offset += 4;
            let shape = unpack_vec_non_aligned::<u64, _>(
                &mmap[offset..offset + 8 * n_dims],
                n_dims,
                |bytes| u64::from_le_bytes(bytes.try_into().unwrap()),
            );
            offset += 8 * n_dims;
            let strides_in_bytes = unpack_vec_non_aligned::<i64, _>(
                &mmap[offset..offset + 8 * n_dims],
                n_dims,
                |bytes| i64::from_le_bytes(bytes.try_into().unwrap()),
            );
            offset += 8 * n_dims;
            let element_type = mmap[offset];
            let data_start = u64::from_le_bytes(mmap[offset + 1..offset + 9].try_into().unwrap());
            offset += 9;
            arrays.push(ArrayProxy::from_mmap(
                &mmap,
                shape,
                strides_in_bytes,
                element_type,
                data_start,
            )?);
        }
        self.file_mmap.write().unwrap().replace(mmap);
        self.arrays.write().unwrap().replace(arrays);
        Ok(())
    }
}

/// Returns the array index (if present) specified in the URI.
///
/// For example:
/// - Returns `Ok(None)` for `"some/path.star"`,
/// - Returns `Ok(None)` for `"some/path.png"`,
/// - Returns `Ok(Some(10))` for `"https://example.com/some/path.star#10"`,
/// - Returns `Err` for `"some/path.png#eleven"`.
pub fn extract_array_index(uri: &ArrayUriReference) -> Result<Option<usize>, Error> {
    // See if there is anything after `#`.
    if let Some(fragment) = uri.fragment() {
        // Convert `&str` to `usize`.
        Ok(Some(fragment.as_str().parse::<usize>().map_err(|e| {
            format!("{fragment} is not a valid integer: {e}")
        })?))
    } else {
        Ok(None)
    }
}

#[derive(Debug, PartialEq)]
struct Version {
    major: u8,
    minor: u8,
}

// C-represented header will introduce some padding. Adding macroses:
// #[repr(packed(1))]
// #[derive(Debug, serde::Deserialize)]
// generates an error:
// cannot move out of `self.version` which is behind a shared reference
//`#[derive(Debug)]` triggers a move because taking references to the fields of a packed
// struct is undefined behaviour.
#[derive(Debug)]
struct Header {
    n_arrays: u32,
}

impl Header {
    const MAGIC_BYTES: &'static [u8; 4] = b"StAr";
    const VERSION: Version = Version { major: 0, minor: 2 };
    const N_BYTES: usize = 10;

    fn read(f: &mut std::fs::File) -> Result<Self, std::io::Error> {
        assert_eq!(Self::N_BYTES, 10);
        let mut magic_bytes = [0; 4];
        f.read_exact(&mut magic_bytes)?;
        if &magic_bytes != Self::MAGIC_BYTES {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "Magic bytes validation failed",
            ));
        }
        let mut version = [0; 2];
        f.read_exact(&mut version)?;
        let version = Version {
            major: version[0],
            minor: version[1],
        };
        if version != Self::VERSION {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("Unsupported version: {version:?}"),
            ));
        }
        let mut n_arrays = [0; 4];
        f.read_exact(&mut n_arrays)?;
        Ok(Self {
            n_arrays: u32::from_le_bytes(n_arrays),
        })
    }

    fn write(&self, f: &mut std::fs::File) -> Result<(), std::io::Error> {
        assert_eq!(Self::N_BYTES, 10);
        f.write_all(Self::MAGIC_BYTES)?;
        f.write_all(&Self::VERSION.major.to_le_bytes())?;
        f.write_all(&Self::VERSION.minor.to_le_bytes())?;
        f.write_all(&self.n_arrays.to_le_bytes())?;
        Ok(())
    }
}

// This method is used only twice: with converter being either `i64::from_le_bytes` or
// `u64::from_le_bytes`. The thing is that, those 2 methods don't have any common trait
// and, thus, needs to be passed as a closure.
fn unpack_vec_non_aligned<T, F>(bytes: &[u8], n_elems: usize, converter: F) -> Vec<T>
where
    F: Fn(&[u8]) -> T,
{
    let elem_size = std::mem::size_of::<T>();
    let mut vec = Vec::<T>::with_capacity(n_elems);
    let mut offset = 0;
    for _ in 0..n_elems {
        vec.push(converter(&bytes[offset..offset + elem_size]));
        offset += elem_size;
    }
    vec
}

fn write_vec<T>(vec: &[T], file: &mut std::fs::File) -> Result<(), std::io::Error> {
    let bytes: &[u8] = unsafe {
        // Note that the from_raw_parts method requires the memory to be nicely
        // aligned, which is guaranteed in this particular case
        std::slice::from_raw_parts(vec.as_ptr() as *const u8, std::mem::size_of_val(vec))
    };
    file.write_all(bytes)
}