1use std::io::{self, Read, Write};
9
10use roaring::RoaringBitmap;
11
12use crate::closure::Closure;
13use crate::csr::{Csr, CsrError};
14use crate::ordinal::to_usize;
15
16const MAGIC: &[u8; 8] = b"FTGRAPH\0";
17const VERSION: u32 = 1;
18
19#[derive(Debug, thiserror::Error)]
21pub enum PersistError {
22 #[error("graph I/O failed")]
24 Io(#[from] io::Error),
25 #[error("not a graph artifact")]
27 Magic,
28 #[error("graph layout version {found}, expected {expected}")]
30 Version {
31 found: u32,
33 expected: u32,
35 },
36 #[error(transparent)]
38 Csr(#[from] CsrError),
39 #[error("{found} closure sets for {nodes} nodes")]
41 Count {
42 found: usize,
44 nodes: u32,
46 },
47}
48
49#[derive(Debug, Clone, PartialEq, Eq)]
51pub struct Hierarchy {
52 pub is_a: Csr,
54 pub closure: Closure,
56}
57
58impl Hierarchy {
59 pub fn write_to(&self, out: &mut impl Write) -> Result<(), PersistError> {
65 out.write_all(MAGIC)?;
66 out.write_all(&VERSION.to_le_bytes())?;
67 write_u32s(out, self.is_a.offsets())?;
68 write_u32s(out, self.is_a.targets())?;
69 write_bitmaps(out, self.closure.ancestor_sets())?;
70 write_bitmaps(out, self.closure.descendant_sets())?;
71 Ok(())
72 }
73
74 pub fn read_from(input: &mut impl Read) -> Result<Self, PersistError> {
80 let mut magic = [0_u8; 8];
81 input.read_exact(&mut magic)?;
82 if &magic != MAGIC {
83 return Err(PersistError::Magic);
84 }
85 let version = read_u32(input)?;
86 if version != VERSION {
87 return Err(PersistError::Version {
88 found: version,
89 expected: VERSION,
90 });
91 }
92 let offsets = read_u32s(input)?;
93 let targets = read_u32s(input)?;
94 let is_a = Csr::from_parts(offsets, targets)?;
95 let nodes = is_a.nodes();
96 let ancestors = read_bitmaps(input)?;
97 let descendants = read_bitmaps(input)?;
98 for list in [&ancestors, &descendants] {
99 if list.len() != to_usize(nodes) {
100 return Err(PersistError::Count {
101 found: list.len(),
102 nodes,
103 });
104 }
105 }
106 Ok(Self {
107 is_a,
108 closure: Closure::from_parts(ancestors, descendants),
109 })
110 }
111}
112
113fn write_u32s(out: &mut impl Write, values: &[u32]) -> io::Result<()> {
114 let len = u32::try_from(values.len()).map_err(|_| io::Error::other("array too long"))?;
115 out.write_all(&len.to_le_bytes())?;
116 for value in values {
117 out.write_all(&value.to_le_bytes())?;
118 }
119 Ok(())
120}
121
122fn read_u32(input: &mut impl Read) -> io::Result<u32> {
123 let mut buffer = [0_u8; 4];
124 input.read_exact(&mut buffer)?;
125 Ok(u32::from_le_bytes(buffer))
126}
127
128fn read_u32s(input: &mut impl Read) -> io::Result<Vec<u32>> {
129 let len = read_u32(input)?;
130 let mut values = Vec::with_capacity(to_usize(len));
131 for _ in 0..len {
132 values.push(read_u32(input)?);
133 }
134 Ok(values)
135}
136
137fn write_bitmaps(out: &mut impl Write, sets: &[RoaringBitmap]) -> io::Result<()> {
138 let len = u32::try_from(sets.len()).map_err(|_| io::Error::other("too many sets"))?;
139 out.write_all(&len.to_le_bytes())?;
140 for set in sets {
141 let size =
142 u32::try_from(set.serialized_size()).map_err(|_| io::Error::other("set too large"))?;
143 out.write_all(&size.to_le_bytes())?;
144 set.serialize_into(&mut *out)?;
145 }
146 Ok(())
147}
148
149fn read_bitmaps(input: &mut impl Read) -> io::Result<Vec<RoaringBitmap>> {
150 let len = read_u32(input)?;
151 let mut sets = Vec::with_capacity(to_usize(len));
152 for _ in 0..len {
153 let size = read_u32(input)?;
154 let mut bytes = vec![0_u8; to_usize(size)];
155 input.read_exact(&mut bytes)?;
156 sets.push(RoaringBitmap::deserialize_from(bytes.as_slice())?);
157 }
158 Ok(sets)
159}
160
161#[cfg(test)]
162mod tests {
163 use super::{Hierarchy, PersistError};
164 use crate::closure::Closure;
165 use crate::csr::Csr;
166 use crate::ordinal::Ordinal;
167
168 #[test]
169 fn the_layout_round_trips_and_rejects_foreign_bytes() {
170 let o = Ordinal::new;
171 let is_a = Csr::build(4, [(o(1), o(0)), (o(2), o(0)), (o(3), o(1)), (o(3), o(2))])
172 .expect("builds");
173 let closure = Closure::compute(&is_a).expect("acyclic");
174 let hierarchy = Hierarchy { is_a, closure };
175 let mut bytes = Vec::new();
176 hierarchy.write_to(&mut bytes).expect("writes");
177 let back = Hierarchy::read_from(&mut bytes.as_slice()).expect("reads");
178 assert_eq!(back, hierarchy);
179 assert!(matches!(
180 Hierarchy::read_from(&mut b"nope".as_slice()),
181 Err(PersistError::Io(_))
182 ));
183 assert!(matches!(
184 Hierarchy::read_from(&mut b"XXXXXXXX\0\0\0\0".as_slice()),
185 Err(PersistError::Magic)
186 ));
187 }
188}