1use rudb_common::{Error, Result};
24use rudb_encoding::bitpack;
25
26use crate::bits::BitVector;
27use crate::rid::{NO_PARENT, Rid};
28use crate::rids::Rids;
29use crate::tail::Tail;
30
31const LAYOUT: u8 = 1;
33
34pub const HEADER_BYTES: usize = 32;
38
39#[derive(Debug, Clone)]
41pub struct Adjacency {
42 children: u64,
44 parents: u64,
45 edges: u64,
47 starts: BitVector,
49 rows: Tail,
51 width: usize,
52}
53
54impl Adjacency {
55 pub fn build(parents_of: &[Rid], parents: u64) -> Result<Self> {
65 let parent_rows = usize::try_from(parents)
66 .map_err(|_| malformed("a parent table larger than fits in memory"))?;
67 let mut starts = vec![0_usize; parent_rows + 1];
68 for &parent in parents_of {
69 if parent == NO_PARENT {
70 continue;
71 }
72 if parent >= parents {
73 return Err(malformed(format!(
74 "a child points at parent {parent} of a table with {parents} rows"
75 )));
76 }
77 starts[parent as usize + 1] += 1;
78 }
79 for at in 1..starts.len() {
80 starts[at] += starts[at - 1];
81 }
82 let edges = starts[parent_rows];
83 let mut placed = starts.clone();
84 let mut grouped = vec![0_u64; edges];
85 for (child, &parent) in parents_of.iter().enumerate() {
86 if parent == NO_PARENT {
87 continue;
88 }
89 let slot = &mut placed[parent as usize];
90 grouped[*slot] = count(child);
91 *slot += 1;
92 }
93 let len = edges + parent_rows;
94 let mut words = vec![0_u64; len.div_ceil(64)];
95 for parent in 0..parent_rows {
96 for at in starts[parent] + parent..starts[parent + 1] + parent {
98 words[at / 64] |= 1 << (at % 64);
99 }
100 }
101 let children = count(parents_of.len());
102 let width = width_for(children);
103 let mut rows = Vec::with_capacity(bitpack::tail_len(edges, width));
104 bitpack::pack_linear(&grouped, width, &mut rows)?;
105 Ok(Self {
106 children,
107 parents,
108 edges: count(edges),
109 starts: BitVector::new(words, len)?,
110 rows: rows.into(),
111 width,
112 })
113 }
114
115 #[must_use]
117 pub fn children(&self) -> u64 {
118 self.children
119 }
120
121 #[must_use]
123 pub fn parents(&self) -> u64 {
124 self.parents
125 }
126
127 #[must_use]
129 pub fn edges(&self) -> u64 {
130 self.edges
131 }
132
133 #[must_use]
135 pub fn bytes(&self) -> usize {
136 self.starts.bytes() + self.rows.len()
137 }
138
139 fn list(&self, parent: Rid) -> Option<std::ops::Range<usize>> {
141 if parent >= self.parents {
142 return None;
143 }
144 let cum = |nth: Rid| -> Option<usize> {
145 self.starts.select0(nth).map(|at| at - usize::try_from(nth).unwrap_or(usize::MAX))
146 };
147 let from = if parent == 0 { 0 } else { cum(parent - 1)? };
148 Some(from..cum(parent)?)
149 }
150
151 pub fn children_of(&self, parent: Rid, out: &mut Vec<Rid>) -> Result<()> {
157 let list = self.list(parent).ok_or_else(|| {
158 malformed(format!("parent {parent} of {} is past the end", self.parents))
159 })?;
160 for at in list {
161 out.push(bitpack::tail_at(&self.rows, self.width, at)?);
162 }
163 Ok(())
164 }
165
166 #[must_use]
170 pub fn reached(&self, held: &Rids) -> u64 {
171 held.iter().filter_map(|parent| self.list(parent)).map(|list| count(list.len())).sum()
172 }
173
174 pub fn push(&self, held: &Rids) -> Result<Rids> {
185 if held.rows() != self.parents {
186 return Err(Error::internal(format!(
187 "a set over {} rows pushed through an adjacency over {} parents",
188 held.rows(),
189 self.parents
190 )));
191 }
192 let mut words = vec![0_u64; usize::try_from(self.children.div_ceil(64)).unwrap_or(0)];
193 for parent in held.iter() {
194 let list = self.list(parent).ok_or_else(|| {
195 malformed(format!("parent {parent} of {} is past the end", self.parents))
196 })?;
197 for at in list {
198 let child = bitpack::tail_at(&self.rows, self.width, at)?;
199 let word = words
200 .get_mut(usize::try_from(child / 64).unwrap_or(usize::MAX))
201 .ok_or_else(|| malformed(format!("child {child} past the end")))?;
202 *word |= 1 << (child % 64);
203 }
204 }
205 Rids::from_words(self.children, words)
206 }
207
208 pub fn write(&self, out: &mut Vec<u8>) -> Result<()> {
214 out.extend_from_slice(&self.children.to_le_bytes());
215 out.extend_from_slice(&self.parents.to_le_bytes());
216 out.extend_from_slice(&self.edges.to_le_bytes());
217 out.push(u8::try_from(self.width).map_err(|_| malformed("a width past a byte"))?);
218 out.push(LAYOUT);
219 out.extend_from_slice(&[0; 6]);
220 self.starts.write(out);
221 out.extend_from_slice(&self.rows);
222 Ok(())
223 }
224
225 pub fn read(bytes: &[u8]) -> Result<Self> {
232 Self::read_from(bytes.to_vec(), 0)
233 }
234
235 pub fn read_from(payload: Vec<u8>, at: usize) -> Result<Self> {
242 let bytes =
243 payload.get(at..).ok_or_else(|| malformed("a payload shorter than its header"))?;
244 if bytes.len() < HEADER_BYTES {
245 return Err(malformed("a payload shorter than its header"));
246 }
247 let children = number(&bytes[0..8])?;
248 let parents = number(&bytes[8..16])?;
249 let edges = number(&bytes[16..24])?;
250 let width = bytes[24] as usize;
251 if bytes[25] != LAYOUT {
252 return Err(malformed(format!("layout {} is not one this build knows", bytes[25])));
253 }
254 if width != width_for(children) || edges > children {
255 return Err(malformed("a header whose numbers do not agree"));
256 }
257 let len = usize::try_from(edges + parents)
258 .map_err(|_| malformed("a list longer than fits in memory"))?;
259 let edge_count =
260 usize::try_from(edges).map_err(|_| malformed("more edges than fit in memory"))?;
261 let rest = &bytes[HEADER_BYTES..];
262 let split = BitVector::bytes_for(len);
263 if rest.len() != split + bitpack::tail_len(edge_count, width) {
264 return Err(malformed("a body that is not the size its header implies"));
265 }
266 let starts = BitVector::read(&rest[..split], len)?;
267 if starts.ones() != edges {
268 return Err(malformed("lists that do not hold the edges the header counts"));
269 }
270 let rows = Tail::of(payload, at + HEADER_BYTES + split)
271 .ok_or_else(|| malformed("a body that is not the size its header implies"))?;
272 Ok(Self { children, parents, edges, starts, rows, width })
273 }
274}
275
276fn width_for(children: u64) -> usize {
278 (u64::BITS - children.saturating_sub(1).leading_zeros()).max(1) as usize
279}
280
281fn count(rows: usize) -> u64 {
282 u64::try_from(rows).unwrap_or(u64::MAX)
283}
284
285fn number(bytes: &[u8]) -> Result<u64> {
286 Ok(u64::from_le_bytes(bytes.try_into().map_err(|_| malformed("a header is torn"))?))
287}
288
289fn malformed(message: impl Into<String>) -> Error {
290 Error::invalid_input(format!("invalid rudb backward adjacency: {}", message.into()))
291}
292
293#[cfg(test)]
294mod tests {
295 use super::*;
296 use crate::link::Link;
297
298 fn scattered() -> (Vec<Rid>, u64) {
301 let parents = 50;
302 let parents_of = (0..3_000_u64)
303 .map(|child| if child % 97 == 0 { NO_PARENT } else { (child * 7919 + 13) % 49 })
304 .collect();
305 (parents_of, parents)
306 }
307
308 #[test]
309 fn each_parent_lists_exactly_the_children_that_point_at_it_in_row_order() {
310 let (parents_of, parents) = scattered();
311 let adjacency = Adjacency::build(&parents_of, parents).expect("build");
312 for parent in 0..parents {
313 let mut listed = Vec::new();
314 adjacency.children_of(parent, &mut listed).expect("list");
315 let expected: Vec<Rid> = (0..count(parents_of.len()))
316 .filter(|&child| parents_of[child as usize] == parent)
317 .collect();
318 assert_eq!(listed, expected, "parent {parent}");
319 }
320 assert!(adjacency.children_of(parents, &mut Vec::new()).is_err(), "past the end");
321 }
322
323 #[test]
324 fn a_push_is_the_set_a_forward_push_through_the_link_gives() {
325 let (parents_of, parents) = scattered();
326 let adjacency = Adjacency::build(&parents_of, parents).expect("build");
327 let link = Link::build(&parents_of, parents).expect("link");
328 let held = Rids::from_sorted(parents, vec![0, 3, 17, 48, 49]).expect("held");
329 let pushed = adjacency.push(&held).expect("push");
330 let forward = held.forward(&link).expect("forward").rids;
331 assert_eq!(pushed.iter().collect::<Vec<_>>(), forward.iter().collect::<Vec<_>>());
332 assert_eq!(adjacency.reached(&held), pushed.len(), "counted without reading a row");
333 let wrong = Rids::from_sorted(parents + 1, vec![0]).expect("wrong");
334 assert!(adjacency.push(&wrong).is_err(), "a set over another table");
335 }
336
337 #[test]
338 fn it_reads_back_what_it_wrote_and_refuses_a_torn_body() {
339 let (parents_of, parents) = scattered();
340 let adjacency = Adjacency::build(&parents_of, parents).expect("build");
341 let mut bytes = Vec::new();
342 adjacency.write(&mut bytes).expect("write");
343 assert_eq!(bytes.len(), HEADER_BYTES + adjacency.bytes());
344 let read = Adjacency::read(&bytes).expect("read");
345 assert_eq!(read.edges(), adjacency.edges());
346 let held = Rids::from_sorted(parents, vec![5, 6, 7]).expect("held");
347 assert_eq!(read.push(&held).expect("push"), adjacency.push(&held).expect("push"));
348 assert!(Adjacency::read(&bytes[..bytes.len() - 1]).is_err(), "a short body");
349 bytes[25] = 9;
350 assert!(Adjacency::read(&bytes).is_err(), "a layout from elsewhere");
351 }
352}