feat: add GraphBLAS serialization primitives

Cargo.toml

@@ -28,7 +28,9 @@ graph = { path = "graph", version = "0.1.0" }
 lazy_static = "1.5.0"
 parking_lot = "0.12.5"
 redis-module = { git = "https://github.com/AviAvni/redismodule-rs", branch = "master" }
+roaring = "0.11.3"
 ryu = "1.0.23"
+thin-vec = "0.2.14"
 orx-tree = "2.2.0"
 
 [build-dependencies]

graph/src/graph/graphblas/matrix.rs

@@ -54,12 +54,25 @@
 
 #![allow(clippy::doc_markdown)]
 
-use std::{mem::MaybeUninit, os::raw::c_void, ptr::null_mut, sync::Arc};
+use std::{
+    marker::PhantomData,
+    mem::{ManuallyDrop, MaybeUninit},
+    os::raw::c_void,
+    ptr::null_mut,
+    sync::Arc,
+};
 
 use parking_lot::Mutex;
 
-use crate::graph::graphblas::lagraph_bindings::{LAGraph_Finalize, LAGraph_Init};
+use crate::graph::graphblas::{
+    lagraph_bindings::{LAGraph_Finalize, LAGraph_Init},
+    serialization::{Decode, Encode, Reader, Writer},
+};
+
+/// Size of the `GxB_Container_struct` in bytes.
+const CONTAINER_STRUCT_SIZE: usize = std::mem::size_of::<super::GxB_Container_struct>();
 
+use super::vector::Vector;
 use super::{
     GrB_BOOL, GrB_DESC_C, GrB_DESC_CT0, GrB_DESC_CT0T1, GrB_DESC_CT1, GrB_DESC_R, GrB_DESC_RC,
     GrB_DESC_RCT0, GrB_DESC_RCT0T1, GrB_DESC_RCT1, GrB_DESC_RS, GrB_DESC_RSC, GrB_DESC_RSCT0,
@@ -68,14 +81,16 @@ use super::{
     GrB_DESC_SCT1, GrB_DESC_ST0, GrB_DESC_ST0T1, GrB_DESC_ST1, GrB_DESC_T0, GrB_DESC_T0T1,
     GrB_DESC_T1, GrB_Descriptor, GrB_GLOBAL, GrB_Global_set_INT32, GrB_Info, GrB_Matrix,
     GrB_Matrix_clear, GrB_Matrix_dup, GrB_Matrix_eWiseAdd_Semiring, GrB_Matrix_eWiseMult_Semiring,
-    GrB_Matrix_extractElement_BOOL, GrB_Matrix_free, GrB_Matrix_get_INT32, GrB_Matrix_ncols,
-    GrB_Matrix_new, GrB_Matrix_nrows, GrB_Matrix_nvals, GrB_Matrix_removeElement,
-    GrB_Matrix_resize, GrB_Matrix_setElement_BOOL, GrB_Matrix_wait, GrB_Mode, GrB_WaitMode,
-    GrB_finalize, GrB_mxm, GrB_transpose, GxB_ANY_BOOL, GxB_ANY_PAIR_BOOL, GxB_Iterator,
-    GxB_Iterator_free, GxB_Iterator_new, GxB_Matrix_fprint, GxB_Matrix_memoryUsage,
-    GxB_Option_Field, GxB_Print_Level, GxB_init, GxB_rowIterator_attach,
-    GxB_rowIterator_getColIndex, GxB_rowIterator_getRowIndex, GxB_rowIterator_nextCol,
-    GxB_rowIterator_nextRow, GxB_rowIterator_seekRow,
+    GrB_Matrix_extractElement_BOOL, GrB_Matrix_extractElement_UINT64, GrB_Matrix_free,
+    GrB_Matrix_get_INT32, GrB_Matrix_ncols, GrB_Matrix_new, GrB_Matrix_nrows, GrB_Matrix_nvals,
+    GrB_Matrix_removeElement, GrB_Matrix_resize, GrB_Matrix_setElement_BOOL,
+    GrB_Matrix_setElement_UINT64, GrB_Matrix_wait, GrB_Mode, GrB_UINT64, GrB_WaitMode,
+    GrB_finalize, GrB_mxm, GrB_transpose, GxB_ANY_BOOL, GxB_ANY_PAIR_BOOL, GxB_Container_free,
+    GxB_Container_new, GxB_Iterator, GxB_Iterator_free, GxB_Iterator_new, GxB_Matrix_fprint,
+    GxB_Matrix_memoryUsage, GxB_Matrix_type, GxB_Option_Field, GxB_Print_Level, GxB_init,
+    GxB_load_Matrix_from_Container, GxB_rowIterator_attach, GxB_rowIterator_getColIndex,
+    GxB_rowIterator_getRowIndex, GxB_rowIterator_nextCol, GxB_rowIterator_nextRow,
+    GxB_rowIterator_seekRow, GxB_unload_Matrix_into_Container,
 };
 
 /// Initializes the GraphBLAS library in non-blocking mode.
@@ -419,6 +434,104 @@ impl Drop for Matrix {
     }
 }
 
+impl Decode<19> for Matrix {
+    fn decode(r: &mut dyn Reader) -> Result<Self, String> {
+        let container_bytes = r.read_buffer()?;
+
+        // Validate container size before copying
+        if container_bytes.len() < CONTAINER_STRUCT_SIZE {
+            return Err(format!(
+                "container buffer too small: {} bytes < {} bytes required",
+                container_bytes.len(),
+                CONTAINER_STRUCT_SIZE
+            ));
+        }
+
+        unsafe {
+            let mut container: MaybeUninit<super::GxB_Container> = MaybeUninit::uninit();
+            let info = GxB_Container_new(container.as_mut_ptr());
+            debug_assert_eq!(info, GrB_Info::GrB_SUCCESS);
+            let container = container.assume_init();
+
+            // Copy struct data into the allocated container
+            std::ptr::copy_nonoverlapping(
+                container_bytes.as_ptr(),
+                container.cast::<u8>(),
+                CONTAINER_STRUCT_SIZE,
+            );
+
+            // Nullify vector/matrix pointers (will be populated below)
+            (*container).x = null_mut();
+            (*container).h = null_mut();
+            (*container).b = null_mut();
+            (*container).i = null_mut();
+            (*container).p = null_mut();
+            (*container).Y = null_mut();
+
+            // Read and load 5 vectors: x, h, p, i, b
+            (*container).x = ManuallyDrop::new(Vector::<bool>::decode(r)?).ptr();
+            (*container).h = ManuallyDrop::new(Vector::<bool>::decode(r)?).ptr();
+            (*container).p = ManuallyDrop::new(Vector::<bool>::decode(r)?).ptr();
+            (*container).i = ManuallyDrop::new(Vector::<bool>::decode(r)?).ptr();
+            (*container).b = ManuallyDrop::new(Vector::<bool>::decode(r)?).ptr();
+
+            // Create matrix and load from container
+            let mut m: MaybeUninit<GrB_Matrix> = MaybeUninit::uninit();
+            let info = GrB_Matrix_new(m.as_mut_ptr(), GrB_BOOL, 0, 0);
+            debug_assert_eq!(info, GrB_Info::GrB_SUCCESS);
+            let m = m.assume_init();
+
+            let info = GxB_load_Matrix_from_Container(m, container, null_mut());
+            debug_assert_eq!(info, GrB_Info::GrB_SUCCESS);
+
+            let mut c = container;
+            let info = GxB_Container_free(&raw mut c);
+            debug_assert_eq!(info, GrB_Info::GrB_SUCCESS);
+
+            Ok(Self {
+                m: Arc::new(m),
+                lock: Arc::new(Mutex::new(())),
+            })
+        }
+    }
+}
+
+impl Encode<19> for Matrix {
+    fn encode(
+        &self,
+        w: &mut dyn Writer,
+    ) {
+        unsafe {
+            let mut container: MaybeUninit<super::GxB_Container> = MaybeUninit::uninit();
+            let info = GxB_Container_new(container.as_mut_ptr());
+            debug_assert_eq!(info, GrB_Info::GrB_SUCCESS);
+            let container = container.assume_init();
+
+            let info = GxB_unload_Matrix_into_Container(self.inner(), container, null_mut());
+            debug_assert_eq!(info, GrB_Info::GrB_SUCCESS);
+
+            // Write container struct bytes
+            let container_bytes =
+                std::slice::from_raw_parts(container.cast::<u8>(), CONTAINER_STRUCT_SIZE);
+            w.write_buffer(container_bytes);
+
+            // Write 5 vectors: x, h, p, i, b
+            ManuallyDrop::new(Vector::<bool>::from((*container).x)).encode(w);
+            ManuallyDrop::new(Vector::<bool>::from((*container).h)).encode(w);
+            ManuallyDrop::new(Vector::<bool>::from((*container).p)).encode(w);
+            ManuallyDrop::new(Vector::<bool>::from((*container).i)).encode(w);
+            ManuallyDrop::new(Vector::<bool>::from((*container).b)).encode(w);
+
+            let info = GxB_load_Matrix_from_Container(self.inner(), container, null_mut());
+            debug_assert_eq!(info, GrB_Info::GrB_SUCCESS);
+
+            let mut c = container;
+            let info = GxB_Container_free(&raw mut c);
+            debug_assert_eq!(info, GrB_Info::GrB_SUCCESS);
+        }
+    }
+}
+
 impl Matrix {
     /// Returns the raw GrB_Matrix handle for FFI calls (e.g. LAGraph).
     /// The caller must NOT free the returned handle.
@@ -427,6 +540,29 @@ impl Matrix {
         *self.m
     }
 
+    /// Iterate entries as `(row, col, value)` UINT64 triples.
+    ///
+    /// Used when loading C-produced relation matrices where single-edge
+    /// entries store the edge ID as a UINT64 value.
+    #[must_use]
+    pub fn uint64_iter(&self) -> Iter<Uint64Extract> {
+        Iter::new(self, 0, u64::MAX)
+    }
+
+    /// Returns true if this matrix has UINT64 element type.
+    ///
+    /// C-produced relation matrices store edge IDs as UINT64, while
+    /// Rust-produced ones use BOOL.
+    #[must_use]
+    pub fn is_uint64(&self) -> bool {
+        unsafe {
+            let mut t: MaybeUninit<super::GrB_Type> = MaybeUninit::uninit();
+            let info = GxB_Matrix_type(t.as_mut_ptr(), *self.m);
+            debug_assert_eq!(info, GrB_Info::GrB_SUCCESS);
+            t.assume_init() == GrB_UINT64
+        }
+    }
+
     #[must_use]
     pub fn pending(&self) -> bool {
         unsafe {
@@ -476,6 +612,17 @@ impl Matrix {
             debug_assert_eq!(info, GrB_Info::GrB_SUCCESS);
         }
     }
+
+    pub fn select(
+        &mut self,
+        mask: &Matrix,
+        a: &Matrix,
+    ) {
+        unsafe {
+            let info = GrB_transpose(*self.m, *mask.m, null_mut(), *a.m, GrB_DESC_RCT0);
+            debug_assert_eq!(info, GrB_Info::GrB_SUCCESS);
+        }
+    }
 }
 
 impl Size for Matrix {
@@ -561,6 +708,36 @@ impl Dup<Self> for Matrix {
 }
 
 impl Matrix {
+    /// Create a new UINT64 matrix (for C-compatible tensor encoding).
+    #[must_use]
+    pub fn new_uint64(
+        nrows: u64,
+        ncols: u64,
+    ) -> Self {
+        unsafe {
+            let mut m: MaybeUninit<GrB_Matrix> = MaybeUninit::uninit();
+            let info = GrB_Matrix_new(m.as_mut_ptr(), GrB_UINT64, nrows, ncols);
+            debug_assert_eq!(info, GrB_Info::GrB_SUCCESS);
+            Self {
+                m: Arc::new(m.assume_init()),
+                lock: Arc::new(Mutex::new(())),
+            }
+        }
+    }
+
+    /// Set a UINT64 value at (i, j).
+    pub fn set_uint64(
+        &mut self,
+        i: u64,
+        j: u64,
+        value: u64,
+    ) {
+        unsafe {
+            let info = GrB_Matrix_setElement_UINT64(*self.m, value, i, j);
+            debug_assert_eq!(info, GrB_Info::GrB_SUCCESS);
+        }
+    }
+
     #[must_use]
     #[allow(clippy::iter_without_into_iter)]
     pub fn iter(
@@ -655,20 +832,71 @@ where
     }
 }
 
-pub struct Iter {
+/// Strategy for extracting values from a GraphBLAS row iterator position.
+///
+/// # Safety
+/// Implementations must only call valid GraphBLAS FFI functions on the provided matrix.
+pub trait IterExtract {
+    type Item;
+
+    /// Extract the item from the current iterator position.
+    ///
+    /// # Safety
+    /// `m` must be a valid `GrB_Matrix` and the iterator must be positioned on a valid entry.
+    unsafe fn extract(
+        m: GrB_Matrix,
+        row: u64,
+        col: u64,
+    ) -> Self::Item;
+}
+
+/// Extracts `(row, col)` pairs from a boolean matrix.
+pub struct BoolExtract;
+
+impl IterExtract for BoolExtract {
+    type Item = (u64, u64);
+
+    unsafe fn extract(
+        _m: GrB_Matrix,
+        row: u64,
+        col: u64,
+    ) -> Self::Item {
+        (row, col)
+    }
+}
+
+/// Extracts `(row, col, value)` triples from a UINT64 matrix.
+pub struct Uint64Extract;
+
+impl IterExtract for Uint64Extract {
+    type Item = (u64, u64, u64);
+
+    unsafe fn extract(
+        m: GrB_Matrix,
+        row: u64,
+        col: u64,
+    ) -> Self::Item {
+        let mut val: u64 = 0;
+        unsafe { GrB_Matrix_extractElement_UINT64(&raw mut val, m, row, col) };
+        (row, col, val)
+    }
+}
+
+pub struct Iter<E: IterExtract = BoolExtract> {
     m: Arc<GrB_Matrix>,
     /// The underlying GraphBLAS iterator.
     inner: GxB_Iterator,
     /// Indicates whether the iterator is depleted.
     depleted: bool,
     /// The maximum row index for the iterator.
     max_row: u64,
+    _extract: PhantomData<E>,
 }
 
-unsafe impl Send for Iter {}
-unsafe impl Sync for Iter {}
+unsafe impl<E: IterExtract> Send for Iter<E> {}
+unsafe impl<E: IterExtract> Sync for Iter<E> {}
 
-impl Drop for Iter {
+impl<E: IterExtract> Drop for Iter<E> {
     /// Frees the GraphBLAS iterator when the `Iter` is dropped.
     fn drop(&mut self) {
         unsafe {
@@ -681,7 +909,7 @@ impl Drop for Iter {
     }
 }
 
-impl Iter {
+impl<E: IterExtract> Iter<E> {
     /// Creates a new iterator for traversing all elements in a matrix.
     ///
     /// # Parameters
@@ -716,18 +944,19 @@ impl Iter {
                 depleted: info != GrB_Info::GrB_SUCCESS
                     || GxB_rowIterator_getRowIndex(iter) > max_row,
                 max_row,
+                _extract: PhantomData,
             }
         }
     }
 }
 
-impl Iterator for Iter {
-    type Item = (u64, u64);
+impl<E: IterExtract> Iterator for Iter<E> {
+    type Item = E::Item;
 
     /// Advances the iterator and returns the next element in the matrix.
     ///
     /// # Returns
-    /// - `Some((u64, u64))`: The next element in the matrix.
+    /// - `Some(E::Item)`: The next element in the matrix.
     /// - `None`: The iterator is depleted.
     fn next(&mut self) -> Option<Self::Item> {
         if self.depleted {
@@ -736,6 +965,7 @@ impl Iterator for Iter {
         unsafe {
             let row = GxB_rowIterator_getRowIndex(self.inner);
             let col = GxB_rowIterator_getColIndex(self.inner);
+            let item = E::extract(*self.m, row, col);
             if GxB_rowIterator_nextCol(self.inner) != GrB_Info::GrB_SUCCESS {
                 let mut info = GxB_rowIterator_nextRow(self.inner);
                 debug_assert!(
@@ -751,7 +981,7 @@ impl Iterator for Iter {
                 self.depleted = info != GrB_Info::GrB_SUCCESS
                     || GxB_rowIterator_getRowIndex(self.inner) > self.max_row;
             }
-            Some((row, col))
+            Some(item)
         }
     }
 }

graph/src/graph/graphblas/mod.rs

@@ -54,6 +54,7 @@
 pub mod lagraph_bindings;
 pub mod lagraphx_bindings;
 pub mod matrix;
+pub mod serialization;
 pub mod tensor;
 pub mod vector;
 pub mod versioned_matrix;