use crate::error::{Error, Result};
use crate::identical::Identical;
use crate::rle::{Rle, RleLike, RleRef};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PlanMany {
pub ends: Vec<i64>,
pub idx: Vec<Vec<usize>>,
}
fn walk(inputs: &[&[i64]], mut visit: impl FnMut(i64, &[usize])) {
let mut pos = vec![0usize; inputs.len()];
loop {
let end = inputs
.iter()
.zip(&pos)
.filter_map(|(ends, &p)| ends.get(p).copied())
.min();
let Some(end) = end else { return };
visit(end, &pos);
for (ends, p) in inputs.iter().zip(pos.iter_mut()) {
if ends.get(*p) == Some(&end) {
*p += 1;
}
}
}
}
pub fn plan_many(inputs: &[&[i64]]) -> PlanMany {
let mut out = PlanMany {
ends: Vec::new(),
idx: vec![Vec::new(); inputs.len()],
};
walk(inputs, |end, pos| {
out.ends.push(end);
for (column, &p) in out.idx.iter_mut().zip(pos) {
column.push(p);
}
});
out
}
pub fn zip_many<V, C, F>(inputs: &[RleRef<'_, V>], fill: Option<&V>, mut f: F) -> Result<Rle<C>>
where
C: Identical,
F: FnMut(&[&V]) -> C,
{
let lengths = inputs.iter().map(|r| r.len());
let (min, max) = (
lengths.clone().min().unwrap_or(0),
lengths.max().unwrap_or(0),
);
if min != max && fill.is_none() {
return Err(Error::UnequalLength {
left: min,
right: max,
});
}
let ends: Vec<&[i64]> = inputs.iter().map(|r| r.ends).collect();
let mut out = Rle::new();
let mut values: Vec<&V> = Vec::with_capacity(inputs.len());
walk(&ends, |end, pos| {
values.clear();
values.extend(inputs.iter().zip(pos).map(|(r, &p)| {
r.values
.get(p)
.or(fill)
.expect("the length check guarantees a fill value")
}));
out.push_merged(end, f(&values));
});
Ok(out)
}