1use std::io::{self, Read, Write};
12
13use roaring::RoaringBitmap;
14
15use crate::closure::Closure;
16use crate::csr::{Csr, CsrError};
17use crate::ordinal::to_usize;
18
19const MAGIC: &[u8; 8] = b"FTGRAPH\0";
20const VERSION: u32 = 1;
21
22#[derive(Debug, thiserror::Error)]
24pub enum PersistError {
25 #[error("graph I/O failed")]
27 Io(#[from] io::Error),
28 #[error("not a graph artifact")]
30 Magic,
31 #[error("graph layout version {found}, expected {expected}")]
33 Version {
34 found: u32,
36 expected: u32,
38 },
39 #[error(transparent)]
41 Csr(#[from] CsrError),
42 #[error("{found} closure sets for {nodes} nodes")]
44 Count {
45 found: usize,
47 nodes: u32,
49 },
50}
51
52#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct Hierarchy {
55 pub is_a: Csr,
57 pub closure: Closure,
59}
60
61impl Hierarchy {
62 #[must_use]
65 pub fn size_in_bytes(&self) -> usize {
66 self.is_a
67 .size_in_bytes()
68 .saturating_add(self.closure.size_in_bytes())
69 }
70
71 pub fn write_to(&self, out: &mut impl Write) -> Result<(), PersistError> {
77 out.write_all(MAGIC)?;
78 out.write_all(&VERSION.to_le_bytes())?;
79 write_u32s(out, self.is_a.offsets())?;
80 write_u32s(out, self.is_a.targets())?;
81 write_bitmaps(out, self.closure.ancestor_sets())?;
82 write_bitmaps(out, self.closure.descendant_sets())?;
83 Ok(())
84 }
85
86 pub fn read_from(input: &mut impl Read) -> Result<Self, PersistError> {
92 let mut magic = [0_u8; 8];
93 input.read_exact(&mut magic)?;
94 if &magic != MAGIC {
95 return Err(PersistError::Magic);
96 }
97 let version = read_u32(input)?;
98 if version != VERSION {
99 return Err(PersistError::Version {
100 found: version,
101 expected: VERSION,
102 });
103 }
104 let offsets = read_u32s(input)?;
105 let targets = read_u32s(input)?;
106 let is_a = Csr::from_parts(offsets, targets)?;
107 let nodes = is_a.nodes();
108 let ancestors = read_bitmaps(input)?;
109 let descendants = read_bitmaps(input)?;
110 for list in [&ancestors, &descendants] {
111 if list.len() != to_usize(nodes) {
112 return Err(PersistError::Count {
113 found: list.len(),
114 nodes,
115 });
116 }
117 }
118 Ok(Self {
119 is_a,
120 closure: Closure::from_parts(ancestors, descendants),
121 })
122 }
123}
124
125#[derive(Debug, thiserror::Error)]
127#[error("{what} exceeds the u32 the artifact layout stores")]
128pub(crate) struct TooLong {
129 what: &'static str,
130 #[source]
131 source: std::num::TryFromIntError,
132}
133
134pub(crate) fn u32_len(len: usize, what: &'static str) -> io::Result<u32> {
136 u32::try_from(len).map_err(|source| io::Error::other(TooLong { what, source }))
137}
138
139fn write_u32s(out: &mut impl Write, values: &[u32]) -> io::Result<()> {
140 let len = u32_len(values.len(), "an array length")?;
141 out.write_all(&len.to_le_bytes())?;
142 for value in values {
143 out.write_all(&value.to_le_bytes())?;
144 }
145 Ok(())
146}
147
148fn read_u32(input: &mut impl Read) -> io::Result<u32> {
149 let mut buffer = [0_u8; 4];
150 input.read_exact(&mut buffer)?;
151 Ok(u32::from_le_bytes(buffer))
152}
153
154fn read_u32s(input: &mut impl Read) -> io::Result<Vec<u32>> {
155 let len = read_u32(input)?;
156 let mut values = Vec::with_capacity(to_usize(len));
157 for _ in 0..len {
158 values.push(read_u32(input)?);
159 }
160 Ok(values)
161}
162
163fn write_bitmaps(out: &mut impl Write, sets: &[RoaringBitmap]) -> io::Result<()> {
164 let len = u32_len(sets.len(), "the set count")?;
165 out.write_all(&len.to_le_bytes())?;
166 for set in sets {
167 let size = u32_len(set.serialized_size(), "a set size")?;
168 out.write_all(&size.to_le_bytes())?;
169 set.serialize_into(&mut *out)?;
170 }
171 Ok(())
172}
173
174fn read_bitmaps(input: &mut impl Read) -> io::Result<Vec<RoaringBitmap>> {
175 let len = read_u32(input)?;
176 let mut sets = Vec::with_capacity(to_usize(len));
177 for _ in 0..len {
178 let size = read_u32(input)?;
179 let mut bytes = vec![0_u8; to_usize(size)];
180 input.read_exact(&mut bytes)?;
181 sets.push(RoaringBitmap::deserialize_from(bytes.as_slice())?);
182 }
183 Ok(sets)
184}
185
186#[cfg(test)]
187mod tests {
188 use super::{Hierarchy, PersistError};
189 use crate::closure::Closure;
190 use crate::csr::Csr;
191 use crate::ordinal::Ordinal;
192
193 #[test]
194 fn the_layout_round_trips_and_rejects_foreign_bytes() {
195 let o = Ordinal::new;
196 let is_a = Csr::build(4, [(o(1), o(0)), (o(2), o(0)), (o(3), o(1)), (o(3), o(2))])
197 .expect("builds");
198 let closure = Closure::compute(&is_a).expect("acyclic");
199 let hierarchy = Hierarchy { is_a, closure };
200 let mut bytes = Vec::new();
201 hierarchy.write_to(&mut bytes).expect("writes");
202 let back = Hierarchy::read_from(&mut bytes.as_slice()).expect("reads");
203 assert_eq!(back, hierarchy);
204 assert!(matches!(
205 Hierarchy::read_from(&mut b"nope".as_slice()),
206 Err(PersistError::Io(_))
207 ));
208 assert!(matches!(
209 Hierarchy::read_from(&mut b"XXXXXXXX\0\0\0\0".as_slice()),
210 Err(PersistError::Magic)
211 ));
212 }
213}