use std::io::{self, Read, Write};
use roaring::RoaringBitmap;
use crate::closure::Closure;
use crate::csr::{Csr, CsrError};
use crate::ordinal::to_usize;
const MAGIC: &[u8; 8] = b"FTGRAPH\0";
const VERSION: u32 = 1;
#[derive(Debug, thiserror::Error)]
pub enum PersistError {
#[error("graph I/O failed")]
Io(#[from] io::Error),
#[error("not a graph artifact")]
Magic,
#[error("graph layout version {found}, expected {expected}")]
Version {
found: u32,
expected: u32,
},
#[error(transparent)]
Csr(#[from] CsrError),
#[error("{found} closure sets for {nodes} nodes")]
Count {
found: usize,
nodes: u32,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Hierarchy {
pub is_a: Csr,
pub closure: Closure,
}
impl Hierarchy {
pub fn write_to(&self, out: &mut impl Write) -> Result<(), PersistError> {
out.write_all(MAGIC)?;
out.write_all(&VERSION.to_le_bytes())?;
write_u32s(out, self.is_a.offsets())?;
write_u32s(out, self.is_a.targets())?;
write_bitmaps(out, self.closure.ancestor_sets())?;
write_bitmaps(out, self.closure.descendant_sets())?;
Ok(())
}
pub fn read_from(input: &mut impl Read) -> Result<Self, PersistError> {
let mut magic = [0_u8; 8];
input.read_exact(&mut magic)?;
if &magic != MAGIC {
return Err(PersistError::Magic);
}
let version = read_u32(input)?;
if version != VERSION {
return Err(PersistError::Version {
found: version,
expected: VERSION,
});
}
let offsets = read_u32s(input)?;
let targets = read_u32s(input)?;
let is_a = Csr::from_parts(offsets, targets)?;
let nodes = is_a.nodes();
let ancestors = read_bitmaps(input)?;
let descendants = read_bitmaps(input)?;
for list in [&ancestors, &descendants] {
if list.len() != to_usize(nodes) {
return Err(PersistError::Count {
found: list.len(),
nodes,
});
}
}
Ok(Self {
is_a,
closure: Closure::from_parts(ancestors, descendants),
})
}
}
#[derive(Debug, thiserror::Error)]
#[error("{what} exceeds the u32 the artifact layout stores")]
pub(crate) struct TooLong {
what: &'static str,
#[source]
source: std::num::TryFromIntError,
}
pub(crate) fn u32_len(len: usize, what: &'static str) -> io::Result<u32> {
u32::try_from(len).map_err(|source| io::Error::other(TooLong { what, source }))
}
fn write_u32s(out: &mut impl Write, values: &[u32]) -> io::Result<()> {
let len = u32_len(values.len(), "an array length")?;
out.write_all(&len.to_le_bytes())?;
for value in values {
out.write_all(&value.to_le_bytes())?;
}
Ok(())
}
fn read_u32(input: &mut impl Read) -> io::Result<u32> {
let mut buffer = [0_u8; 4];
input.read_exact(&mut buffer)?;
Ok(u32::from_le_bytes(buffer))
}
fn read_u32s(input: &mut impl Read) -> io::Result<Vec<u32>> {
let len = read_u32(input)?;
let mut values = Vec::with_capacity(to_usize(len));
for _ in 0..len {
values.push(read_u32(input)?);
}
Ok(values)
}
fn write_bitmaps(out: &mut impl Write, sets: &[RoaringBitmap]) -> io::Result<()> {
let len = u32_len(sets.len(), "the set count")?;
out.write_all(&len.to_le_bytes())?;
for set in sets {
let size = u32_len(set.serialized_size(), "a set size")?;
out.write_all(&size.to_le_bytes())?;
set.serialize_into(&mut *out)?;
}
Ok(())
}
fn read_bitmaps(input: &mut impl Read) -> io::Result<Vec<RoaringBitmap>> {
let len = read_u32(input)?;
let mut sets = Vec::with_capacity(to_usize(len));
for _ in 0..len {
let size = read_u32(input)?;
let mut bytes = vec![0_u8; to_usize(size)];
input.read_exact(&mut bytes)?;
sets.push(RoaringBitmap::deserialize_from(bytes.as_slice())?);
}
Ok(sets)
}
#[cfg(test)]
mod tests {
use super::{Hierarchy, PersistError};
use crate::closure::Closure;
use crate::csr::Csr;
use crate::ordinal::Ordinal;
#[test]
fn the_layout_round_trips_and_rejects_foreign_bytes() {
let o = Ordinal::new;
let is_a = Csr::build(4, [(o(1), o(0)), (o(2), o(0)), (o(3), o(1)), (o(3), o(2))])
.expect("builds");
let closure = Closure::compute(&is_a).expect("acyclic");
let hierarchy = Hierarchy { is_a, closure };
let mut bytes = Vec::new();
hierarchy.write_to(&mut bytes).expect("writes");
let back = Hierarchy::read_from(&mut bytes.as_slice()).expect("reads");
assert_eq!(back, hierarchy);
assert!(matches!(
Hierarchy::read_from(&mut b"nope".as_slice()),
Err(PersistError::Io(_))
));
assert!(matches!(
Hierarchy::read_from(&mut b"XXXXXXXX\0\0\0\0".as_slice()),
Err(PersistError::Magic)
));
}
}