use core::cmp;
use core::ops::{Index, IndexMut, RangeBounds};
use crate::local_prelude::*;
use crate::util::round_up_to_next;
#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(
feature = "miniserde",
derive(miniserde::Serialize, miniserde::Deserialize)
)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct BitMatrix<B: BitBlock = u32> {
bit_vec: BitVec<B>,
row_bits: usize,
}
impl<B: BitBlock> BitMatrix<B> {
pub fn new(rows: usize, row_bits: usize) -> Self {
BitMatrix {
bit_vec: BitVec::from_elem_general(round_up_to_next(row_bits, B::bits()) * rows, false),
row_bits,
}
}
#[inline]
fn num_rows(&self) -> usize {
if self.row_bits == 0 {
0
} else {
let row_blocks = round_up_to_next(self.row_bits, B::bits()) / B::bits();
self.bit_vec.storage().len() / row_blocks
}
}
#[inline]
pub fn num_cols(&self) -> usize {
self.row_bits
}
pub fn size(&self) -> (usize, usize) {
(self.num_rows(), self.row_bits)
}
#[inline]
pub fn set(&mut self, row: usize, col: usize, enabled: bool) {
let row_size_in_bits = round_up_to_next(self.row_bits, B::bits());
self.bit_vec.set(row * row_size_in_bits + col, enabled);
}
#[inline]
pub fn fill(&mut self, enabled: bool) {
self.bit_vec.fill(enabled);
}
pub fn grow(&mut self, num_rows: usize, value: bool) {
self.bit_vec
.grow(round_up_to_next(self.row_bits, B::bits()) * num_rows, value);
}
pub fn truncate(&mut self, num_rows: usize) {
self.bit_vec
.truncate(round_up_to_next(self.row_bits, B::bits()) * num_rows);
}
#[inline]
pub fn sub_matrix<R: RangeBounds<usize>>(&self, range: R) -> BitSubMatrix<'_, B> {
let row_size = round_up_to_next(self.row_bits, B::bits()) / B::bits();
BitSubMatrix {
slice: &self.bit_vec.storage()[(
range.start_bound().map(|&s| s * row_size),
range.end_bound().map(|&e| e * row_size),
)],
row_bits: self.row_bits,
}
}
#[inline]
pub fn sub_matrix_mut<R: RangeBounds<usize>>(&mut self, range: R) -> BitSubMatrixMut<'_, B> {
let row_size = self.row_size();
unsafe {
BitSubMatrixMut {
slice: &mut self.bit_vec.storage_mut()[(
range.start_bound().map(|&s| s * row_size),
range.end_bound().map(|&e| e * row_size),
)],
row_bits: self.row_bits,
}
}
}
fn row_size(&self) -> usize {
round_up_to_next(self.row_bits, B::bits()) / B::bits()
}
#[inline]
pub fn split_at(&self, row: usize) -> (BitSubMatrix<'_, B>, BitSubMatrix<'_, B>) {
(
self.sub_matrix(0..row),
self.sub_matrix(row..self.num_rows()),
)
}
#[inline]
pub fn split_at_mut(&mut self, row: usize) -> (BitSubMatrixMut<'_, B>, BitSubMatrixMut<'_, B>) {
let row_size = round_up_to_next(self.row_bits, B::bits()) / B::bits();
let (first, second) = unsafe { self.bit_vec.storage_mut().split_at_mut(row * row_size) };
(
BitSubMatrixMut::new(first, self.row_bits),
BitSubMatrixMut::new(second, self.row_bits),
)
}
pub fn iter_row(&self, row: usize) -> impl Iterator<Item = bool> + '_ {
BitSlice::new(&self[row].slice).iter_bits(self.row_bits)
}
pub fn transitive_closure(&mut self) {
Into::<BitSubMatrixMut<B>>::into(self).transitive_closure();
}
pub fn is_square(&self) -> bool {
self.num_rows() == self.row_bits
}
pub fn is_empty(&self) -> bool {
self.size() == (0, 0)
}
pub fn reflexive_closure(&mut self) {
for i in 0..cmp::min(self.row_bits, self.num_rows()) {
self.set(i, i, true);
}
}
}
impl<B: BitBlock> Index<usize> for BitMatrix<B> {
type Output = BitSlice<B>;
#[inline]
fn index(&self, row: usize) -> &Self::Output {
let row_size = round_up_to_next(self.row_bits, B::bits()) / B::bits();
BitSlice::new(&self.bit_vec.storage()[row * row_size..(row + 1) * row_size])
}
}
impl<B: BitBlock> IndexMut<usize> for BitMatrix<B> {
#[inline]
fn index_mut(&mut self, row: usize) -> &mut Self::Output {
let row_size = round_up_to_next(self.row_bits, B::bits()) / B::bits();
unsafe {
BitSlice::new_mut(&mut self.bit_vec.storage_mut()[row * row_size..(row + 1) * row_size])
}
}
}
impl<B: BitBlock> Index<(usize, usize)> for BitMatrix<B> {
type Output = bool;
#[inline]
fn index(&self, (row, col): (usize, usize)) -> &bool {
let row_size_in_bits = round_up_to_next(self.row_bits, B::bits());
if self.bit_vec.get(row * row_size_in_bits + col).unwrap() {
&TRUE
} else {
&FALSE
}
}
}
impl<'a, B: BitBlock> From<&'a mut BitMatrix<B>> for BitSubMatrixMut<'a, B> {
fn from(value: &'a mut BitMatrix<B>) -> Self {
unsafe { BitSubMatrixMut::new(value.bit_vec.storage_mut(), value.row_bits) }
}
}
#[test]
fn test_empty() {
let mut matrix = <BitMatrix>::new(0, 0);
for _ in 0..3 {
assert_eq!(matrix.num_rows(), 0);
assert_eq!(matrix.size(), (0, 0));
assert!(matrix.is_square());
assert!(matrix.is_empty());
matrix.transitive_closure();
}
}