steel_utils/codec/
bit_set.rs

1use std::io::{Cursor, Result, Write};
2
3use crate::serial::{PrefixedRead, PrefixedWrite, ReadFrom, WriteTo};
4
5use super::VarInt;
6
7/// A simple bit set implementation.
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub struct BitSet(pub Box<[u64]>);
10
11impl BitSet {
12    /// Sets the bit at the given index.
13    pub fn set(&mut self, index: usize, value: bool) {
14        let u64_index = index / 64;
15        let bit_index = index % 64;
16
17        if u64_index >= self.0.len() {
18            return;
19        }
20
21        if value {
22            self.0[u64_index] |= 1 << bit_index;
23        } else {
24            self.0[u64_index] &= !(1 << bit_index);
25        }
26    }
27}
28
29impl ReadFrom for BitSet {
30    fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
31        Ok(Self(Vec::read_prefixed::<VarInt>(data)?.into_boxed_slice()))
32    }
33}
34
35impl WriteTo for BitSet {
36    fn write(&self, writer: &mut impl Write) -> Result<()> {
37        let trimmed_len = self
38            .0
39            .iter()
40            .rposition(|word| *word != 0)
41            .map_or(0, |index| index + 1);
42        self.0[..trimmed_len].write_prefixed::<VarInt>(writer)
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use crate::serial::WriteTo;
49
50    use super::BitSet;
51
52    #[test]
53    fn write_trims_empty_bit_set_to_zero_longs() {
54        let bit_set = BitSet(vec![0].into_boxed_slice());
55        let mut data = Vec::new();
56
57        bit_set.write(&mut data).expect("bit set should encode");
58
59        assert_eq!(data, vec![0]);
60    }
61
62    #[test]
63    fn write_trims_only_trailing_zero_longs() {
64        let bit_set = BitSet(vec![5, 0, 7, 0, 0].into_boxed_slice());
65        let mut data = Vec::new();
66
67        bit_set.write(&mut data).expect("bit set should encode");
68
69        let mut expected = vec![3];
70        expected.extend_from_slice(&5_u64.to_be_bytes());
71        expected.extend_from_slice(&0_u64.to_be_bytes());
72        expected.extend_from_slice(&7_u64.to_be_bytes());
73        assert_eq!(data, expected);
74    }
75}
steel_utils/codec/bit_set.rs

steel_utils/codec/
bit_set.rs
