nordisk/nornir

crates/mimir-graph/src/vector.rs

use std::marker::PhantomData;
use std::ptr;

use graphblas_sys as sys;

use crate::error::{Error, Result, check, check_value, grb};
use crate::runtime;
use crate::types::Scalar;

/// An owned `GrB_Vector` of element type `T` (freed on drop).
pub struct Vector<T: Scalar> {
    pub(crate) raw: sys::GrB_Vector,
    _t: PhantomData<T>,
}

unsafe impl<T: Scalar> Send for Vector<T> {}

impl<T: Scalar> Drop for Vector<T> {
    fn drop(&mut self) {
        if !self.raw.is_null() {
            unsafe { sys::GrB_Vector_free(&mut self.raw) };
        }
    }
}

impl<T: Scalar> Vector<T> {
    /// Empty vector of size `n`.
    pub fn new(n: u64) -> Result<Self> {
        runtime::init()?;
        let mut raw = ptr::null_mut();
        grb!(GrB_Vector_new(&mut raw, T::grb_type(), n))?;
        Ok(Self {
            raw,
            _t: PhantomData,
        })
    }

    /// From `(index, value)` tuples; duplicates combined (OR / sum).
    pub fn from_tuples(n: u64, idx: &[u64], vals: &[T]) -> Result<Self> {
        if idx.len() != vals.len() {
            return Err(Error::invalid(format!(
                "idx {} != vals {}",
                idx.len(),
                vals.len()
            )));
        }
        let v = Self::new(n)?;
        check("GrB_Vector_build", unsafe {
            T::vector_build(v.raw, idx.as_ptr(), vals.as_ptr(), idx.len() as u64)
        })?;
        Ok(v)
    }

    /// Every listed index set to `T::ONE` (a seed set / scope set).
    pub fn from_indices(n: u64, idx: &[u64]) -> Result<Self> {
        Self::from_tuples(n, idx, &vec![T::ONE; idx.len()])
    }

    /// Dense vector: entry `i` = `vals[i]`.
    pub fn from_dense(vals: &[T]) -> Result<Self> {
        let idx: Vec<u64> = (0..vals.len() as u64).collect();
        Self::from_tuples(vals.len() as u64, &idx, vals)
    }

    pub fn size(&self) -> Result<u64> {
        let mut n = 0;
        grb!(GrB_Vector_size(&mut n, self.raw))?;
        Ok(n)
    }

    pub fn nvals(&self) -> Result<u64> {
        let mut n = 0;
        grb!(GrB_Vector_nvals(&mut n, self.raw))?;
        Ok(n)
    }

    pub fn set(&mut self, i: u64, x: T) -> Result<()> {
        check("GrB_Vector_setElement", unsafe {
            T::vector_set_element(self.raw, x, i)
        })
    }

    pub fn get(&self, i: u64) -> Result<Option<T>> {
        let mut x = T::default();
        let present = check_value("GrB_Vector_extractElement", unsafe {
            T::vector_extract_element(&mut x, self.raw, i)
        })?;
        Ok(present.then_some(x))
    }

    /// Stored entries as `(indices, values)`, indices ascending.
    pub fn to_tuples(&self) -> Result<(Vec<u64>, Vec<T>)> {
        let n = self.nvals()?;
        let mut i = vec![0u64; n as usize];
        let mut x = vec![T::default(); n as usize];
        let mut got = n;
        check("GrB_Vector_extractTuples", unsafe {
            T::vector_extract_tuples(i.as_mut_ptr(), x.as_mut_ptr(), &mut got, self.raw)
        })?;
        i.truncate(got as usize);
        x.truncate(got as usize);
        // SuiteSparse returns sorted indices for a finished vector; make it a contract.
        if !i.is_sorted() {
            let mut p: Vec<(u64, T)> = i.into_iter().zip(x).collect();
            p.sort_unstable_by_key(|e| e.0);
            return Ok(p.into_iter().unzip());
        }
        Ok((i, x))
    }

    /// Indices of the stored entries, ascending.
    pub fn indices(&self) -> Result<Vec<u64>> {
        Ok(self.to_tuples()?.0)
    }

    pub fn dup(&self) -> Result<Self> {
        let mut raw = ptr::null_mut();
        grb!(GrB_Vector_dup(&mut raw, self.raw))?;
        Ok(Self {
            raw,
            _t: PhantomData,
        })
    }

    pub fn as_raw(&self) -> sys::GrB_Vector {
        self.raw
    }
}

impl<T: Scalar> std::fmt::Debug for Vector<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Vector")
            .field("type", &std::any::type_name::<T>())
            .field("size", &self.size().ok())
            .field("nvals", &self.nvals().ok())
            .finish()
    }
}