1use core::cmp;
4use core::ops::{Index, IndexMut, RangeBounds};
5
6use crate::local_prelude::*;
7use crate::util::round_up_to_next;
8
9#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
11#[cfg_attr(
12 feature = "miniserde",
13 derive(miniserde::Serialize, miniserde::Deserialize)
14)]
15#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
16pub struct BitMatrix<B: BitBlock = u32> {
17 bit_vec: BitVec<B>,
18 row_bits: usize,
19}
20
21impl<B: BitBlock> BitMatrix<B> {
24 pub fn new(rows: usize, row_bits: usize) -> Self {
26 BitMatrix {
27 bit_vec: BitVec::from_elem_general(round_up_to_next(row_bits, B::bits()) * rows, false),
28 row_bits,
29 }
30 }
31
32 #[inline]
34 fn num_rows(&self) -> usize {
35 if self.row_bits == 0 {
36 0
37 } else {
38 let row_blocks = round_up_to_next(self.row_bits, B::bits()) / B::bits();
39 self.bit_vec.storage().len() / row_blocks
40 }
41 }
42
43 #[inline]
45 pub fn num_cols(&self) -> usize {
46 self.row_bits
47 }
48
49 pub fn size(&self) -> (usize, usize) {
51 (self.num_rows(), self.row_bits)
52 }
53
54 #[inline]
60 pub fn set(&mut self, row: usize, col: usize, enabled: bool) {
61 let row_size_in_bits = round_up_to_next(self.row_bits, B::bits());
62 self.bit_vec.set(row * row_size_in_bits + col, enabled);
63 }
64
65 #[inline]
67 pub fn fill(&mut self, enabled: bool) {
68 self.bit_vec.fill(enabled);
69 }
70
71 pub fn grow(&mut self, num_rows: usize, value: bool) {
73 self.bit_vec
74 .grow(round_up_to_next(self.row_bits, B::bits()) * num_rows, value);
75 }
76
77 pub fn truncate(&mut self, num_rows: usize) {
79 self.bit_vec
80 .truncate(round_up_to_next(self.row_bits, B::bits()) * num_rows);
81 }
82
83 #[inline]
85 pub fn sub_matrix<R: RangeBounds<usize>>(&self, range: R) -> BitSubMatrix<'_, B> {
86 let row_size = round_up_to_next(self.row_bits, B::bits()) / B::bits();
87 BitSubMatrix {
88 slice: &self.bit_vec.storage()[(
89 range.start_bound().map(|&s| s * row_size),
90 range.end_bound().map(|&e| e * row_size),
91 )],
92 row_bits: self.row_bits,
93 }
94 }
95
96 #[inline]
98 pub fn sub_matrix_mut<R: RangeBounds<usize>>(&mut self, range: R) -> BitSubMatrixMut<'_, B> {
99 let row_size = self.row_size();
100 unsafe {
103 BitSubMatrixMut {
104 slice: &mut self.bit_vec.storage_mut()[(
105 range.start_bound().map(|&s| s * row_size),
106 range.end_bound().map(|&e| e * row_size),
107 )],
108 row_bits: self.row_bits,
109 }
110 }
111 }
112
113 fn row_size(&self) -> usize {
114 round_up_to_next(self.row_bits, B::bits()) / B::bits()
115 }
116
117 #[inline]
123 pub fn split_at(&self, row: usize) -> (BitSubMatrix<'_, B>, BitSubMatrix<'_, B>) {
124 (
125 self.sub_matrix(0..row),
126 self.sub_matrix(row..self.num_rows()),
127 )
128 }
129
130 #[inline]
133 pub fn split_at_mut(&mut self, row: usize) -> (BitSubMatrixMut<'_, B>, BitSubMatrixMut<'_, B>) {
134 let row_size = round_up_to_next(self.row_bits, B::bits()) / B::bits();
135 let (first, second) = unsafe { self.bit_vec.storage_mut().split_at_mut(row * row_size) };
136 (
137 BitSubMatrixMut::new(first, self.row_bits),
138 BitSubMatrixMut::new(second, self.row_bits),
139 )
140 }
141
142 pub fn iter_row(&self, row: usize) -> impl Iterator<Item = bool> + '_ {
144 BitSlice::new(&self[row].slice).iter_bits(self.row_bits)
145 }
146
147 pub fn transitive_closure(&mut self) {
165 Into::<BitSubMatrixMut<B>>::into(self).transitive_closure();
166 }
167
168 pub fn is_square(&self) -> bool {
172 self.num_rows() == self.row_bits
173 }
174
175 pub fn is_empty(&self) -> bool {
177 self.size() == (0, 0)
178 }
179
180 pub fn reflexive_closure(&mut self) {
189 for i in 0..cmp::min(self.row_bits, self.num_rows()) {
190 self.set(i, i, true);
191 }
192 }
193}
194
195impl<B: BitBlock> Index<usize> for BitMatrix<B> {
197 type Output = BitSlice<B>;
198
199 #[inline]
200 fn index(&self, row: usize) -> &Self::Output {
201 let row_size = round_up_to_next(self.row_bits, B::bits()) / B::bits();
202 BitSlice::new(&self.bit_vec.storage()[row * row_size..(row + 1) * row_size])
203 }
204}
205
206impl<B: BitBlock> IndexMut<usize> for BitMatrix<B> {
208 #[inline]
209 fn index_mut(&mut self, row: usize) -> &mut Self::Output {
210 let row_size = round_up_to_next(self.row_bits, B::bits()) / B::bits();
211 unsafe {
214 BitSlice::new_mut(&mut self.bit_vec.storage_mut()[row * row_size..(row + 1) * row_size])
215 }
216 }
217}
218
219impl<B: BitBlock> Index<(usize, usize)> for BitMatrix<B> {
224 type Output = bool;
225
226 #[inline]
227 fn index(&self, (row, col): (usize, usize)) -> &bool {
228 let row_size_in_bits = round_up_to_next(self.row_bits, B::bits());
229 if self.bit_vec.get(row * row_size_in_bits + col).unwrap() {
230 &TRUE
231 } else {
232 &FALSE
233 }
234 }
235}
236
237impl<'a, B: BitBlock> From<&'a mut BitMatrix<B>> for BitSubMatrixMut<'a, B> {
238 fn from(value: &'a mut BitMatrix<B>) -> Self {
239 unsafe { BitSubMatrixMut::new(value.bit_vec.storage_mut(), value.row_bits) }
240 }
241}
242
243#[test]
246fn test_empty() {
247 let mut matrix = <BitMatrix>::new(0, 0);
248 for _ in 0..3 {
249 assert_eq!(matrix.num_rows(), 0);
250 assert_eq!(matrix.size(), (0, 0));
251 assert!(matrix.is_square());
252 assert!(matrix.is_empty());
253 matrix.transitive_closure();
254 }
255}