#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct FragmentId(pub u32);
impl FragmentId {
pub fn index(self) -> usize {
self.0 as usize
}
}
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct TreeNode {
pub parent: Option<FragmentId>,
pub children: Option<[FragmentId; 2]>,
pub depth: u16,
pub split_at: Option<u32>,
}
impl TreeNode {
pub fn is_leaf(&self) -> bool {
self.children.is_none()
}
}
#[derive(Clone, Debug, Default, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct FragmentTree {
nodes: Vec<TreeNode>,
roots: Vec<FragmentId>,
cuts: u32,
}
impl FragmentTree {
pub(crate) fn from_nodes(nodes: Vec<TreeNode>, cuts: u32) -> Self {
let roots = nodes
.iter()
.enumerate()
.filter(|(_, n)| n.parent.is_none())
.map(|(i, _)| FragmentId(i as u32))
.collect();
FragmentTree { nodes, roots, cuts }
}
pub fn len(&self) -> usize {
self.nodes.len()
}
pub fn is_empty(&self) -> bool {
self.nodes.is_empty()
}
pub fn cuts(&self) -> u32 {
self.cuts
}
pub fn node(&self, id: FragmentId) -> Option<&TreeNode> {
self.nodes.get(id.index())
}
pub fn roots(&self) -> &[FragmentId] {
&self.roots
}
pub fn iter(&self) -> impl Iterator<Item = (FragmentId, &TreeNode)> {
self.nodes.iter().enumerate().map(|(i, n)| (FragmentId(i as u32), n))
}
pub fn leaves(&self) -> Vec<FragmentId> {
self.iter().filter(|(_, n)| n.is_leaf()).map(|(id, _)| id).collect()
}
pub fn frontier_after(&self, cuts: u32) -> Vec<FragmentId> {
self.iter()
.filter(|(_, n)| {
let born = match n.parent.and_then(|p| self.node(p)).and_then(|p| p.split_at) {
Some(t) => t < cuts,
None => n.parent.is_none(),
};
born && n.split_at.is_none_or(|t| t >= cuts)
})
.map(|(id, _)| id)
.collect()
}
pub fn frontier_of(&self, count: usize) -> Vec<FragmentId> {
let floor = self.roots.len();
let cuts = count.saturating_sub(floor).min(self.cuts as usize) as u32;
self.frontier_after(cuts)
}
pub fn at_depth(&self, depth: u16) -> Vec<FragmentId> {
self.iter()
.filter(|(_, n)| n.depth <= depth && (n.is_leaf() || n.depth == depth))
.map(|(id, _)| id)
.collect()
}
pub fn descendants(&self, id: FragmentId) -> Vec<FragmentId> {
let mut out = Vec::new();
let mut stack = vec![id];
while let Some(n) = stack.pop() {
let Some(node) = self.node(n) else { continue };
if let Some(kids) = node.children {
for k in kids {
out.push(k);
stack.push(k);
}
}
}
out.sort_unstable();
out
}
pub fn root_of(&self, id: FragmentId) -> Option<FragmentId> {
let mut cur = id;
for _ in 0..=self.nodes.len() {
let node = self.node(cur)?;
match node.parent {
Some(p) => cur = p,
None => return Some(cur),
}
}
None
}
pub fn is_ancestor(&self, ancestor: FragmentId, id: FragmentId) -> bool {
let mut cur = id;
for _ in 0..=self.nodes.len() {
let Some(node) = self.node(cur) else { return false };
match node.parent {
Some(p) if p == ancestor => return true,
Some(p) => cur = p,
None => return false,
}
}
false
}
}
#[cfg(test)]
mod tests {
use super::*;
fn forest() -> FragmentTree {
let n = |parent, children, depth, split_at| TreeNode { parent, children, depth, split_at };
FragmentTree::from_nodes(
vec![
n(None, Some([FragmentId(1), FragmentId(2)]), 0, Some(0)),
n(Some(FragmentId(0)), Some([FragmentId(3), FragmentId(4)]), 1, Some(1)),
n(Some(FragmentId(0)), Some([FragmentId(5), FragmentId(6)]), 1, Some(2)),
n(Some(FragmentId(1)), None, 2, None),
n(Some(FragmentId(1)), None, 2, None),
n(Some(FragmentId(2)), None, 2, None),
n(Some(FragmentId(2)), None, 2, None),
],
3,
)
}
#[test]
fn every_frontier_is_an_antichain_that_covers_once() {
let t = forest();
for cuts in 0..=t.cuts() {
let f = t.frontier_after(cuts);
assert_eq!(f.len(), t.roots().len() + cuts as usize, "cuts={cuts} wrong piece count");
for &a in &f {
for &b in &f {
assert!(
a == b || !t.is_ancestor(a, b),
"cuts={cuts}: {a:?} is an ancestor of {b:?} — the frontier double-covers"
);
}
}
}
}
#[test]
fn the_coarsest_frontier_is_the_roots_and_the_finest_is_the_leaves() {
let t = forest();
assert_eq!(t.frontier_after(0), t.roots());
assert_eq!(t.frontier_after(t.cuts()), t.leaves());
}
#[test]
fn frontier_of_clamps_instead_of_failing() {
let t = forest();
assert_eq!(t.frontier_of(0).len(), 1, "below the root count clamps up to the roots");
assert_eq!(t.frontier_of(3).len(), 3);
assert_eq!(t.frontier_of(9_999), t.leaves(), "above the cut count clamps down to the leaves");
}
#[test]
fn at_depth_cuts_every_branch_to_the_same_level() {
let t = forest();
assert_eq!(t.at_depth(0), vec![FragmentId(0)]);
assert_eq!(t.at_depth(1), vec![FragmentId(1), FragmentId(2)]);
assert_eq!(t.at_depth(2), t.leaves());
assert_eq!(t.at_depth(9), t.leaves(), "past the deepest level is the leaves, not empty");
}
#[test]
fn descendants_and_ancestry_agree() {
let t = forest();
assert_eq!(
t.descendants(FragmentId(0)),
vec![FragmentId(1), FragmentId(2), FragmentId(3), FragmentId(4), FragmentId(5), FragmentId(6)]
);
assert_eq!(t.descendants(FragmentId(3)), vec![]);
for d in t.descendants(FragmentId(1)) {
assert!(t.is_ancestor(FragmentId(1), d));
assert!(t.is_ancestor(FragmentId(0), d), "ancestry is transitive to the root");
}
assert!(!t.is_ancestor(FragmentId(1), FragmentId(5)));
assert_eq!(t.root_of(FragmentId(6)), Some(FragmentId(0)));
}
#[test]
fn an_out_of_range_id_is_refused_rather_than_fatal() {
let t = forest();
assert!(t.node(FragmentId(99)).is_none());
assert!(t.root_of(FragmentId(99)).is_none());
assert!(!t.is_ancestor(FragmentId(0), FragmentId(99)));
assert_eq!(t.descendants(FragmentId(99)), vec![]);
}
}