steel_core/chunk/paletted_container/
encode.rs

1//! Network encoding of paletted container entries.
2//!
3//! Entries are packed like vanilla `SimpleBitStorage`: `64 / bits` entries per long, low
4//! bits first, no entry spanning two longs, longs written big-endian.
5
6use std::cell::RefCell;
7use std::io::{Error, ErrorKind, Result};
8
9use super::ToGlobalId;
10
11/// Largest packed payload: a 16×16×16 section at 16 bits per entry.
12pub(super) const MAX_PACKED_BYTES: usize = super::BlockPalette::VOLUME.div_ceil(64 / 16) * 8;
13/// Global ids are `u16`-backed, so a table indexed by global id has this many slots.
14const INDEX_TABLE_LEN: usize = 1 << 16;
15/// Marks an index-table slot whose global id is not in the palette being encoded.
16const EMPTY_SLOT: u16 = u16::MAX;
17
18thread_local! {
19    /// Palette index per global id, reused across encodes on this thread. It takes 128 KiB on
20    /// each thread that encodes a section, allocated on first use.
21    ///
22    /// Only the slots of the palette being encoded hold an index; every other slot holds
23    /// [`EMPTY_SLOT`], so a value missing from its palette is detected instead of being
24    /// encoded with an index left over from an earlier palette.
25    static PALETTE_INDEX: RefCell<Box<[u16]>> =
26        RefCell::new(vec![EMPTY_SLOT; INDEX_TABLE_LEN].into_boxed_slice());
27}
28
29/// Packs each value's global id into `buffer` and returns the packed bytes.
30///
31/// # Errors
32/// Returns [`ErrorKind::InvalidData`] if a global id does not fit in `bits`.
33pub(super) fn pack_global_ids<'buffer, V: ToGlobalId + Copy>(
34    values: &[V],
35    bits: u8,
36    buffer: &'buffer mut [u8; MAX_PACKED_BYTES],
37) -> Result<&'buffer [u8]> {
38    // Checked once up front so the packing loop has no per-entry branch.
39    let all_ids = values
40        .iter()
41        .fold(0, |all_ids, value| all_ids | value.to_global_id());
42    if all_ids >> u32::from(bits) != 0 {
43        return Err(invalid_data(
44            "container holds a global id wider than the global palette",
45        ));
46    }
47    pack(values, bits, |value| Some(value.to_global_id()), buffer)
48}
49
50/// Packs each value's index in `palette` into `buffer` and returns the packed bytes.
51///
52/// The thread's index table is released before this returns, so writing the bytes can
53/// encode other containers on the same thread.
54///
55/// # Errors
56/// Returns [`ErrorKind::InvalidData`] if a value is missing from `palette`, which means the
57/// container's palette invariant is broken.
58pub(super) fn pack_palette_indices<'buffer, V: ToGlobalId + Copy>(
59    values: &[V],
60    palette: &[(V, u16)],
61    bits: u8,
62    buffer: &'buffer mut [u8; MAX_PACKED_BYTES],
63) -> Result<&'buffer [u8]> {
64    PALETTE_INDEX.with_borrow_mut(move |table| {
65        let table = PaletteIndexTable::fill(table, palette)?;
66        pack(values, bits, |value| table.index_of(value), buffer)
67    })
68}
69
70/// The thread's index table filled with one palette, reset to empty when dropped (including
71/// on unwind) so the next palette starts from a clean table.
72struct PaletteIndexTable<'a, V: ToGlobalId> {
73    slots: &'a mut [u16],
74    palette: &'a [(V, u16)],
75}
76
77impl<'a, V: ToGlobalId> PaletteIndexTable<'a, V> {
78    fn fill(slots: &'a mut [u16], palette: &'a [(V, u16)]) -> Result<Self> {
79        // Built before filling so an early return still resets the slots already written.
80        let table = Self { slots, palette };
81        for (index, (value, _)) in palette.iter().enumerate() {
82            let index = u16::try_from(index)
83                .ok()
84                .filter(|index| *index != EMPTY_SLOT)
85                .ok_or_else(|| invalid_data("palette has too many entries to index"))?;
86            let slot = table
87                .slots
88                .get_mut(value.to_global_id() as usize)
89                .ok_or_else(|| invalid_data("palette value has an out-of-range global id"))?;
90            debug_assert_eq!(*slot, EMPTY_SLOT, "palette entries must be distinct");
91            *slot = index;
92        }
93        Ok(table)
94    }
95
96    fn index_of(&self, value: V) -> Option<u32> {
97        let index = *self.slots.get(value.to_global_id() as usize)?;
98        (index != EMPTY_SLOT).then_some(u32::from(index))
99    }
100}
101
102impl<V: ToGlobalId> Drop for PaletteIndexTable<'_, V> {
103    fn drop(&mut self) {
104        for (value, _) in self.palette {
105            if let Some(slot) = self.slots.get_mut(value.to_global_id() as usize) {
106                *slot = EMPTY_SLOT;
107            }
108        }
109    }
110}
111
112/// Packs `values` into `buffer`, failing if a value has no entry.
113fn pack<'buffer, V: Copy>(
114    values: &[V],
115    bits: u8,
116    lookup: impl FnMut(V) -> Option<u32>,
117    buffer: &'buffer mut [u8; MAX_PACKED_BYTES],
118) -> Result<&'buffer [u8]> {
119    // A width known at compile time lets each long be filled by a fixed-length loop.
120    macro_rules! pack_with_width {
121        ($($width:literal)*) => {
122            match bits {
123                $($width => pack_longs::<V, $width>(values, lookup, buffer),)*
124                _ => Err(invalid_data("unsupported bits per entry")),
125            }
126        };
127    }
128    pack_with_width!(1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16)
129}
130
131/// Packs `values` into the front of `buffer`, failing if a value has no entry.
132fn pack_longs<'buffer, V: Copy, const BITS: usize>(
133    values: &[V],
134    mut lookup: impl FnMut(V) -> Option<u32>,
135    buffer: &'buffer mut [u8; MAX_PACKED_BYTES],
136) -> Result<&'buffer [u8]> {
137    let values_per_long = 64 / BITS;
138    let len = values.len().div_ceil(values_per_long) * 8;
139    let packed = buffer
140        .get_mut(..len)
141        .ok_or_else(|| invalid_data("container is too large to pack"))?;
142    // Accumulate misses instead of branching per entry; the whole container is rejected.
143    let mut missing = false;
144    for (entries, long_bytes) in values
145        .chunks(values_per_long)
146        .zip(packed.as_chunks_mut::<8>().0)
147    {
148        let mut long = 0u64;
149        for (slot, value) in entries.iter().enumerate() {
150            let id = lookup(*value);
151            missing |= id.is_none();
152            long |= u64::from(id.unwrap_or(0)) << (slot * BITS);
153        }
154        *long_bytes = long.to_be_bytes();
155    }
156    if missing {
157        return Err(invalid_data(
158            "container holds a value its palette cannot encode",
159        ));
160    }
161    Ok(packed)
162}
163
164fn invalid_data(message: &'static str) -> Error {
165    Error::new(ErrorKind::InvalidData, message)
166}
167
168#[cfg(test)]
169mod tests {
170    use super::{MAX_PACKED_BYTES, pack_global_ids};
171
172    /// Packs values at every width from 1 to 16 bits, including 9 to 15, which no section
173    /// uses with vanilla's registries, and reads them back the way vanilla's
174    /// `SimpleBitStorage` does: `64 / bits` values per big-endian long, low bits first.
175    #[test]
176    fn packs_every_width_like_simple_bit_storage() {
177        let mut state = 0x9E37_79B9_7F4A_7C15_u64;
178        for bits in 1..=16u8 {
179            let mask = (1u64 << bits) - 1;
180            let values: Vec<u16> = (0..4096)
181                .map(|_| {
182                    state ^= state << 13;
183                    state ^= state >> 7;
184                    state ^= state << 17;
185                    (state & mask) as u16
186                })
187                .collect();
188            let mut buffer = [0; MAX_PACKED_BYTES];
189            let packed =
190                pack_global_ids(&values, bits, &mut buffer).expect("every value fits the width");
191
192            let values_per_long = 64 / usize::from(bits);
193            let (longs, rest) = packed.as_chunks::<8>();
194            assert!(
195                rest.is_empty(),
196                "{bits} bits: packed bytes should be whole longs"
197            );
198            assert_eq!(
199                longs.len(),
200                values.len().div_ceil(values_per_long),
201                "{bits} bits"
202            );
203            for (index, &value) in values.iter().enumerate() {
204                let long = u64::from_be_bytes(longs[index / values_per_long]);
205                let entry = (long >> (index % values_per_long * usize::from(bits))) & mask;
206                assert_eq!(entry, u64::from(value), "{bits} bits, value {index}");
207            }
208        }
209    }
210}
steel_core/chunk/paletted_container/encode.rs

steel_core/chunk/paletted_container/
encode.rs
