use std::ops::Range;
pub(crate) struct ConsumerMatrix {
words: usize,
bits: Vec<u64>,
}
impl ConsumerMatrix {
pub(crate) fn new(units: usize, consumers: usize) -> Option<Self> {
let words = consumers.div_ceil(u64::BITS as usize);
let len = units.checked_mul(words)?;
len.checked_mul(size_of::<u64>())?;
Some(Self {
words,
bits: vec![0u64; len],
})
}
pub(crate) fn insert(&mut self, unit: usize, consumer: usize) {
let offset = unit * self.words + consumer / u64::BITS as usize;
self.bits[offset] |= 1u64 << (consumer % u64::BITS as usize);
}
pub(crate) fn row(&self, unit: usize) -> &[u64] {
let start = unit * self.words;
&self.bits[start..start + self.words]
}
pub(crate) fn count(&self, unit: usize) -> u32 {
self.row(unit).iter().map(|word| word.count_ones()).sum()
}
pub(crate) fn single(&self, unit: usize) -> usize {
for (index, word) in self.row(unit).iter().copied().enumerate() {
if word != 0 {
return index * u64::BITS as usize + word.trailing_zeros() as usize;
}
}
0
}
pub(crate) fn expand(&self, unit: usize) -> Vec<usize> {
let mut consumers = Vec::with_capacity(self.count(unit) as usize);
for (index, mut word) in self.row(unit).iter().copied().enumerate() {
while word != 0 {
consumers.push(index * u64::BITS as usize + word.trailing_zeros() as usize);
word &= word - 1;
}
}
consumers
}
}
pub(crate) fn group_runs(units: &[usize], matrix: &ConsumerMatrix) -> Vec<Range<usize>> {
let mut runs: Vec<Range<usize>> = Vec::new();
let mut start = 0usize;
while start < units.len() {
let row = matrix.row(units[start]);
let mut end = start + 1;
while end < units.len() && matrix.row(units[end]) == row {
end += 1;
}
runs.push(start..end);
start = end;
}
runs
}
#[cfg(test)]
mod tests {
use super::*;
fn matrix(units: usize, consumers: usize, edges: &[(usize, usize)]) -> ConsumerMatrix {
let mut matrix = ConsumerMatrix::new(units, consumers).expect("matrix");
for (unit, consumer) in edges {
matrix.insert(*unit, *consumer);
}
matrix
}
#[test]
fn counts_and_expands_consumer_sets() {
let matrix = matrix(3, 3, &[(0, 0), (1, 0), (1, 2), (2, 1)]);
assert_eq!(matrix.count(0), 1);
assert_eq!(matrix.single(0), 0);
assert_eq!(matrix.expand(0), [0]);
assert_eq!(matrix.count(1), 2);
assert_eq!(matrix.expand(1), [0, 2]);
assert_eq!(matrix.count(2), 1);
assert_eq!(matrix.single(2), 1);
}
#[test]
fn reports_no_consumers_for_unreached_units() {
let matrix = matrix(2, 2, &[(0, 1)]);
assert_eq!(matrix.count(1), 0);
assert!(matrix.expand(1).is_empty());
assert_eq!(matrix.single(1), 0);
}
#[test]
fn spans_consumer_sets_wider_than_one_word() {
let matrix = matrix(2, 130, &[(0, 0), (0, 129), (1, 129)]);
assert_eq!(matrix.count(0), 2);
assert_eq!(matrix.expand(0), [0, 129]);
assert_eq!(matrix.count(1), 1);
assert_eq!(matrix.single(1), 129);
assert_ne!(matrix.row(0), matrix.row(1));
}
#[test]
fn groups_only_adjacent_units_with_equal_consumer_sets() {
let matrix = matrix(3, 2, &[(0, 0), (0, 1), (1, 0), (2, 0), (2, 1)]);
assert_eq!(group_runs(&[0, 1, 2], &matrix), [0..1, 1..2, 2..3]);
assert_eq!(group_runs(&[0, 2, 1], &matrix), [0..2, 2..3]);
}
#[test]
fn groups_an_empty_slice_into_no_runs() {
let matrix = matrix(1, 1, &[]);
assert!(group_runs(&[], &matrix).is_empty());
}
#[test]
fn rejects_a_matrix_that_cannot_be_indexed() {
assert!(ConsumerMatrix::new(usize::MAX, 64).is_none());
}
}