96 lines
2.7 KiB
Rust
96 lines
2.7 KiB
Rust
use either::Either;
|
|
use std::{iter::Rev, ops::RangeInclusive};
|
|
|
|
pub fn split_from<T>(range: RangeInclusive<T>, x: T) -> [impl Iterator<Item = T> + Clone; 2]
|
|
where
|
|
T: num::Integer + Copy,
|
|
RangeInclusive<T>: Iterator<Item = T> + DoubleEndedIterator,
|
|
Rev<RangeInclusive<T>>: Iterator<Item = T>,
|
|
{
|
|
let in_range = range.contains(&x);
|
|
let (start, end) = (*range.start(), *range.end());
|
|
|
|
// RangeInclusive (1..=0), has 0 elements
|
|
let base = Either::Right(T::one()..=T::zero());
|
|
[
|
|
if in_range && x > start + T::one() {
|
|
Either::Left((start..=(x - T::one())).rev())
|
|
} else {
|
|
base.clone()
|
|
},
|
|
if in_range && x + T::one() < end {
|
|
Either::Right((x + T::one())..=end)
|
|
} else {
|
|
base
|
|
},
|
|
]
|
|
}
|
|
|
|
pub fn diag_raw<T>(
|
|
i_chains: [impl Iterator<Item = T> + Clone; 2],
|
|
j_chains: [impl Iterator<Item = T> + Clone; 2],
|
|
) -> [impl Iterator<Item = (T, T)> + Clone; 4]
|
|
where
|
|
T: num::Integer + Copy,
|
|
RangeInclusive<T>: Iterator<Item = T> + DoubleEndedIterator,
|
|
Rev<RangeInclusive<T>>: Iterator<Item = T>,
|
|
{
|
|
[(0, 0), (1, 1), (1, 0), (0, 1)].map(move |(a, b)| i_chains[a].clone().zip(j_chains[b].clone()))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod test {
|
|
use super::*;
|
|
|
|
fn diag_test_helper<T>(
|
|
i: T,
|
|
j: T,
|
|
range_i: RangeInclusive<T>,
|
|
range_j: RangeInclusive<T>,
|
|
) -> [impl Iterator<Item = (T, T)> + Clone; 4]
|
|
where
|
|
T: num::Integer + Copy,
|
|
RangeInclusive<T>: Iterator<Item = T> + DoubleEndedIterator,
|
|
Rev<RangeInclusive<T>>: Iterator<Item = T>,
|
|
{
|
|
diag_raw(split_from(range_i, i), split_from(range_j, j))
|
|
}
|
|
|
|
#[test]
|
|
fn split_test() {
|
|
assert_eq!(
|
|
split_from(0..=6, 2).map(Iterator::collect::<Vec<usize>>),
|
|
[vec![1, 0], vec![3, 4, 5, 6]]
|
|
);
|
|
|
|
assert_eq!(
|
|
split_from(0..=6, 0).map(Iterator::collect::<Vec<usize>>),
|
|
[vec![], vec![1, 2, 3, 4, 5, 6]]
|
|
);
|
|
|
|
assert_eq!(
|
|
split_from(0..=6, 6).map(Iterator::collect::<Vec<usize>>),
|
|
[vec![5, 4, 3, 2, 1, 0], vec![]]
|
|
);
|
|
|
|
// test out-of-bounds and also generics
|
|
assert_eq!(
|
|
split_from::<i16>(-1i16..=4i16, 10i16).map(Iterator::collect::<Vec<i16>>),
|
|
[const { Vec::new() }; 2]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn diag_test() {
|
|
assert_eq!(
|
|
diag_test_helper(2, 3, 0..=7, 0..=7).map(Iterator::collect::<Vec<(usize, usize)>>),
|
|
[
|
|
vec![(1, 2), (0, 1)],
|
|
vec![(3, 4), (4, 5), (5, 6), (6, 7)],
|
|
vec![(3, 2), (4, 1), (5, 0)],
|
|
vec![(1, 4), (0, 5)]
|
|
]
|
|
);
|
|
}
|
|
}
|