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    // Matches Java's `BitSet.valueOf(byte[])`: little-endian bytes packed into words.
31    fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
32        let bytes = Vec::<u8>::read_prefixed::<VarInt>(data)?;
33        let mut words = vec![0_u64; bytes.len().div_ceil(8)];
34        for (index, byte) in bytes.iter().enumerate() {
35            words[index / 8] |= u64::from(*byte) << ((index % 8) * 8);
36        }
37        Ok(Self(words.into_boxed_slice()))
38    }
39}
40
41impl WriteTo for BitSet {
42    // Matches Java's `BitSet.toByteArray()`: little-endian bytes per word, trimmed
43    // to the last non-zero byte (not word) — vanilla's `ByteBufCodecs.BIT_SET` wire
44    // format is a VarInt byte-length prefix followed by those bytes.
45    fn write(&self, writer: &mut impl Write) -> Result<()> {
46        let mut bytes = Vec::with_capacity(self.0.len() * 8);
47        for word in &self.0 {
48            bytes.extend_from_slice(&word.to_le_bytes());
49        }
50        let trimmed_len = bytes.iter().rposition(|&b| b != 0).map_or(0, |i| i + 1);
51        bytes.truncate(trimmed_len);
52        bytes.write_prefixed::<VarInt>(writer)
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use crate::serial::WriteTo;
59
60    use super::BitSet;
61
62    #[test]
63    fn write_trims_empty_bit_set_to_zero_bytes() {
64        let bit_set = BitSet(vec![0].into_boxed_slice());
65        let mut data = Vec::new();
66
67        bit_set.write(&mut data).expect("bit set should encode");
68
69        assert_eq!(data, vec![0]);
70    }
71
72    #[test]
73    fn write_trims_to_last_non_zero_byte_little_endian() {
74        let bit_set = BitSet(vec![5, 0, 7, 0, 0].into_boxed_slice());
75        let mut data = Vec::new();
76
77        bit_set.write(&mut data).expect("bit set should encode");
78
79        // 3 non-empty words -> 24 little-endian bytes, trimmed to the last
80        // non-zero byte: word 0 = 5 (byte 0), word 2 = 7 (byte 16).
81        let expected = vec![17_u8, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7];
82        assert_eq!(data, expected);
83    }
84
85    #[test]
86    fn round_trips_through_read() {
87        use std::io::Cursor;
88
89        use crate::serial::ReadFrom;
90
91        let bit_set =
92            BitSet(vec![0x0102_0304_0506_0708, 0, 0x0000_0000_0000_00ff].into_boxed_slice());
93        let mut data = Vec::new();
94        bit_set.write(&mut data).expect("bit set should encode");
95
96        let mut cursor = Cursor::new(data.as_slice());
97        let read_back = BitSet::read(&mut cursor).expect("bit set should decode");
98
99        assert_eq!(read_back.0.as_ref(), bit_set.0.as_ref());
100    }
101}
steel_utils/codec/bit_set.rs

steel_utils/codec/
bit_set.rs
