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
113#[derive(Debug, thiserror::Error)]
115#[error("{what} exceeds the u32 the artifact layout stores")]
116pub(crate) struct TooLong {
117 what: &'static str,
118 #[source]
119 source: std::num::TryFromIntError,
120}
121
122pub(crate) fn u32_len(len: usize, what: &'static str) -> io::Result<u32> {
124 u32::try_from(len).map_err(|source| io::Error::other(TooLong { what, source }))
125}
126
127fn write_u32s(out: &mut impl Write, values: &[u32]) -> io::Result<()> {
128 let len = u32_len(values.len(), "an array length")?;
129 out.write_all(&len.to_le_bytes())?;
130 for value in values {
131 out.write_all(&value.to_le_bytes())?;
132 }
133 Ok(())
134}
135
136fn read_u32(input: &mut impl Read) -> io::Result<u32> {
137 let mut buffer = [0_u8; 4];
138 input.read_exact(&mut buffer)?;
139 Ok(u32::from_le_bytes(buffer))
140}
141
142fn read_u32s(input: &mut impl Read) -> io::Result<Vec<u32>> {
143 let len = read_u32(input)?;
144 let mut values = Vec::with_capacity(to_usize(len));
145 for _ in 0..len {
146 values.push(read_u32(input)?);
147 }
148 Ok(values)
149}
150
151fn write_bitmaps(out: &mut impl Write, sets: &[RoaringBitmap]) -> io::Result<()> {
152 let len = u32_len(sets.len(), "the set count")?;
153 out.write_all(&len.to_le_bytes())?;
154 for set in sets {
155 let size = u32_len(set.serialized_size(), "a set size")?;
156 out.write_all(&size.to_le_bytes())?;
157 set.serialize_into(&mut *out)?;
158 }
159 Ok(())
160}
161
162fn read_bitmaps(input: &mut impl Read) -> io::Result<Vec<RoaringBitmap>> {
163 let len = read_u32(input)?;
164 let mut sets = Vec::with_capacity(to_usize(len));
165 for _ in 0..len {
166 let size = read_u32(input)?;
167 let mut bytes = vec![0_u8; to_usize(size)];
168 input.read_exact(&mut bytes)?;
169 sets.push(RoaringBitmap::deserialize_from(bytes.as_slice())?);
170 }
171 Ok(sets)
172}
173
174#[cfg(test)]
175mod tests {
176 use super::{Hierarchy, PersistError};
177 use crate::closure::Closure;
178 use crate::csr::Csr;
179 use crate::ordinal::Ordinal;
180
181 #[test]
182 fn the_layout_round_trips_and_rejects_foreign_bytes() {
183 let o = Ordinal::new;
184 let is_a = Csr::build(4, [(o(1), o(0)), (o(2), o(0)), (o(3), o(1)), (o(3), o(2))])
185 .expect("builds");
186 let closure = Closure::compute(&is_a).expect("acyclic");
187 let hierarchy = Hierarchy { is_a, closure };
188 let mut bytes = Vec::new();
189 hierarchy.write_to(&mut bytes).expect("writes");
190 let back = Hierarchy::read_from(&mut bytes.as_slice()).expect("reads");
191 assert_eq!(back, hierarchy);
192 assert!(matches!(
193 Hierarchy::read_from(&mut b"nope".as_slice()),
194 Err(PersistError::Io(_))
195 ));
196 assert!(matches!(
197 Hierarchy::read_from(&mut b"XXXXXXXX\0\0\0\0".as_slice()),
198 Err(PersistError::Magic)
199 ));
200 }
201}