use crate::canonical::ir::{drop_subsumed, Bounds, Divisors, IntegerBounds, IntegerLeaf};
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct IntegerLeaves {
leaves: Vec<IntegerLeaf>,
canonical: bool,
}
impl Default for IntegerLeaves {
fn default() -> Self {
Self {
leaves: Vec::new(),
canonical: true,
}
}
}
impl IntegerLeaves {
pub(crate) fn insert(&mut self, leaf: IntegerLeaf) {
self.leaves.push(leaf);
self.canonical = false;
}
fn canonicalize(&mut self) {
if self.canonical {
return;
}
let was_empty = self.leaves.is_empty();
self.leaves = merge(std::mem::take(&mut self.leaves));
drop_subsumed(&mut self.leaves, |outer, inner| {
outer.bounds.covers(&inner.bounds) && outer.multiple_of.divide_all(&inner.multiple_of)
});
self.canonical = true;
debug_assert_eq!(
self.leaves.is_empty(),
was_empty,
"merging emptied the leaves"
);
}
pub(crate) fn clear(&mut self) {
self.leaves.clear();
self.canonical = true;
}
pub(crate) fn retain(&mut self, keep: impl FnMut(&IntegerLeaf) -> bool) {
self.canonicalize();
self.leaves.retain(keep);
}
pub(crate) fn is_empty(&self) -> bool {
self.leaves.is_empty()
}
pub(crate) fn as_slice(&mut self) -> &[IntegerLeaf] {
self.canonicalize();
&self.leaves
}
}
impl IntoIterator for IntegerLeaves {
type Item = IntegerLeaf;
type IntoIter = std::vec::IntoIter<IntegerLeaf>;
fn into_iter(mut self) -> Self::IntoIter {
self.canonicalize();
self.leaves.into_iter()
}
}
struct Group {
divisor: Divisors,
}
fn merge(mut leaves: Vec<IntegerLeaf>) -> Vec<IntegerLeaf> {
if leaves.len() < 2 {
return leaves;
}
leaves.sort_by(|left, right| left.multiple_of.cmp(&right.multiple_of));
let mut merged: Vec<IntegerLeaf> = Vec::with_capacity(leaves.len());
let mut windows: Vec<IntegerBounds> = Vec::new();
let mut group: Option<Group> = None;
for leaf in leaves {
if group
.as_ref()
.is_none_or(|open| open.divisor != leaf.multiple_of)
{
flush_group(&mut merged, group.take(), &mut windows);
group = Some(Group {
divisor: leaf.multiple_of,
});
}
windows.push(leaf.bounds);
}
flush_group(&mut merged, group, &mut windows);
merged
}
fn flush_group(
merged: &mut Vec<IntegerLeaf>,
group: Option<Group>,
windows: &mut Vec<IntegerBounds>,
) {
let Some(Group {
divisor: multiple_of,
}) = group
else {
return;
};
for bounds in Bounds::merge_all(std::mem::take(windows)) {
merged.push(IntegerLeaf {
bounds,
multiple_of: multiple_of.clone(),
});
}
}