arrow_buffer/util/
bit_mask.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Utils for working with packed bit masks
19
20use crate::bit_util::ceil;
21
22/// Util function to set bits in a slice of bytes.
23///
24/// This will sets all bits on `write_data` in the range `[offset_write..offset_write+len]`
25/// to be equal to the bits in `data` in the range `[offset_read..offset_read+len]`
26/// returns the number of `0` bits `data[offset_read..offset_read+len]`
27/// `offset_write`, `offset_read`, and `len` are in terms of bits
28///
29/// # Panics
30///
31/// Panics if `offset_write + len` exceeds `write_data.len() * 8`, or if
32/// `offset_read + len` exceeds `data.len() * 8`
33pub fn set_bits(
34    write_data: &mut [u8],
35    data: &[u8],
36    offset_write: usize,
37    offset_read: usize,
38    len: usize,
39) -> usize {
40    assert!(
41        offset_write
42            .checked_add(len)
43            .expect("operation will overflow write buffer")
44            <= write_data.len() * 8
45    );
46    assert!(
47        offset_read
48            .checked_add(len)
49            .expect("operation will overflow read buffer")
50            <= data.len() * 8
51    );
52    let mut null_count = 0;
53    let mut acc = 0;
54    while len > acc {
55        // SAFETY: the arguments to `set_upto_64bits` are within the valid range because
56        // (offset_write + acc) + (len - acc) == offset_write + len <= write_data.len() * 8
57        // (offset_read + acc) + (len - acc) == offset_read + len <= data.len() * 8
58        let (n, len_set) = unsafe {
59            set_upto_64bits(
60                write_data,
61                data,
62                offset_write + acc,
63                offset_read + acc,
64                len - acc,
65            )
66        };
67        null_count += n;
68        acc += len_set;
69    }
70
71    null_count
72}
73
74/// Similar to `set_bits` but sets only upto 64 bits, actual number of bits set may vary.
75/// Returns a pair of the number of `0` bits and the number of bits set
76///
77/// # Safety
78/// The caller must ensure all arguments are within the valid range.
79#[inline]
80unsafe fn set_upto_64bits(
81    write_data: &mut [u8],
82    data: &[u8],
83    offset_write: usize,
84    offset_read: usize,
85    len: usize,
86) -> (usize, usize) {
87    let read_byte = offset_read / 8;
88    let read_shift = offset_read % 8;
89    let write_byte = offset_write / 8;
90    let write_shift = offset_write % 8;
91
92    if len >= 64 {
93        let chunk = unsafe { data.as_ptr().add(read_byte).cast::<u64>().read_unaligned() };
94        if read_shift == 0 {
95            if write_shift == 0 {
96                // no shifting necessary
97                let len = 64;
98                let null_count = chunk.count_zeros() as usize;
99                unsafe { write_u64_bytes(write_data, write_byte, chunk) };
100                (null_count, len)
101            } else {
102                // only write shifting necessary
103                let len = 64 - write_shift;
104                let chunk = chunk << write_shift;
105                let null_count = len - chunk.count_ones() as usize;
106                unsafe { or_write_u64_bytes(write_data, write_byte, chunk) };
107                (null_count, len)
108            }
109        } else if write_shift == 0 {
110            // only read shifting necessary
111            let len = 64 - 8; // 56 bits so the next set_upto_64bits call will see write_shift == 0
112            let chunk = (chunk >> read_shift) & 0x00FFFFFFFFFFFFFF; // 56 bits mask
113            let null_count = len - chunk.count_ones() as usize;
114            unsafe { write_u64_bytes(write_data, write_byte, chunk) };
115            (null_count, len)
116        } else {
117            let len = 64 - std::cmp::max(read_shift, write_shift);
118            let chunk = (chunk >> read_shift) << write_shift;
119            let null_count = len - chunk.count_ones() as usize;
120            unsafe { or_write_u64_bytes(write_data, write_byte, chunk) };
121            (null_count, len)
122        }
123    } else if len == 1 {
124        let byte_chunk = (unsafe { data.get_unchecked(read_byte) } >> read_shift) & 1;
125        unsafe { *write_data.get_unchecked_mut(write_byte) |= byte_chunk << write_shift };
126        ((byte_chunk ^ 1) as usize, 1)
127    } else {
128        let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
129        let bytes = ceil(len + read_shift, 8);
130        // SAFETY: the args of `read_bytes_to_u64` are valid as read_byte + bytes <= data.len()
131        let chunk = unsafe { read_bytes_to_u64(data, read_byte, bytes) };
132        let mask = u64::MAX >> (64 - len);
133        let chunk = (chunk >> read_shift) & mask; // masking to read `len` bits only
134        let chunk = chunk << write_shift; // shifting back to align with `write_data`
135        let null_count = len - chunk.count_ones() as usize;
136        let bytes = ceil(len + write_shift, 8);
137        for (i, c) in chunk.to_le_bytes().iter().enumerate().take(bytes) {
138            unsafe { *write_data.get_unchecked_mut(write_byte + i) |= c };
139        }
140        (null_count, len)
141    }
142}
143
144/// # Safety
145/// The caller must ensure `data` has `offset..(offset + 8)` range, and `count <= 8`.
146#[inline]
147unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
148    debug_assert!(count <= 8);
149    let mut tmp: u64 = 0;
150    let src = unsafe { data.as_ptr().add(offset) };
151    unsafe { std::ptr::copy_nonoverlapping(src, std::ptr::from_mut(&mut tmp).cast::<u8>(), count) };
152    tmp
153}
154
155/// # Safety
156/// The caller must ensure `data` has `offset..(offset + 8)` range
157#[inline]
158unsafe fn write_u64_bytes(data: &mut [u8], offset: usize, chunk: u64) {
159    let ptr = unsafe { data.as_mut_ptr().add(offset) }.cast::<u64>();
160    unsafe { ptr.write_unaligned(chunk) };
161}
162
163/// Similar to `write_u64_bytes`, but this method ORs the offset addressed `data` and `chunk`
164/// instead of overwriting
165///
166/// # Safety
167/// The caller must ensure `data` has `offset..(offset + 8)` range
168#[inline]
169unsafe fn or_write_u64_bytes(data: &mut [u8], offset: usize, chunk: u64) {
170    let ptr = unsafe { data.as_mut_ptr().add(offset) };
171    let chunk = chunk | (unsafe { *ptr }) as u64;
172    unsafe { ptr.cast::<u64>().write_unaligned(chunk) };
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178    use crate::bit_util::{get_bit, set_bit, unset_bit};
179    use rand::prelude::StdRng;
180    use rand::{RngExt, SeedableRng, TryRng};
181    use std::fmt::Display;
182
183    #[test]
184    fn test_set_bits_aligned() {
185        SetBitsTest {
186            write_data: vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
187            data: vec![
188                0b11100111, 0b10100101, 0b10011001, 0b11011011, 0b11101011, 0b11000011, 0b11100111,
189                0b10100101,
190            ],
191            offset_write: 8,
192            offset_read: 0,
193            len: 64,
194            expected_data: vec![
195                0, 0b11100111, 0b10100101, 0b10011001, 0b11011011, 0b11101011, 0b11000011,
196                0b11100111, 0b10100101, 0,
197            ],
198            expected_null_count: 24,
199        }
200        .verify();
201    }
202
203    #[test]
204    fn test_set_bits_unaligned_destination_start() {
205        SetBitsTest {
206            write_data: vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
207            data: vec![
208                0b11100111, 0b10100101, 0b10011001, 0b11011011, 0b11101011, 0b11000011, 0b11100111,
209                0b10100101,
210            ],
211            offset_write: 3,
212            offset_read: 0,
213            len: 64,
214            expected_data: vec![
215                0b00111000, 0b00101111, 0b11001101, 0b11011100, 0b01011110, 0b00011111, 0b00111110,
216                0b00101111, 0b00000101, 0b00000000,
217            ],
218            expected_null_count: 24,
219        }
220        .verify();
221    }
222
223    #[test]
224    fn test_set_bits_unaligned_destination_end() {
225        SetBitsTest {
226            write_data: vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
227            data: vec![
228                0b11100111, 0b10100101, 0b10011001, 0b11011011, 0b11101011, 0b11000011, 0b11100111,
229                0b10100101,
230            ],
231            offset_write: 8,
232            offset_read: 0,
233            len: 62,
234            expected_data: vec![
235                0, 0b11100111, 0b10100101, 0b10011001, 0b11011011, 0b11101011, 0b11000011,
236                0b11100111, 0b00100101, 0,
237            ],
238            expected_null_count: 23,
239        }
240        .verify();
241    }
242
243    #[test]
244    fn test_set_bits_unaligned() {
245        SetBitsTest {
246            write_data: vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
247            data: vec![
248                0b11100111, 0b10100101, 0b10011001, 0b11011011, 0b11101011, 0b11000011, 0b11100111,
249                0b10100101, 0b10011001, 0b11011011, 0b11101011, 0b11000011, 0b11100111, 0b10100101,
250                0b10011001, 0b11011011, 0b11101011, 0b11000011,
251            ],
252            offset_write: 3,
253            offset_read: 5,
254            len: 95,
255            expected_data: vec![
256                0b01111000, 0b01101001, 0b11100110, 0b11110110, 0b11111010, 0b11110000, 0b01111001,
257                0b01101001, 0b11100110, 0b11110110, 0b11111010, 0b11110000, 0b00000001,
258            ],
259            expected_null_count: 35,
260        }
261        .verify();
262    }
263
264    #[test]
265    fn set_bits_fuzz() {
266        let mut rng = StdRng::seed_from_u64(42);
267        let mut data = SetBitsTest::new();
268        for _ in 0..100 {
269            data.regen(&mut rng);
270            data.verify();
271        }
272    }
273
274    #[derive(Debug, Default)]
275    struct SetBitsTest {
276        /// target write data
277        write_data: Vec<u8>,
278        /// source data
279        data: Vec<u8>,
280        offset_write: usize,
281        offset_read: usize,
282        len: usize,
283        /// the expected contents of write_data after the test
284        expected_data: Vec<u8>,
285        /// the expected number of nulls copied at the end of the test
286        expected_null_count: usize,
287    }
288
289    /// prints a byte slice as a binary string like "01010101 10101010"
290    struct BinaryFormatter<'a>(&'a [u8]);
291    impl Display for BinaryFormatter<'_> {
292        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
293            for byte in self.0 {
294                write!(f, "{byte:08b} ")?;
295            }
296            write!(f, " ")?;
297            Ok(())
298        }
299    }
300
301    impl Display for SetBitsTest {
302        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
303            writeln!(f, "SetBitsTest {{")?;
304            writeln!(f, "  write_data:    {}", BinaryFormatter(&self.write_data))?;
305            writeln!(f, "  data:          {}", BinaryFormatter(&self.data))?;
306            writeln!(
307                f,
308                "  expected_data: {}",
309                BinaryFormatter(&self.expected_data)
310            )?;
311            writeln!(f, "  offset_write: {}", self.offset_write)?;
312            writeln!(f, "  offset_read: {}", self.offset_read)?;
313            writeln!(f, "  len: {}", self.len)?;
314            writeln!(f, "  expected_null_count: {}", self.expected_null_count)?;
315            writeln!(f, "}}")
316        }
317    }
318
319    impl SetBitsTest {
320        /// create a new instance of FuzzData
321        fn new() -> Self {
322            Self::default()
323        }
324
325        /// Update this instance's fields with randomly selected values and expected data
326        fn regen(&mut self, rng: &mut StdRng) {
327            //  (read) data
328            // ------------------+-----------------+-------
329            // .. offset_read .. | data            | ...
330            // ------------------+-----------------+-------
331
332            // Write data
333            // -------------------+-----------------+-------
334            // .. offset_write .. | (data to write) | ...
335            // -------------------+-----------------+-------
336
337            // length of data to copy
338            let len = rng.random_range(0..=200);
339
340            // randomly pick where we will write to
341            let offset_write_bits = rng.random_range(0..=200);
342            let offset_write_bytes = if offset_write_bits % 8 == 0 {
343                offset_write_bits / 8
344            } else {
345                (offset_write_bits / 8) + 1
346            };
347            let extra_write_data_bytes = rng.random_range(0..=5); // ensure 0 shows up often
348
349            // randomly decide where we will read from
350            let extra_read_data_bytes = rng.random_range(0..=5); // make sure 0 shows up often
351            let offset_read_bits = rng.random_range(0..=200);
352            let offset_read_bytes = if offset_read_bits % 8 != 0 {
353                (offset_read_bits / 8) + 1
354            } else {
355                offset_read_bits / 8
356            };
357
358            // create space for writing
359            self.write_data.clear();
360            self.write_data
361                .resize(offset_write_bytes + len + extra_write_data_bytes, 0);
362
363            // interestingly set_bits seems to assume the output is already zeroed
364            // the fuzz tests fail when this is uncommented
365            //self.write_data.try_fill(rng).unwrap();
366            self.offset_write = offset_write_bits;
367
368            // make source data
369            self.data
370                .resize(offset_read_bytes + len + extra_read_data_bytes, 0);
371            // fill source data with random bytes
372            rng.try_fill_bytes(self.data.as_mut_slice()).unwrap();
373            self.offset_read = offset_read_bits;
374
375            self.len = len;
376
377            // generated expectated output (not efficient)
378            self.expected_data.resize(self.write_data.len(), 0);
379            self.expected_data.copy_from_slice(&self.write_data);
380
381            self.expected_null_count = 0;
382            for i in 0..self.len {
383                let bit = get_bit(&self.data, self.offset_read + i);
384                if bit {
385                    set_bit(&mut self.expected_data, self.offset_write + i);
386                } else {
387                    unset_bit(&mut self.expected_data, self.offset_write + i);
388                    self.expected_null_count += 1;
389                }
390            }
391        }
392
393        /// call set_bits with the given parameters and compare with the expected output
394        fn verify(&self) {
395            // call set_bits and compare
396            let mut actual = self.write_data.clone();
397            let null_count = set_bits(
398                &mut actual,
399                &self.data,
400                self.offset_write,
401                self.offset_read,
402                self.len,
403            );
404
405            assert_eq!(actual, self.expected_data, "self: {self}");
406            assert_eq!(null_count, self.expected_null_count, "self: {self}");
407        }
408    }
409
410    #[test]
411    fn test_set_upto_64bits() {
412        // len >= 64
413        let write_data: &mut [u8] = &mut [0; 9];
414        let data: &[u8] = &[
415            0b00000001, 0b00000001, 0b00000001, 0b00000001, 0b00000001, 0b00000001, 0b00000001,
416            0b00000001, 0b00000001,
417        ];
418        let offset_write = 1;
419        let offset_read = 0;
420        let len = 65;
421        let (n, len_set) =
422            unsafe { set_upto_64bits(write_data, data, offset_write, offset_read, len) };
423        assert_eq!(n, 55);
424        assert_eq!(len_set, 63);
425        assert_eq!(
426            write_data,
427            &[
428                0b00000010, 0b00000010, 0b00000010, 0b00000010, 0b00000010, 0b00000010, 0b00000010,
429                0b00000010, 0b00000000
430            ]
431        );
432
433        // len = 1
434        let write_data: &mut [u8] = &mut [0b00000000];
435        let data: &[u8] = &[0b00000001];
436        let offset_write = 1;
437        let offset_read = 0;
438        let len = 1;
439        let (n, len_set) =
440            unsafe { set_upto_64bits(write_data, data, offset_write, offset_read, len) };
441        assert_eq!(n, 0);
442        assert_eq!(len_set, 1);
443        assert_eq!(write_data, &[0b00000010]);
444    }
445
446    #[test]
447    #[should_panic(expected = "operation will overflow read buffer")]
448    fn test_overflow_read_buffer_bounds() {
449        // Tiny buffers so any huge computed index is out-of-bounds.
450        let data = [0u8; 1];
451        let mut write_data = [0u8; 1];
452
453        // Choose values so (offset_read + len) wraps to a small number in release builds.
454        // offset_read = usize::MAX - 7, len = 8 => wraps to 0.
455        // This can bypass `assert!(offset_read + len <= data.len() * 8)`.
456        let offset_write: usize = 0;
457        let offset_read: usize = usize::MAX - 7;
458        let len: usize = 8;
459
460        // should panic on bounds check overflow
461        let _nulls = set_bits(&mut write_data, &data, offset_write, offset_read, len);
462    }
463
464    #[test]
465    #[should_panic(expected = "operation will overflow write buffer")]
466    fn test_overflow_write_buffer_bounds() {
467        // Tiny buffers so any huge computed index is out-of-bounds.
468        let data = [0u8; 1];
469        let mut write_data = [0u8; 1];
470
471        // Choose values so (offset_write + len) wraps to a small number in release builds.
472        // offset_write = usize::MAX - 7, len = 8 => wraps to 0.
473        // This can bypass `assert!(offset_write + len <= write_data.len() * 8)`.
474        let offset_write: usize = usize::MAX - 7;
475        let offset_read: usize = 0;
476        let len: usize = 8;
477
478        // should panic on bounds check overflow
479        let _nulls = set_bits(&mut write_data, &data, offset_write, offset_read, len);
480    }
481}
arrow_buffer/util/bit_mask.rs

arrow_buffer/util/
bit_mask.rs
