use rudb_common::{Error, Result};
use rudb_encoding::bitpack;
use crate::bits::BitVector;
use crate::rid::{NO_PARENT, PART_ROWS, Rid};
const LAYOUT: u8 = 1;
pub const HEADER_BYTES: usize = 32;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Counts {
pub children: u64,
pub parents: u64,
pub linked: u64,
pub form: Form,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Form {
Packed,
Monotone,
}
impl Form {
#[must_use]
pub fn tag(self) -> u8 {
match self {
Self::Packed => 0,
Self::Monotone => 1,
}
}
#[must_use]
pub fn label(self) -> &'static str {
match self {
Self::Packed => "packed",
Self::Monotone => "monotone",
}
}
pub fn from_tag(tag: u8) -> Result<Self> {
match tag {
0 => Ok(Self::Packed),
1 => Ok(Self::Monotone),
_ => Err(malformed(format!("forward link form {tag} is not one this build knows"))),
}
}
}
pub type Bounds = Option<(Rid, Rid)>;
#[derive(Debug, Clone)]
enum Body {
Packed {
bytes: Vec<u8>,
width: usize,
heads: Vec<Bounds>,
},
Monotone {
vector: BitVector,
},
}
#[derive(Debug, Clone)]
pub struct Link {
children: u64,
parents: u64,
linked: u64,
body: Body,
}
impl Link {
pub fn build(parents_of: &[Rid], parents: u64) -> Result<Self> {
let children = count(parents_of.len());
let mut linked = 0_u64;
let mut monotone = true;
let mut previous = 0_u64;
for parent in parents_of {
if *parent == NO_PARENT {
monotone = false;
continue;
}
if *parent >= parents {
return Err(malformed(format!(
"a forward link points at parent {parent} of a table with {parents} rows"
)));
}
if *parent < previous {
monotone = false;
}
previous = *parent;
linked += 1;
}
let body = if monotone && parents > 0 {
Body::Monotone { vector: runs(parents_of, parents)? }
} else {
packed(parents_of, parents)?
};
Ok(Self { children, parents, linked, body })
}
#[must_use]
pub fn form(&self) -> Form {
match self.body {
Body::Packed { .. } => Form::Packed,
Body::Monotone { .. } => Form::Monotone,
}
}
#[must_use]
pub fn children(&self) -> u64 {
self.children
}
#[must_use]
pub fn parents(&self) -> u64 {
self.parents
}
#[must_use]
pub fn linked(&self) -> u64 {
self.linked
}
pub(crate) fn runs(&self) -> Option<&BitVector> {
match &self.body {
Body::Monotone { vector } => Some(vector),
Body::Packed { .. } => None,
}
}
#[must_use]
pub fn bytes(&self) -> usize {
match &self.body {
Body::Packed { bytes, heads, .. } => bytes.len() + heads.len() * 16,
Body::Monotone { vector } => vector.bytes(),
}
}
#[must_use]
pub fn forward(&self, child: Rid) -> Option<Rid> {
if child >= self.children {
return None;
}
match &self.body {
Body::Packed { bytes, width, .. } => {
let value = bitpack::tail_at(bytes, *width, usize::try_from(child).ok()?).ok()?;
(value != reserved(*width)).then_some(value)
}
Body::Monotone { vector } => {
let at = vector.select1(child)?;
Some(vector.rank0(at))
}
}
}
pub fn forward_run(&self, first: Rid, out: &mut [Rid]) -> Result<()> {
let end = first.checked_add(count(out.len()));
if end.is_none_or(|end| end > self.children) {
return Err(Error::internal(format!(
"a run of {} children from {first} goes past the {} the link has",
out.len(),
self.children
)));
}
if out.is_empty() {
return Ok(());
}
match &self.body {
Body::Packed { bytes, width, .. } => {
let absent = reserved(*width);
let start = usize::try_from(first)
.map_err(|_| malformed("a child past what fits in memory"))?;
for (at, slot) in out.iter_mut().enumerate() {
let value = bitpack::tail_at(bytes, *width, start + at)?;
*slot = if value == absent { NO_PARENT } else { value };
}
}
Body::Monotone { vector } => {
let Some(mut at) = vector.select1(first) else {
return Err(malformed("a monotone link has fewer ones than children"));
};
let mut parent = vector.rank0(at);
let words = vector.words();
for slot in out.iter_mut() {
loop {
let word = words.get(at / 64).copied().unwrap_or(0) >> (at % 64);
if word == 0 {
let skipped = 64 - at % 64;
parent += count(skipped);
at += skipped;
if at >= vector.len() {
return Err(malformed("a monotone link ran out of ones"));
}
continue;
}
let zeros = word.trailing_zeros() as usize;
parent += count(zeros);
at += zeros;
break;
}
*slot = parent;
at += 1;
}
}
}
Ok(())
}
pub fn forward_each(&self, children: &[Rid], out: &mut Vec<Rid>) {
out.clear();
out.reserve(children.len());
let Body::Monotone { vector } = &self.body else {
out.extend(children.iter().map(|&child| self.forward(child).unwrap_or(NO_PARENT)));
return;
};
let words = vector.words();
let mut last: Option<(Rid, usize, Rid)> = None;
for &child in children {
if child >= self.children {
out.push(NO_PARENT);
continue;
}
let near = last.filter(|&(from, ..)| child >= from && child - from <= Self::WALK);
let (at, parent) = match near {
Some((from, at, parent)) => {
walk_ones(words, at, parent, child - from).unwrap_or((usize::MAX, NO_PARENT))
}
None => match vector.select1(child) {
Some(at) => (at, vector.rank0(at)),
None => (usize::MAX, NO_PARENT),
},
};
if parent == NO_PARENT {
last = None;
} else {
last = Some((child, at, parent));
}
out.push(parent);
}
}
const WALK: Rid = 1024;
#[must_use]
pub fn backward(&self, parent: Rid) -> Option<std::ops::Range<Rid>> {
let Body::Monotone { vector } = &self.body else { return None };
if parent >= self.parents {
return None;
}
let cum = |nth: Rid| -> Option<u64> { vector.select0(nth).map(|at| count(at) - nth) };
let from = if parent == 0 { 0 } else { cum(parent - 1)? };
Some(from..cum(parent)?)
}
#[must_use]
pub fn part_bounds(&self, part: usize) -> Option<Bounds> {
let first = count(part * PART_ROWS);
if first >= self.children {
return None;
}
match &self.body {
Body::Packed { heads, .. } => heads.get(part).copied(),
Body::Monotone { .. } => {
let last = (first + count(PART_ROWS) - 1).min(self.children - 1);
match (self.forward(first), self.forward(last)) {
(Some(low), Some(high)) => Some(Some((low, high))),
_ => Some(None),
}
}
}
}
pub fn write(&self, out: &mut Vec<u8>) -> Result<()> {
let start = out.len();
out.extend_from_slice(&self.children.to_le_bytes());
out.extend_from_slice(&self.parents.to_le_bytes());
out.extend_from_slice(&self.linked.to_le_bytes());
out.push(self.form().tag());
out.push(match &self.body {
Body::Packed { width, .. } => u8::try_from(*width)
.map_err(|_| malformed("a forward link wider than a byte can name"))?,
Body::Monotone { .. } => 0,
});
out.push(LAYOUT);
out.extend_from_slice(&[0; 5]);
debug_assert_eq!(
out.len() - start,
HEADER_BYTES,
"the forward link header is thirty two bytes"
);
match &self.body {
Body::Packed { bytes, heads, .. } => {
for head in heads {
let (low, high) = head.unwrap_or((NO_PARENT, NO_PARENT));
out.extend_from_slice(&low.to_le_bytes());
out.extend_from_slice(&high.to_le_bytes());
}
out.extend_from_slice(bytes);
}
Body::Monotone { vector } => vector.write(out),
}
Ok(())
}
pub fn counts(bytes: &[u8]) -> Result<Counts> {
if bytes.len() < HEADER_BYTES {
return Err(malformed("a forward link payload is shorter than its header"));
}
let form = Form::from_tag(bytes[24])?;
if bytes[26] != LAYOUT {
return Err(malformed(format!(
"forward link layout {} is not one this build knows",
bytes[26]
)));
}
Ok(Counts {
children: number(&bytes[0..8])?,
parents: number(&bytes[8..16])?,
linked: number(&bytes[16..24])?,
form,
})
}
pub fn read(bytes: &[u8]) -> Result<Self> {
let Counts { children, parents, linked, form } = Self::counts(bytes)?;
let width = bytes[25] as usize;
let rest = &bytes[HEADER_BYTES..];
let body = match form {
Form::Packed => {
if width != width_for(parents) {
return Err(malformed(
"a forward link's width is not the one its parents imply",
));
}
let parts = usize::try_from(children.div_ceil(count(PART_ROWS)))
.map_err(|_| malformed("a forward link with more parts than fit in memory"))?;
let head = parts * 16;
let rows = usize::try_from(children)
.map_err(|_| malformed("a forward link longer than fits in memory"))?;
let packed = bitpack::tail_len(rows, width);
if rest.len() != head + packed {
return Err(malformed(
"a forward link's body is not the size its header implies",
));
}
let mut heads = Vec::with_capacity(parts);
for part in 0..parts {
let low = number(&rest[part * 16..part * 16 + 8])?;
let high = number(&rest[part * 16 + 8..part * 16 + 16])?;
heads.push((low != NO_PARENT).then_some((low, high)));
}
Body::Packed { bytes: rest[head..].to_vec(), width, heads }
}
Form::Monotone => {
let len = usize::try_from(linked + parents)
.map_err(|_| malformed("a forward link longer than fits in memory"))?;
Body::Monotone { vector: BitVector::read(rest, len)? }
}
};
Ok(Self { children, parents, linked, body })
}
}
fn runs(parents_of: &[Rid], parents: u64) -> Result<BitVector> {
let len = usize::try_from(count(parents_of.len()) + parents)
.map_err(|_| malformed("a forward link longer than fits in memory"))?;
let mut words = vec![0_u64; len.div_ceil(64)];
let mut at = 0_usize;
let mut child = 0_usize;
for parent in 0..parents {
while child < parents_of.len() && parents_of[child] == parent {
words[at / 64] |= 1 << (at % 64);
at += 1;
child += 1;
}
at += 1;
}
debug_assert_eq!(at, len, "every child is a one and every parent is a zero");
BitVector::new(words, len)
}
fn packed(parents_of: &[Rid], parents: u64) -> Result<Body> {
let width = width_for(parents);
let absent = reserved(width);
let mut heads = Vec::with_capacity(parents_of.len().div_ceil(PART_ROWS));
for rows in parents_of.chunks(PART_ROWS) {
let mut bounds: Bounds = None;
for parent in rows {
if *parent == NO_PARENT {
continue;
}
bounds = Some(match bounds {
None => (*parent, *parent),
Some((low, high)) => (low.min(*parent), high.max(*parent)),
});
}
heads.push(bounds);
}
let values = parents_of
.iter()
.map(|parent| if *parent == NO_PARENT { absent } else { *parent })
.collect::<Vec<u64>>();
let mut bytes = Vec::with_capacity(bitpack::tail_len(values.len(), width));
bitpack::pack_linear(&values, width, &mut bytes)?;
Ok(Body::Packed { bytes, width, heads })
}
fn width_for(parents: u64) -> usize {
(u64::BITS - parents.leading_zeros()).max(1) as usize
}
fn reserved(width: usize) -> u64 {
if width >= 64 { u64::MAX } else { (1_u64 << width) - 1 }
}
fn walk_ones(words: &[u64], at: usize, mut parent: Rid, skip: Rid) -> Option<(usize, Rid)> {
if skip == 0 {
return Some((at, parent));
}
let mut left = skip - 1;
let mut from = at + 1;
loop {
let index = from / 64;
let offset = from % 64;
let word = *words.get(index)? >> offset;
let span = 64 - offset;
let ones = u64::from(word.count_ones());
if ones > left {
#[expect(clippy::cast_possible_truncation, reason = "under the ones in one word")]
let within = crate::bits::nth_set(word, left as u32) as usize;
parent += count(within) - left;
return Some((from + within, parent));
}
parent += count(span) - ones;
left -= ones;
from += span;
}
}
fn count(rows: usize) -> u64 {
u64::try_from(rows).unwrap_or(u64::MAX)
}
fn number(bytes: &[u8]) -> Result<u64> {
Ok(u64::from_le_bytes(
bytes.try_into().map_err(|_| malformed("a forward link header is torn"))?,
))
}
fn malformed(message: impl Into<String>) -> Error {
Error::invalid_input(format!("invalid rudb forward link: {}", message.into()))
}
#[cfg(test)]
mod tests {
use super::*;
fn resolves(parents_of: &[Rid], parents: u64) -> Link {
let built = Link::build(parents_of, parents).expect("build");
let mut bytes = Vec::new();
built.write(&mut bytes).expect("write");
let read = Link::read(&bytes).expect("read");
let counts = Link::counts(&bytes[..HEADER_BYTES]).expect("the header alone");
assert_eq!(
counts,
Counts {
children: built.children(),
parents: built.parents(),
linked: built.linked(),
form: built.form()
}
);
assert_eq!(read.form(), built.form(), "the form survives the round trip");
assert_eq!(read.children(), built.children());
assert_eq!(read.parents(), built.parents());
assert_eq!(read.linked(), built.linked());
for link in [&built, &read] {
for (child, parent) in parents_of.iter().enumerate() {
let want = (*parent != NO_PARENT).then_some(*parent);
assert_eq!(link.forward(child as Rid), want, "child {child}");
}
assert_eq!(link.forward(parents_of.len() as Rid), None, "past the last child");
}
built
}
#[test]
fn a_clustered_child_takes_the_monotone_form_and_answers_both_directions() {
let link = resolves(&[0, 0, 0, 2, 2], 3);
assert_eq!(link.form(), Form::Monotone);
assert_eq!(link.backward(0), Some(0..3));
assert_eq!(link.backward(1), Some(3..3), "a parent with no children, not a missing parent");
assert_eq!(link.backward(2), Some(3..5));
assert_eq!(link.backward(3), None, "past the last parent");
}
#[test]
fn an_unclustered_child_takes_the_packed_form_and_answers_one_direction() {
let link = resolves(&[4, 1, 4, 0, 2], 5);
assert_eq!(link.form(), Form::Packed);
assert_eq!(link.backward(0), None, "the packed form does not answer backward");
}
#[test]
fn a_child_with_no_parent_keeps_the_link_out_of_the_monotone_form() {
let link = resolves(&[0, 1, NO_PARENT, 2], 3);
assert_eq!(link.form(), Form::Packed);
assert_eq!(link.linked(), 3, "the orphan is not linked and the other three are");
}
#[test]
fn every_child_pointing_at_one_parent_is_one_run() {
let link = resolves(&[7; 50], 8);
assert_eq!(link.form(), Form::Monotone);
assert_eq!(link.backward(6), Some(0..0));
assert_eq!(link.backward(7), Some(0..50));
}
#[test]
fn a_link_with_no_children_builds_and_resolves_nothing() {
let link = resolves(&[], 10);
assert_eq!(link.children(), 0);
assert_eq!(link.forward(0), None);
assert_eq!(link.part_bounds(0), None, "there is no part zero of an empty table");
}
#[test]
fn a_link_whose_parent_table_is_empty_is_packed_and_matches_nothing() {
let link = resolves(&[NO_PARENT, NO_PARENT], 0);
assert_eq!(link.form(), Form::Packed);
assert_eq!(link.linked(), 0);
}
#[test]
fn a_parent_rid_past_the_parent_table_is_refused_rather_than_stored() {
let error = Link::build(&[0, 9], 5).expect_err("refused");
assert!(error.to_string().contains("parent 9"), "{error}");
}
#[test]
fn the_reserved_value_is_not_a_parent_rid_even_at_the_width_boundary() {
assert_eq!(width_for(3), 2);
assert_eq!(width_for(4), 3);
assert_eq!(reserved(2), 3);
let link = resolves(&[2, 0, NO_PARENT], 3);
assert_eq!(link.form(), Form::Packed);
}
#[test]
fn a_packed_link_carries_the_bounds_of_every_part() {
let mut parents_of = vec![0_u64; PART_ROWS * 2 + 5];
for (child, parent) in parents_of.iter_mut().enumerate() {
*parent = (PART_ROWS - child % PART_ROWS) as u64;
}
let link = Link::build(&parents_of, PART_ROWS as u64 + 1).expect("build");
assert_eq!(link.form(), Form::Packed);
assert_eq!(link.part_bounds(0), Some(Some((1, PART_ROWS as u64))));
assert_eq!(link.part_bounds(2), Some(Some((PART_ROWS as u64 - 4, PART_ROWS as u64))));
assert_eq!(link.part_bounds(3), None, "there is no fourth part");
}
#[test]
fn a_part_in_which_nothing_matched_is_reported_as_skippable() {
let mut parents_of = vec![NO_PARENT; PART_ROWS * 2];
parents_of[PART_ROWS] = 3;
let link = Link::build(&parents_of, 10).expect("build");
assert_eq!(link.part_bounds(0), Some(None), "a part a reduction can skip outright");
assert_eq!(link.part_bounds(1), Some(Some((3, 3))));
}
#[test]
fn a_monotone_link_derives_its_part_bounds_from_its_ends() {
let parents_of = (0..PART_ROWS as u64 * 2).map(|child| child / 4).collect::<Vec<Rid>>();
let link = Link::build(&parents_of, PART_ROWS as u64).expect("build");
assert_eq!(link.form(), Form::Monotone);
assert_eq!(link.part_bounds(0), Some(Some((0, (PART_ROWS as u64 - 1) / 4))));
assert_eq!(
link.part_bounds(1),
Some(Some((PART_ROWS as u64 / 4, (PART_ROWS as u64 * 2 - 1) / 4)))
);
}
#[test]
fn the_monotone_form_is_a_bit_per_child_and_the_packed_form_is_a_rid_per_child() {
let ordered = (0..10_000_u64).map(|child| child / 10).collect::<Vec<Rid>>();
let monotone = Link::build(&ordered, 1000).expect("build");
assert_eq!(monotone.form(), Form::Monotone);
let mut shuffled = ordered.clone();
shuffled.swap(0, 9999);
let packed = Link::build(&shuffled, 1000).expect("build");
assert_eq!(packed.form(), Form::Packed);
assert!(
monotone.bytes() * 4 < packed.bytes(),
"monotone {} is not far below packed {}",
monotone.bytes(),
packed.bytes()
);
}
#[test]
fn a_payload_shorter_than_its_header_is_refused() {
let link = Link::build(&[0, 1], 2).expect("build");
let mut bytes = Vec::new();
link.write(&mut bytes).expect("write");
for cut in [0, 1, HEADER_BYTES - 1] {
assert!(Link::read(&bytes[..cut]).is_err(), "a payload of {cut} bytes is refused");
assert!(Link::counts(&bytes[..cut]).is_err(), "a header of {cut} bytes is refused");
}
}
#[test]
fn a_form_or_a_layout_this_build_does_not_know_is_refused() {
let link = Link::build(&[0, 1], 2).expect("build");
let mut bytes = Vec::new();
link.write(&mut bytes).expect("write");
let mut wrong = bytes.clone();
wrong[24] = 9;
assert!(Link::read(&wrong).is_err(), "an unknown form is refused");
assert!(Link::counts(&wrong).is_err(), "and its counts are not read");
let mut wrong = bytes;
wrong[26] = LAYOUT + 1;
assert!(Link::read(&wrong).is_err(), "an unknown layout is refused");
}
#[test]
fn a_width_that_does_not_match_the_parents_is_refused_rather_than_read_at() {
let link = Link::build(&[1, 0], 2).expect("build");
let mut bytes = Vec::new();
link.write(&mut bytes).expect("write");
bytes[25] = 7;
let error = Link::read(&bytes).expect_err("refused");
assert!(error.to_string().contains("width"), "{error}");
}
#[test]
fn a_truncated_body_is_refused_for_either_form() {
for parents_of in [vec![0_u64, 0, 1, 2], vec![2_u64, 0, 1, 0]] {
let link = Link::build(&parents_of, 3).expect("build");
let mut bytes = Vec::new();
link.write(&mut bytes).expect("write");
let short = &bytes[..bytes.len() - 1];
assert!(Link::read(short).is_err(), "a truncated {:?} body is refused", link.form());
}
}
#[test]
fn a_run_agrees_with_the_per_child_lookup_in_both_forms() {
let mut clustered = Vec::new();
for parent in 0..400_u64 {
let children = if parent % 50 < 45 { 0 } else { parent % 7 + 1 };
clustered.extend(std::iter::repeat_n(parent, children as usize));
}
let scattered: Vec<Rid> = (0..3000_u64)
.map(|child| if child % 13 == 0 { NO_PARENT } else { (child * 37) % 500 })
.collect();
for (parents_of, parents) in [(clustered, 400), (scattered, 500)] {
let link = Link::build(&parents_of, parents).expect("build");
let children = link.children();
for first in [0, 1, 63, 64, 65, children / 2, children - 1] {
for len in [0, 1, 2, 100, children - first] {
let len = len.min(children - first);
let mut out = vec![0; len as usize];
link.forward_run(first, &mut out).expect("in range");
let want: Vec<Rid> = (first..first + len)
.map(|child| link.forward(child).unwrap_or(NO_PARENT))
.collect();
assert_eq!(out, want, "{:?} from {first} for {len}", link.form());
}
}
let mut out = vec![0; 2];
assert!(link.forward_run(children - 1, &mut out).is_err(), "past the last child");
}
}
#[test]
fn a_list_agrees_with_the_per_child_lookup_in_both_forms() {
let mut clustered = Vec::new();
for parent in 0..3000_u64 {
let children = if parent % 50 < 45 { parent % 3 } else { parent % 7 + 1 };
clustered.extend(std::iter::repeat_n(parent, children as usize));
}
let scattered: Vec<Rid> = (0..3000_u64)
.map(|child| if child % 13 == 0 { NO_PARENT } else { (child * 37) % 500 })
.collect();
for (parents_of, parents) in [(clustered, 3000), (scattered, 500)] {
let link = Link::build(&parents_of, parents).expect("build");
let children = link.children();
let mut lists: Vec<Vec<Rid>> = vec![
Vec::new(),
(0..children).collect(),
(0..children).step_by(7).collect(),
(0..children).step_by(1500).collect(),
(0..children).rev().step_by(11).collect(),
vec![5, 5, 4, children - 1, children, children + 9, 0, 63, 64, 65],
];
let mut state = 7_u64;
let mut sparse = Vec::new();
for child in 0..children {
state = state.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1);
if state >> 60 == 0 {
sparse.push(child);
}
}
lists.push(sparse);
let mut out = Vec::new();
for list in &lists {
link.forward_each(list, &mut out);
let want: Vec<Rid> =
list.iter().map(|&child| link.forward(child).unwrap_or(NO_PARENT)).collect();
assert_eq!(out, want, "{:?} over {} children", link.form(), list.len());
}
}
}
}