vortex_mask/
bitops.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::{BitAnd, BitOr, Not};
5
6use vortex_error::vortex_panic;
7
8use crate::{AllOr, Mask};
9
10impl BitAnd for &Mask {
11    type Output = Mask;
12
13    fn bitand(self, rhs: Self) -> Self::Output {
14        if self.len() != rhs.len() {
15            vortex_panic!("Masks must have the same length");
16        }
17
18        match (self.boolean_buffer(), rhs.boolean_buffer()) {
19            (AllOr::All, _) => rhs.clone(),
20            (_, AllOr::All) => self.clone(),
21            (AllOr::None, _) => Mask::new_false(self.len()),
22            (_, AllOr::None) => Mask::new_false(self.len()),
23            (AllOr::Some(lhs), AllOr::Some(rhs)) => Mask::from_buffer(lhs & rhs),
24        }
25    }
26}
27
28impl BitOr for &Mask {
29    type Output = Mask;
30
31    fn bitor(self, rhs: Self) -> Self::Output {
32        if self.len() != rhs.len() {
33            vortex_panic!("Masks must have the same length");
34        }
35
36        match (self.boolean_buffer(), rhs.boolean_buffer()) {
37            (AllOr::All, _) => Mask::new_true(self.len()),
38            (_, AllOr::All) => Mask::new_true(self.len()),
39            (AllOr::None, _) => rhs.clone(),
40            (_, AllOr::None) => self.clone(),
41            (AllOr::Some(lhs), AllOr::Some(rhs)) => Mask::from_buffer(lhs | rhs),
42        }
43    }
44}
45
46impl Not for &Mask {
47    type Output = Mask;
48
49    fn not(self) -> Self::Output {
50        match self.boolean_buffer() {
51            AllOr::All => Mask::new_false(self.len()),
52            AllOr::None => Mask::new_true(self.len()),
53            AllOr::Some(buffer) => Mask::from_buffer(!buffer),
54        }
55    }
56}
57
58#[cfg(test)]
59#[allow(clippy::many_single_char_names)]
60mod tests {
61    use arrow_buffer::BooleanBuffer;
62
63    use super::*;
64
65    #[test]
66    fn test_bitand_all_combinations() {
67        let len = 5;
68
69        // Test AllTrue & AllTrue
70        let all_true = Mask::new_true(len);
71        let result = &all_true & &all_true;
72        assert!(result.all_true());
73        assert_eq!(result.true_count(), len);
74
75        // Test AllTrue & AllFalse
76        let all_false = Mask::new_false(len);
77        let result = &all_true & &all_false;
78        assert!(result.all_false());
79        assert_eq!(result.true_count(), 0);
80
81        // Test AllFalse & AllTrue
82        let result = &all_false & &all_true;
83        assert!(result.all_false());
84        assert_eq!(result.true_count(), 0);
85
86        // Test AllFalse & AllFalse
87        let result = &all_false & &all_false;
88        assert!(result.all_false());
89        assert_eq!(result.true_count(), 0);
90    }
91
92    #[test]
93    fn test_bitand_with_values() {
94        let mask1 = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false, true]));
95        let mask2 = Mask::from_buffer(BooleanBuffer::from_iter([true, true, false, false, true]));
96
97        let result = &mask1 & &mask2;
98        assert_eq!(result.len(), 5);
99        assert_eq!(result.true_count(), 2);
100        assert!(result.value(0)); // true & true = true
101        assert!(!result.value(1)); // false & true = false
102        assert!(!result.value(2)); // true & false = false
103        assert!(!result.value(3)); // false & false = false
104        assert!(result.value(4)); // true & true = true
105    }
106
107    #[test]
108    fn test_bitand_all_true_with_values() {
109        let all_true = Mask::new_true(5);
110        let values = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false, true]));
111
112        // AllTrue & Values should return Values
113        let result = &all_true & &values;
114        assert_eq!(result.true_count(), 3);
115        assert_eq!(result.len(), 5);
116        assert!(result.value(0));
117        assert!(!result.value(1));
118        assert!(result.value(2));
119        assert!(!result.value(3));
120        assert!(result.value(4));
121    }
122
123    #[test]
124    fn test_bitand_all_false_with_values() {
125        let all_false = Mask::new_false(5);
126        let values = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false, true]));
127
128        // AllFalse & Values should return AllFalse
129        let result = &all_false & &values;
130        assert!(result.all_false());
131        assert_eq!(result.true_count(), 0);
132    }
133
134    #[test]
135    fn test_bitand_values_with_all_true() {
136        let values = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false, true]));
137        let all_true = Mask::new_true(5);
138
139        // Values & AllTrue should return Values
140        let result = &values & &all_true;
141        assert_eq!(result.true_count(), 3);
142        assert_eq!(result.len(), 5);
143        assert!(result.value(0));
144        assert!(!result.value(1));
145        assert!(result.value(2));
146        assert!(!result.value(3));
147        assert!(result.value(4));
148    }
149
150    #[test]
151    fn test_bitand_values_with_all_false() {
152        let values = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false, true]));
153        let all_false = Mask::new_false(5);
154
155        // Values & AllFalse should return AllFalse
156        let result = &values & &all_false;
157        assert!(result.all_false());
158        assert_eq!(result.true_count(), 0);
159    }
160
161    #[test]
162    fn test_bitand_empty_masks() {
163        let empty1 = Mask::new_true(0);
164        let empty2 = Mask::new_false(0);
165
166        let result = &empty1 & &empty2;
167        assert_eq!(result.len(), 0);
168        assert!(result.is_empty());
169    }
170
171    #[test]
172    #[should_panic(expected = "Masks must have the same length")]
173    fn test_bitand_different_lengths() {
174        let mask1 = Mask::new_true(5);
175        let mask2 = Mask::new_true(3);
176        let _ = &mask1 & &mask2;
177    }
178
179    #[test]
180    fn test_not_all_true() {
181        let all_true = Mask::new_true(5);
182        let result = !&all_true;
183        assert!(result.all_false());
184        assert_eq!(result.true_count(), 0);
185        assert_eq!(result.len(), 5);
186    }
187
188    #[test]
189    fn test_not_all_false() {
190        let all_false = Mask::new_false(5);
191        let result = !&all_false;
192        assert!(result.all_true());
193        assert_eq!(result.true_count(), 5);
194        assert_eq!(result.len(), 5);
195    }
196
197    #[test]
198    fn test_not_values() {
199        let values = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false, true]));
200        let result = !&values;
201
202        assert_eq!(result.len(), 5);
203        assert_eq!(result.true_count(), 2);
204        assert!(!result.value(0)); // !true = false
205        assert!(result.value(1)); // !false = true
206        assert!(!result.value(2)); // !true = false
207        assert!(result.value(3)); // !false = true
208        assert!(!result.value(4)); // !true = false
209    }
210
211    #[test]
212    fn test_not_empty() {
213        let empty_true = Mask::new_true(0);
214        let result = !&empty_true;
215        assert_eq!(result.len(), 0);
216        assert!(result.is_empty());
217
218        let empty_false = Mask::new_false(0);
219        let result = !&empty_false;
220        assert_eq!(result.len(), 0);
221        assert!(result.is_empty());
222    }
223
224    #[test]
225    fn test_double_not() {
226        let original =
227            Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false, true]));
228        let double_not = !&(!&original);
229
230        // Double negation should return the original
231        assert_eq!(double_not.true_count(), original.true_count());
232        for i in 0..5 {
233            assert_eq!(double_not.value(i), original.value(i));
234        }
235    }
236
237    #[test]
238    fn test_demorgan_law() {
239        // Test De Morgan's law: !(A & B) = !A | !B
240        let a = Mask::from_buffer(BooleanBuffer::from_iter([true, true, false, false]));
241        let b = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false]));
242
243        let and_result = &a & &b;
244        let not_and = !&and_result;
245
246        let not_a = !&a;
247        let not_b = !&b;
248        let or_result = &not_a | &not_b;
249
250        assert_eq!(not_and.len(), 4);
251        assert!(!not_and.value(0)); // !(true & true) = false
252        assert!(not_and.value(1)); // !(true & false) = true
253        assert!(not_and.value(2)); // !(false & true) = true
254        assert!(not_and.value(3)); // !(false & false) = true
255
256        assert_eq!(or_result.len(), 4);
257        assert_eq!(or_result, not_and)
258    }
259
260    #[test]
261    fn test_bitand_associativity() {
262        // Test (A & B) & C = A & (B & C)
263        let a = Mask::from_buffer(BooleanBuffer::from_iter([true, true, false, true]));
264        let b = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, true]));
265        let c = Mask::from_buffer(BooleanBuffer::from_iter([false, true, true, true]));
266
267        let left_assoc = &(&a & &b) & &c;
268        let right_assoc = &a & &(&b & &c);
269
270        assert_eq!(left_assoc.true_count(), right_assoc.true_count());
271        for i in 0..4 {
272            assert_eq!(left_assoc.value(i), right_assoc.value(i));
273        }
274    }
275
276    #[test]
277    fn test_bitand_commutativity() {
278        // Test A & B = B & A
279        let a = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false]));
280        let b = Mask::from_buffer(BooleanBuffer::from_iter([false, true, false, true]));
281
282        let ab = &a & &b;
283        let ba = &b & &a;
284
285        assert_eq!(ab.true_count(), ba.true_count());
286        for i in 0..4 {
287            assert_eq!(ab.value(i), ba.value(i));
288        }
289    }
290
291    #[test]
292    fn test_bitand_identity() {
293        // Test A & AllTrue = A
294        let mask = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false]));
295        let all_true = Mask::new_true(4);
296
297        let result = &mask & &all_true;
298        assert_eq!(result.true_count(), mask.true_count());
299        for i in 0..4 {
300            assert_eq!(result.value(i), mask.value(i));
301        }
302    }
303
304    #[test]
305    fn test_bitand_annihilator() {
306        // Test A & AllFalse = AllFalse
307        let mask = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false]));
308        let all_false = Mask::new_false(4);
309
310        let result = &mask & &all_false;
311        assert!(result.all_false());
312        assert_eq!(result.true_count(), 0);
313    }
314
315    #[test]
316    fn test_bitand_idempotence() {
317        // Test A & A = A
318        let mask = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false, true]));
319        let result = &mask & &mask;
320
321        assert_eq!(result.true_count(), mask.true_count());
322        for i in 0..5 {
323            assert_eq!(result.value(i), mask.value(i));
324        }
325    }
326
327    #[test]
328    fn test_complex_expression() {
329        // Test a more complex expression: (!(!A) | B) & !C
330        let a = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false]));
331        let b = Mask::from_buffer(BooleanBuffer::from_iter([true, true, false, false]));
332        let c = Mask::from_buffer(BooleanBuffer::from_iter([false, true, false, true]));
333
334        let not_not_a = !(&(!&a));
335        let not_not_a_or_b = &not_not_a | &b;
336        let not_c = !&c;
337        let result = &not_not_a_or_b & &not_c;
338
339        // Verify the result manually
340        assert!(result.value(0)); // (!(!true) | true) & !false = (true | true) & true = true
341        assert!(!result.value(1)); // (!(!false) | true) & !true = (false | true) & false = false
342        assert!(result.value(2)); // (!(!true) | false) & !false = (true | false) & true = true
343        assert!(!result.value(3)); // (!(!false) | false) & !true = (false | false) & false = false
344    }
345
346    #[test]
347    fn test_bitor() {
348        // Test basic OR operations
349        let mask1 = Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, false, true]));
350        let mask2 = Mask::from_buffer(BooleanBuffer::from_iter([true, true, false, false, true]));
351
352        let result = &mask1 | &mask2;
353        assert_eq!(result.len(), 5);
354        assert_eq!(result.true_count(), 4);
355        assert!(result.value(0)); // true | true = true
356        assert!(result.value(1)); // false | true = true
357        assert!(result.value(2)); // true | false = true
358        assert!(!result.value(3)); // false | false = false
359        assert!(result.value(4)); // true | true = true
360
361        // Test with AllTrue
362        let all_true = Mask::new_true(5);
363        let result = &mask1 | &all_true;
364        assert!(result.all_true());
365
366        // Test with AllFalse
367        let all_false = Mask::new_false(5);
368        let result = &mask1 | &all_false;
369        assert_eq!(result.true_count(), mask1.true_count());
370    }
371}