Skip to main content

bixverse_sparse/
serialise.rs

1//! Byte serialisation for the sparse structures and anything built on them.
2//!
3//! A blob is a short header followed by a bincode payload. The header is
4//! 8 magic bytes naming the structure, a `u32` format version and two bytes
5//! holding the element widths of the two data layers, all little-endian.
6//! bincode is not self-describing, so the header is the only thing that stops
7//! a blob of the wrong structure, an older layout or another float width from
8//! decoding into garbage. There is no checksum and no compression.
9
10use bincode::config::{Configuration, standard};
11use serde::{Serialize, de::DeserializeOwned};
12
13use crate::errors::SparseErrors;
14use crate::sparse::CompressedSparseData2;
15use crate::traits::BixverseNumeric;
16
17////////////
18// Header //
19////////////
20
21/// Version of the byte layout. Bump whenever a change to a serialised struct
22/// (reordering, removing or retyping a field) stops older blobs decoding.
23/// Nothing checks this automatically.
24pub const SERIALISE_FORMAT_VERSION: u32 = 1;
25
26/// Magic bytes of a serialised [`CompressedSparseData2`].
27const CSR_MAGIC: &[u8; 8] = b"BXSPCSD2";
28
29/// Header length: magic, version, two element widths.
30const HEADER_LEN: usize = 8 + 4 + 2;
31
32/// The one bincode configuration used in both directions. Little-endian,
33/// variable-length integers, no decode limit. Varints keep the `usize` shape
34/// fields portable and shrink the `u32` indices.
35const CONFIG: Configuration = standard();
36
37/// Encode a value behind a header.
38///
39/// ### Params
40///
41/// * `magic` - Magic bytes naming the structure
42/// * `widths` - Element widths of the value's data layers, in bytes
43/// * `value` - The value to encode
44///
45/// ### Returns
46///
47/// The header followed by the bincode payload.
48pub fn encode_with_header<S>(
49    magic: &[u8; 8],
50    widths: [u8; 2],
51    value: &S,
52) -> Result<Vec<u8>, SparseErrors>
53where
54    S: Serialize,
55{
56    let payload = bincode::serde::encode_to_vec(value, CONFIG)
57        .map_err(|e| SparseErrors::SerialisationFailed(e.to_string()))?;
58
59    let mut out = Vec::with_capacity(HEADER_LEN + payload.len());
60    out.extend_from_slice(magic);
61    out.extend_from_slice(&SERIALISE_FORMAT_VERSION.to_le_bytes());
62    out.extend_from_slice(&widths);
63    out.extend_from_slice(&payload);
64    Ok(out)
65}
66
67/// Check the header and decode the payload behind it.
68///
69/// Checks the magic, the version and the widths, in that order, so the error
70/// names the first thing that is actually wrong. Trailing bytes after the
71/// payload are an error too: they mean the blob is not what the header says.
72///
73/// ### Params
74///
75/// * `magic` - Magic bytes the caller expects
76/// * `widths` - Element widths the caller expects
77/// * `name` - Structure name, for the error message
78/// * `bytes` - The blob
79///
80/// ### Returns
81///
82/// The decoded value, or the first mismatch.
83pub fn decode_with_header<D>(
84    magic: &[u8; 8],
85    widths: [u8; 2],
86    name: &'static str,
87    bytes: &[u8],
88) -> Result<D, SparseErrors>
89where
90    D: DeserializeOwned,
91{
92    if bytes.len() < 8 || &bytes[..8] != magic {
93        return Err(SparseErrors::NotSerialised(name));
94    }
95    if bytes.len() < HEADER_LEN {
96        return Err(SparseErrors::SerialisationFailed(format!(
97            "{name}: the header is truncated"
98        )));
99    }
100
101    let version = u32::from_le_bytes(bytes[8..12].try_into().expect("slice of four bytes"));
102    if version != SERIALISE_FORMAT_VERSION {
103        return Err(SparseErrors::SerialisedVersionMismatch {
104            found: version,
105            supported: SERIALISE_FORMAT_VERSION,
106        });
107    }
108
109    let found = [bytes[12], bytes[13]];
110    if found != widths {
111        return Err(SparseErrors::SerialisedWidthMismatch {
112            expected: widths,
113            found,
114        });
115    }
116
117    let (value, read) = bincode::serde::decode_from_slice(&bytes[HEADER_LEN..], CONFIG)
118        .map_err(|e| SparseErrors::SerialisationFailed(format!("{name}: {e}")))?;
119    if HEADER_LEN + read != bytes.len() {
120        return Err(SparseErrors::SerialisationFailed(format!(
121            "{name}: {} trailing bytes after the payload",
122            bytes.len() - HEADER_LEN - read
123        )));
124    }
125
126    Ok(value)
127}
128
129/// Element width of a type, as stored in the header.
130///
131/// ### Returns
132///
133/// `size_of::<T>()` as a byte.
134#[inline]
135pub fn elem_width<T>() -> u8 {
136    std::mem::size_of::<T>() as u8
137}
138
139///////////////////////////
140// CompressedSparseData2 //
141///////////////////////////
142
143impl<T, U> CompressedSparseData2<T, U>
144where
145    T: BixverseNumeric + Serialize + DeserializeOwned,
146    U: BixverseNumeric + Serialize + DeserializeOwned,
147{
148    /// Serialise the matrix to bytes.
149    ///
150    /// ### Returns
151    ///
152    /// A blob that [`Self::from_bytes`] reads back.
153    pub fn to_bytes(&self) -> Result<Vec<u8>, SparseErrors> {
154        encode_with_header(CSR_MAGIC, [elem_width::<T>(), elem_width::<U>()], self)
155    }
156
157    /// Read a matrix written by [`Self::to_bytes`].
158    ///
159    /// The decoded structure is validated, since the bytes may come from disk.
160    ///
161    /// ### Params
162    ///
163    /// * `bytes` - The blob
164    ///
165    /// ### Returns
166    ///
167    /// The matrix, or a header, decode or validation error.
168    pub fn from_bytes(bytes: &[u8]) -> Result<Self, SparseErrors> {
169        let csr: Self = decode_with_header(
170            CSR_MAGIC,
171            [elem_width::<T>(), elem_width::<U>()],
172            "CompressedSparseData2",
173            bytes,
174        )?;
175        csr.validate()?;
176        Ok(csr)
177    }
178}
179
180///////////
181// Tests //
182///////////
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187    use crate::sparse::CompressedSparseFormat;
188
189    /// A 3 x 4 CSR with one empty row and a second layer.
190    fn small_csr() -> CompressedSparseData2<f32, f64> {
191        CompressedSparseData2 {
192            data: vec![1.0, 2.0, 3.0, 4.0],
193            indices: vec![0, 3, 1, 2],
194            indptr: vec![0, 2, 2, 4],
195            cs_type: CompressedSparseFormat::Csr,
196            data_2: Some(vec![10.0, 20.0, 30.0, 40.0]),
197            shape: (3, 4),
198        }
199    }
200
201    #[test]
202    fn test_csr_bytes_round_trip() {
203        let csr = small_csr();
204        let back = CompressedSparseData2::<f32, f64>::from_bytes(&csr.to_bytes().unwrap()).unwrap();
205        assert_eq!(back.data, csr.data);
206        assert_eq!(back.indices, csr.indices);
207        assert_eq!(back.indptr, csr.indptr);
208        assert_eq!(back.data_2, csr.data_2);
209        assert_eq!(back.shape, csr.shape);
210        assert!(back.cs_type.is_csr());
211    }
212
213    #[test]
214    fn test_csr_bytes_bad_magic() {
215        let mut bytes = small_csr().to_bytes().unwrap();
216        bytes[0] = b'X';
217        let res = CompressedSparseData2::<f32, f64>::from_bytes(&bytes);
218        assert!(matches!(res, Err(SparseErrors::NotSerialised(_))));
219        let res = CompressedSparseData2::<f32, f64>::from_bytes(&[1, 2, 3]);
220        assert!(matches!(res, Err(SparseErrors::NotSerialised(_))));
221    }
222
223    #[test]
224    fn test_csr_bytes_version_mismatch() {
225        let mut bytes = small_csr().to_bytes().unwrap();
226        bytes[8..12].copy_from_slice(&(SERIALISE_FORMAT_VERSION + 1).to_le_bytes());
227        let res = CompressedSparseData2::<f32, f64>::from_bytes(&bytes);
228        assert!(matches!(
229            res,
230            Err(SparseErrors::SerialisedVersionMismatch { found, supported })
231                if found == SERIALISE_FORMAT_VERSION + 1 && supported == SERIALISE_FORMAT_VERSION
232        ));
233    }
234
235    #[test]
236    fn test_csr_bytes_width_mismatch() {
237        let bytes = small_csr().to_bytes().unwrap();
238        let res = CompressedSparseData2::<f64, f64>::from_bytes(&bytes);
239        assert!(matches!(
240            res,
241            Err(SparseErrors::SerialisedWidthMismatch { .. })
242        ));
243    }
244
245    #[test]
246    fn test_csr_bytes_truncated_and_trailing() {
247        let bytes = small_csr().to_bytes().unwrap();
248        let res = CompressedSparseData2::<f32, f64>::from_bytes(&bytes[..bytes.len() - 3]);
249        assert!(matches!(res, Err(SparseErrors::SerialisationFailed(_))));
250
251        let mut long = bytes.clone();
252        long.push(0);
253        let res = CompressedSparseData2::<f32, f64>::from_bytes(&long);
254        assert!(matches!(res, Err(SparseErrors::SerialisationFailed(_))));
255    }
256
257    #[test]
258    fn test_csr_bytes_rejects_invalid_structure() {
259        let mut csr = small_csr();
260        csr.indices[1] = 9;
261        let res = CompressedSparseData2::<f32, f64>::from_bytes(&csr.to_bytes().unwrap());
262        assert!(matches!(
263            res,
264            Err(SparseErrors::SliceIndexOutOfBounds { .. })
265        ));
266    }
267}