Skip to main content

steel_core/chunk_saver/
bit_pack.rs

1//! Bit packing utilities for chunk persistence.
2//!
3//! Uses only power-of-2 bit widths (1, 2, 4, 8, 16) to avoid entries spanning
4//! u64 boundaries, which simplifies encoding/decoding and improves performance.
5
6/// Calculates the number of bits needed to represent indices into a palette.
7/// Only returns power-of-2 values: 1, 2, 4, 8, or 16.
8///
9/// Returns `None` for homogeneous containers (palette length 0 or 1).
10#[must_use]
11pub const fn bits_for_palette_len(palette_len: usize) -> Option<u8> {
12    match palette_len {
13        0..=1 => None, // Homogeneous, no bit array needed
14        2 => Some(1),
15        3..=4 => Some(2),
16        5..=16 => Some(4),
17        17..=256 => Some(8),
18        _ => Some(16),
19    }
20}
21
22/// Packs indices from an iterator into a compact bit array using power-of-2
23/// bit widths, without an intermediate buffer.
24///
25/// The entry count is taken from the iterator's `ExactSizeIterator` length, so
26/// the caller cannot desync the packed width from the packed data.
27///
28/// # Arguments
29/// * `indices` - The indices to pack (values must fit in `bits` bits)
30/// * `bits` - Bits per entry (must be 1, 2, 4, 8, or 16)
31///
32/// # Panics
33/// Panics if `bits` is not a power of 2 or is greater than 16.
34#[must_use]
35pub fn pack_indices_from_iter(indices: impl ExactSizeIterator<Item = u32>, bits: u8) -> Box<[u64]> {
36    debug_assert!(
37        bits.is_power_of_two() && bits <= 16,
38        "bits must be 1, 2, 4, 8, or 16"
39    );
40    let entry_count = indices.len();
41    if entry_count == 0 {
42        return Box::new([]);
43    }
44    let bits = bits as usize;
45    let values_per_u64 = 64 / bits;
46    let num_u64s = entry_count.div_ceil(values_per_u64);
47    let mut data = vec![0u64; num_u64s];
48    // Place each index in its little-endian slot. Supported widths divide 64,
49    // so entries stay within one word and OR preserves the other slots.
50    for (i, index) in indices.enumerate() {
51        data[i / values_per_u64] |= u64::from(index) << ((i % values_per_u64) * bits);
52    }
53    data.into_boxed_slice()
54}
55
56/// Unpacks indices from a compact bit array.
57///
58/// # Arguments
59/// * `data` - The packed bit array
60/// * `bits` - Bits per entry (must be 1, 2, 4, 8, or 16)
61///
62/// # Panics
63/// Panics if `bits` is not a power of 2 or is greater than 16.
64#[inline]
65pub fn unpack_indices(data: &[u64], bits: u8) -> impl Iterator<Item = u32> {
66    let bits = bits as usize;
67    let mask = (1u64 << bits) - 1;
68    let values_per_u64 = 64 / bits;
69    data.iter().flat_map(move |&word| {
70        (0..values_per_u64).map(move |j| ((word >> (j * bits)) & mask) as u32)
71    })
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[test]
79    fn test_bits_for_palette_len() {
80        assert_eq!(bits_for_palette_len(0), None);
81        assert_eq!(bits_for_palette_len(1), None);
82        assert_eq!(bits_for_palette_len(2), Some(1));
83        assert_eq!(bits_for_palette_len(3), Some(2));
84        assert_eq!(bits_for_palette_len(4), Some(2));
85        assert_eq!(bits_for_palette_len(5), Some(4));
86        assert_eq!(bits_for_palette_len(16), Some(4));
87        assert_eq!(bits_for_palette_len(17), Some(8));
88        assert_eq!(bits_for_palette_len(256), Some(8));
89        assert_eq!(bits_for_palette_len(257), Some(16));
90    }
91
92    #[test]
93    fn test_pack_unpack_roundtrip() {
94        for bits in [1, 2, 4, 8, 16] {
95            let max_value = (1u32 << bits) - 1;
96            let indices: Vec<u32> = (0..100).map(|i| i % (max_value + 1)).collect();
97
98            let packed = pack_indices_from_iter(indices.iter().copied(), bits);
99            let unpacked: Vec<u32> = unpack_indices(&packed, bits).take(indices.len()).collect();
100
101            assert_eq!(indices, unpacked, "Failed for bits={bits}");
102        }
103    }
104
105    #[test]
106    fn test_pack_4096_entries() {
107        // Simulate a chunk section with 4096 blocks
108        let indices: Vec<u32> = (0..4096).map(|i| (i % 16) as u32).collect();
109
110        let packed = pack_indices_from_iter(indices.iter().copied(), 4);
111        assert_eq!(packed.len(), 4096 / 16); // 16 values per u64 with 4 bits
112
113        let unpacked: Vec<u32> = unpack_indices(&packed, 4).take(4096).collect();
114        assert_eq!(indices, unpacked);
115    }
116}