steel_core/chunk/
paletted_container.rs

1//! A paletted container is a container that can be either homogeneous or heterogeneous.
2use std::{
3    fmt::Debug,
4    hash::Hash,
5    io::{Result, Write},
6    mem, slice,
7    sync::OnceLock,
8};
9
10use steel_registry::blocks::block_state_ext::BlockStateExt;
11use steel_registry::{REGISTRY, RegistryExt as _};
12use steel_utils::{BlockStateId, codec::VarInt, mth::ceil_log2, serial::WriteTo};
13
14/// A trait for converting a value to a global ID.
15pub trait ToGlobalId {
16    /// Converts the value to a global ID.
17    fn to_global_id(&self) -> u32;
18}
19
20impl ToGlobalId for BlockStateId {
21    fn to_global_id(&self) -> u32 {
22        u32::from(self.0)
23    }
24}
25
26impl ToGlobalId for u16 {
27    fn to_global_id(&self) -> u32 {
28        u32::from(*self)
29    }
30}
31
32/// 3d array indexed by y,z,x
33type Cube<T, const DIM: usize> = [[[T; DIM]; DIM]; DIM];
34
35/// A heterogeneous palette container.
36#[derive(Debug, Clone)]
37pub struct HeterogeneousPalette<V: Hash + Eq + Copy, const DIM: usize> {
38    pub(crate) cube: Box<Cube<V, DIM>>,
39    // Keeps track of how many different times each value appears in the cube. (value, count)
40    pub(crate) palette: Vec<(V, u16)>,
41}
42
43impl<V: Hash + Eq + Copy, const DIM: usize> HeterogeneousPalette<V, DIM> {
44    fn get(&self, x: usize, y: usize, z: usize) -> V {
45        debug_assert!(x < DIM);
46        debug_assert!(y < DIM);
47        debug_assert!(z < DIM);
48
49        self.cube[y][z][x]
50    }
51
52    fn get_at_index(&self, index: usize) -> V {
53        debug_assert!(index < DIM * DIM * DIM);
54
55        let y = index / (DIM * DIM);
56        let z = (index / DIM) % DIM;
57        let x = index % DIM;
58        self.cube[y][z][x]
59    }
60
61    /// Returns an iterator over all values in the cube in y, z, x order.
62    pub fn iter_values(&self) -> impl Iterator<Item = &V> {
63        self.cube.iter().flatten().flatten()
64    }
65
66    fn set(&mut self, x: usize, y: usize, z: usize, value: V) -> V {
67        debug_assert!(x < DIM);
68        debug_assert!(y < DIM);
69        debug_assert!(z < DIM);
70
71        let old_value = self.cube[y][z][x];
72
73        if let Some((_, count)) = self.palette.iter_mut().find(|(v, _)| *v == value) {
74            *count += 1;
75        } else {
76            self.palette.push((value, 1));
77        }
78
79        if let Some((index, (_, count))) = self
80            .palette
81            .iter_mut()
82            .enumerate()
83            .find(|(_, (v, _))| *v == old_value)
84        {
85            *count -= 1;
86            if *count == 0 {
87                self.palette.swap_remove(index);
88            }
89        }
90
91        self.cube[y][z][x] = value;
92
93        old_value
94    }
95}
96
97/// A paletted container.
98///
99/// `Building` is a transient mode used during worldgen: it's a raw cube without
100/// palette tracking, so writes are O(1) stores. Must be finalized via
101/// [`Self::finalize_building`] (or implicitly when the parent section recalculates
102/// its counters) before any serialization or paletted access. Mirrors
103/// `FastNoise`'s `FastChunkSection` write-only fill mode.
104#[derive(Debug, Clone)]
105pub enum PalettedContainer<V: Hash + Eq + Copy + Default, const DIM: usize> {
106    /// A homogeneous container, where all values are the same.
107    Homogeneous(V),
108    /// A heterogeneous container, where values can be different.
109    Heterogeneous(HeterogeneousPalette<V, DIM>),
110    /// Write-only build mode: raw cube without palette tracking.
111    /// `set` is a single store; `get` is a direct read.
112    /// Convert back via [`Self::finalize_building`].
113    Building(Box<Cube<V, DIM>>),
114}
115
116enum PaletteMode {
117    Linear,
118    Hash,
119    Global,
120}
121
122impl<V: Hash + Eq + Copy + Default + Debug, const DIM: usize> PalettedContainer<V, DIM> {
123    /// The size of the container in one dimension.
124    pub const SIZE: usize = DIM;
125    /// The volume of the container.
126    pub const VOLUME: usize = DIM * DIM * DIM;
127
128    /// Creates a `PalettedContainer` from a pre-built cube.
129    ///
130    /// Will automatically determine if the result should be homogeneous or heterogeneous.
131    ///
132    /// Walks the cube as a flat slice (it's `[[[V; DIM]; DIM]; DIM]` so memory
133    /// is contiguous) and counts identical cells in runs. The inner "find run
134    /// end" loop is a vectorizable equality scan, so long stone columns collapse
135    /// to a single palette increment.
136    #[must_use]
137    pub fn from_cube(cube: Box<Cube<V, DIM>>) -> Self {
138        let mut palette: Vec<(V, u16)> = Vec::new();
139        let total = DIM * DIM * DIM;
140        // SAFETY: `[[[V; DIM]; DIM]; DIM]` is a fully-contiguous array of
141        // `DIM*DIM*DIM` `V`s, so casting its base pointer to `*const V` and
142        // building a slice of that length is sound. `cube` is a live `Box`, so
143        // the pointer is valid for the lifetime of the slice.
144        let flat: &[V] = unsafe { slice::from_raw_parts(cube.as_ptr().cast::<V>(), total) };
145
146        let mut i = 0;
147        while i < total {
148            let v = flat[i];
149            let mut j = i + 1;
150            while j < total && flat[j] == v {
151                j += 1;
152            }
153            let run_len = (j - i) as u16;
154            if let Some(pos) = palette.iter().position(|(value, _)| *value == v) {
155                palette[pos].1 += run_len;
156            } else {
157                palette.push((v, run_len));
158            }
159            i = j;
160        }
161
162        if palette.len() == 1 {
163            Self::Homogeneous(palette[0].0)
164        } else {
165            Self::Heterogeneous(HeterogeneousPalette { cube, palette })
166        }
167    }
168
169    /// Gets the value at the given coordinates.
170    pub fn get(&self, x: usize, y: usize, z: usize) -> V {
171        match self {
172            Self::Homogeneous(value) => *value,
173            Self::Heterogeneous(data) => data.get(x, y, z),
174            Self::Building(cube) => {
175                debug_assert!(x < DIM);
176                debug_assert!(y < DIM);
177                debug_assert!(z < DIM);
178                cube[y][z][x]
179            }
180        }
181    }
182
183    /// Gets the value at a y,z,x linear index.
184    ///
185    /// The index layout is `x + z * DIM + y * DIM * DIM`, matching the flat
186    /// order used when serializing palette data and by `ScalableLux` light propagation.
187    pub fn get_at_index(&self, index: usize) -> V {
188        debug_assert!(index < Self::VOLUME);
189
190        match self {
191            Self::Homogeneous(value) => *value,
192            Self::Heterogeneous(data) => data.get_at_index(index),
193            Self::Building(cube) => {
194                let y = index / (DIM * DIM);
195                let z = (index / DIM) % DIM;
196                let x = index % DIM;
197                cube[y][z][x]
198            }
199        }
200    }
201
202    /// Copies the full vertical column at `(x, z)` into `out`.
203    pub(crate) fn copy_column_into(&self, x: usize, z: usize, out: &mut [V]) {
204        debug_assert!(x < DIM);
205        debug_assert!(z < DIM);
206        debug_assert!(out.len() >= DIM);
207
208        match self {
209            Self::Homogeneous(value) => {
210                for slot in &mut out[..DIM] {
211                    *slot = *value;
212                }
213            }
214            Self::Heterogeneous(data) => {
215                for (y, slot) in out[..DIM].iter_mut().enumerate() {
216                    *slot = data.cube[y][z][x];
217                }
218            }
219            Self::Building(cube) => {
220                for (y, slot) in out[..DIM].iter_mut().enumerate() {
221                    *slot = cube[y][z][x];
222                }
223            }
224        }
225    }
226
227    /// Returns whether this container's palette may contain a matching value.
228    ///
229    /// This checks palette entries instead of every cell, matching vanilla and
230    /// `ScalableLux`'s fast pre-scan before doing a full section pass.
231    #[must_use]
232    pub fn maybe_has(&self, mut predicate: impl FnMut(V) -> bool) -> bool {
233        match self {
234            Self::Homogeneous(value) => predicate(*value),
235            Self::Heterogeneous(data) => data.palette.iter().any(|(value, _)| predicate(*value)),
236            Self::Building(cube) => cube
237                .iter()
238                .flatten()
239                .flatten()
240                .any(|value| predicate(*value)),
241        }
242    }
243
244    /// Collects all values in the container in y, z, x order.
245    #[must_use]
246    pub fn collect_values(&self) -> Vec<V> {
247        match self {
248            Self::Homogeneous(value) => vec![*value; Self::VOLUME],
249            Self::Heterogeneous(data) => data.iter_values().copied().collect(),
250            Self::Building(cube) => cube.iter().flatten().flatten().copied().collect(),
251        }
252    }
253
254    /// Switches the container into write-only build mode for fast bulk writes.
255    /// Idempotent: a no-op if already in [`Self::Building`].
256    ///
257    /// Allocates a `Cube` if currently `Homogeneous`. For `Heterogeneous` it
258    /// reuses the existing cube allocation.
259    pub fn enter_building_mode(&mut self) {
260        match self {
261            Self::Building(_) => {}
262            Self::Homogeneous(value) => {
263                let cube: Box<Cube<V, DIM>> = Box::new([[[*value; DIM]; DIM]; DIM]);
264                *self = Self::Building(cube);
265            }
266            Self::Heterogeneous(_) => {
267                let taken = mem::replace(self, Self::Homogeneous(V::default()));
268                let Self::Heterogeneous(data) = taken else {
269                    unreachable!()
270                };
271                *self = Self::Building(data.cube);
272            }
273        }
274    }
275
276    /// Returns the raw cube backing this container as a flat mutable slice if
277    /// it's currently in [`Self::Building`] mode. Indexing is `[y * DIM*DIM + z * DIM + x]`.
278    ///
279    /// Used by the chunk fill path to bypass the 3-arm `set` dispatch and the
280    /// unused old-value load — write-only worldgen never reads back what it
281    /// just wrote, so the read in `set` is wasted memory traffic.
282    #[inline]
283    pub fn as_building_slice_mut(&mut self) -> Option<&mut [V]> {
284        if let Self::Building(cube) = self {
285            // SAFETY: `[[[V; DIM]; DIM]; DIM]` is a contiguous array of
286            // `DIM*DIM*DIM` `V`s; the cast preserves the live `&mut Box`'s
287            // borrow because the returned slice cannot outlive `self`.
288            Some(unsafe {
289                slice::from_raw_parts_mut(cube.as_mut_ptr().cast::<V>(), DIM * DIM * DIM)
290            })
291        } else {
292            None
293        }
294    }
295
296    /// Finalizes a [`Self::Building`] container back to `Homogeneous` or
297    /// `Heterogeneous` by scanning the cube once and constructing the palette.
298    /// No-op if not in build mode.
299    pub fn finalize_building(&mut self) {
300        if !matches!(self, Self::Building(_)) {
301            return;
302        }
303        let taken = mem::replace(self, Self::Homogeneous(V::default()));
304        let Self::Building(cube) = taken else {
305            unreachable!()
306        };
307        *self = Self::from_cube(cube);
308    }
309
310    /// Sets the value at the given coordinates.
311    pub fn set(&mut self, x: usize, y: usize, z: usize, value: V) -> V {
312        debug_assert!(x < Self::SIZE);
313        debug_assert!(y < Self::SIZE);
314        debug_assert!(z < Self::SIZE);
315
316        match self {
317            Self::Homogeneous(original) => {
318                let original = *original;
319                if value != original {
320                    let mut cube = Box::new([[[original; DIM]; DIM]; DIM]);
321                    cube[y][z][x] = value;
322                    *self = Self::from_cube(cube);
323                }
324                original
325            }
326            Self::Heterogeneous(data) => {
327                let original = data.set(x, y, z, value);
328                if data.palette.len() == 1 {
329                    *self = Self::Homogeneous(data.palette[0].0);
330                }
331                original
332            }
333            Self::Building(cube) => {
334                let old = cube[y][z][x];
335                cube[y][z][x] = value;
336                old
337            }
338        }
339    }
340
341    /// Writes the container to the given writer.
342    ///
343    /// # Errors
344    /// - If the writer fails to write.
345    #[expect(
346        clippy::missing_panics_doc,
347        clippy::unwrap_used,
348        reason = "position() is guaranteed to exist: palette was built from the cube's own values"
349    )]
350    pub fn write(&self, writer: &mut impl Write) -> Result<()>
351    where
352        V: ToGlobalId,
353    {
354        match self {
355            Self::Homogeneous(value) => {
356                // bits per entry = 0 (ZeroBitStorage)
357                0u8.write(writer)?;
358                // Single-value palette
359                VarInt(value.to_global_id() as i32).write(writer)?;
360                // writeFixedSizeLongArray(new long[0]) writes nothing
361            }
362            Self::Heterogeneous(data) => {
363                let (bits, mode) = Self::calculate_strategy(data.palette.len());
364
365                // Write bits per entry
366                bits.write(writer)?;
367
368                // Write Palette
369                match mode {
370                    PaletteMode::Linear | PaletteMode::Hash => {
371                        VarInt(data.palette.len() as i32).write(writer)?;
372                        for (val, _) in &data.palette {
373                            VarInt(val.to_global_id() as i32).write(writer)?;
374                        }
375                    }
376                    PaletteMode::Global => {}
377                }
378
379                // Pack data
380                let indices: Vec<u32> = data
381                    .cube
382                    .iter()
383                    .flatten()
384                    .flatten()
385                    .map(|val| {
386                        if matches!(mode, PaletteMode::Global) {
387                            val.to_global_id()
388                        } else {
389                            data.palette.iter().position(|(v, _)| v == val).unwrap() as u32
390                        }
391                    })
392                    .collect();
393
394                let packed = pack_bits(&indices, bits as usize);
395
396                // writeFixedSizeLongArray: raw longs, no VarInt length prefix
397                for long in packed {
398                    long.write(writer)?;
399                }
400            }
401            Self::Building(_) => {
402                panic!(
403                    "PalettedContainer in Building mode cannot be serialized; \
404                     call finalize_building() first"
405                );
406            }
407        }
408        Ok(())
409    }
410
411    fn calculate_strategy(count: usize) -> (u8, PaletteMode) {
412        if DIM == 16 {
413            // Block states
414            match count {
415                0..=1 => unreachable!("Homogeneous handled separately"),
416                2..=16 => (4, PaletteMode::Linear),
417                17..=32 => (5, PaletteMode::Hash),
418                33..=64 => (6, PaletteMode::Hash),
419                65..=128 => (7, PaletteMode::Hash),
420                129..=256 => (8, PaletteMode::Hash),
421                _ => (block_state_global_bits(), PaletteMode::Global),
422            }
423        } else {
424            // Biomes
425            match count {
426                0..=1 => unreachable!("Homogeneous handled separately"),
427                2 => (1, PaletteMode::Linear),
428                3..=4 => (2, PaletteMode::Linear),
429                5..=8 => (3, PaletteMode::Hash),
430                _ => (biome_global_bits(), PaletteMode::Global),
431            }
432        }
433    }
434}
435
436/// Bits per entry for the block-state global palette.
437///
438/// Derived from the registered block-state count and cached on first use, so a
439/// version bump that changes the state count doesn't need a matching hardcoded
440/// constant here. Matches vanilla, which sizes the global palette from the
441/// registry rather than a fixed width.
442fn block_state_global_bits() -> u8 {
443    ceil_log2(usize::from(REGISTRY.blocks.next_state_id))
444}
445
446/// Bits per entry for the biome global palette. See [`block_state_global_bits`].
447fn biome_global_bits() -> u8 {
448    static BITS: OnceLock<u8> = OnceLock::new();
449    *BITS.get_or_init(|| ceil_log2(REGISTRY.biomes.len()))
450}
451
452fn pack_bits(indices: &[u32], bits: usize) -> Vec<u64> {
453    let values_per_long = 64 / bits;
454    let len = indices.len().div_ceil(values_per_long);
455    let mut data = vec![0u64; len];
456
457    for (i, &index) in indices.iter().enumerate() {
458        let array_index = i / values_per_long;
459        let offset = (i % values_per_long) * bits;
460        data[array_index] |= u64::from(index) << offset;
461    }
462
463    data
464}
465
466/// A palette container for blocks.
467pub type BlockPalette = PalettedContainer<BlockStateId, 16>;
468/// A palette container for biomes.
469pub type BiomePalette = PalettedContainer<u16, 4>;
470
471impl BlockPalette {
472    /// Gets the number of non-empty blocks in the container.
473    #[must_use]
474    pub fn non_empty_block_count(&self) -> u16 {
475        match self {
476            Self::Homogeneous(v) => {
477                if v.0 == 0 {
478                    0
479                } else {
480                    #[expect(
481                        clippy::cast_possible_truncation,
482                        reason = "VOLUME = 16^3 = 4096, fits in u16"
483                    )]
484                    {
485                        Self::VOLUME as u16
486                    }
487                }
488            }
489            Self::Heterogeneous(data) => {
490                let mut count = 0;
491                for (v, c) in &data.palette {
492                    if v.0 != 0 {
493                        count += c;
494                    }
495                }
496                count
497            }
498            Self::Building(cube) => {
499                let mut count: u16 = 0;
500                for slab in cube {
501                    for row in slab {
502                        for v in row {
503                            if v.0 != 0 {
504                                count += 1;
505                            }
506                        }
507                    }
508                }
509                count
510            }
511        }
512    }
513
514    /// Returns `true` if this palette contains only air blocks.
515    #[must_use]
516    pub fn has_only_air(&self) -> bool {
517        match self {
518            Self::Homogeneous(v) => v.is_air(),
519            //TODO: Use a nonEmpty counter?
520            Self::Heterogeneous(_data) => false,
521            Self::Building(cube) => cube
522                .iter()
523                .flatten()
524                .flatten()
525                .all(steel_utils::BlockStateId::is_air),
526        }
527    }
528}
529
530#[cfg(test)]
531mod tests {
532    use super::{BiomePalette, BlockPalette, biome_global_bits, block_state_global_bits};
533    use steel_registry::init_vanilla_registry;
534    use steel_utils::BlockStateId;
535
536    fn assert_column_matches_get(container: &BlockPalette, x: usize, z: usize) {
537        let mut column = [BlockStateId::default(); 16];
538        container.copy_column_into(x, z, &mut column);
539        for (y, state) in column.into_iter().enumerate() {
540            assert_eq!(state, container.get(x, y, z));
541        }
542    }
543
544    #[test]
545    fn get_at_index_reads_homogeneous_palette_values() {
546        let container = BlockPalette::Homogeneous(BlockStateId(42));
547
548        assert_eq!(container.get_at_index(0), BlockStateId(42));
549        assert_eq!(
550            container.get_at_index(BlockPalette::VOLUME - 1),
551            BlockStateId(42)
552        );
553    }
554
555    #[test]
556    fn get_at_index_uses_y_z_x_linear_order_for_blocks() {
557        let mut container = BlockPalette::Homogeneous(BlockStateId(0));
558
559        container.set(3, 4, 5, BlockStateId(99));
560
561        assert_eq!(
562            container.get_at_index(3 + 5 * 16 + 4 * 16 * 16),
563            BlockStateId(99)
564        );
565    }
566
567    #[test]
568    fn maybe_has_checks_palette_values_without_scanning_cells() {
569        let mut container = BlockPalette::Homogeneous(BlockStateId(0));
570        assert!(!container.maybe_has(|state| state == BlockStateId(7)));
571
572        container.set(1, 2, 3, BlockStateId(7));
573        assert!(container.maybe_has(|state| state == BlockStateId(7)));
574        assert!(!container.maybe_has(|state| state == BlockStateId(9)));
575    }
576
577    #[test]
578    fn copy_column_into_matches_get_for_homogeneous_container() {
579        let container = BlockPalette::Homogeneous(BlockStateId(7));
580        assert_column_matches_get(&container, 3, 12);
581    }
582
583    #[test]
584    fn copy_column_into_matches_get_for_heterogeneous_container() {
585        let x = 5;
586        let z = 9;
587        let mut cube = Box::new([[[BlockStateId::default(); 16]; 16]; 16]);
588        for y in 0..16 {
589            cube[y][z][x] = BlockStateId((y + 1) as u16);
590        }
591
592        let container = BlockPalette::from_cube(cube);
593        assert_column_matches_get(&container, x, z);
594    }
595
596    #[test]
597    fn copy_column_into_matches_get_for_building_container() {
598        let x = 11;
599        let z = 2;
600        let mut container = BlockPalette::Homogeneous(BlockStateId(3));
601        container.enter_building_mode();
602        for y in 0..16 {
603            container.set(x, y, z, BlockStateId((31 + y) as u16));
604        }
605
606        assert_column_matches_get(&container, x, z);
607    }
608
609    /// The written long array has no length prefix, so a wrong bit width
610    /// desyncs every section parsed after this one.
611    #[test]
612    fn block_global_palette_uses_registry_derived_bit_width_and_long_count() {
613        init_vanilla_registry();
614        let mut cube = Box::new([[[BlockStateId::default(); 16]; 16]; 16]);
615        let mut next_id = 1u16;
616        for slab in &mut cube {
617            for row in slab.iter_mut() {
618                for value in row.iter_mut() {
619                    *value = BlockStateId(next_id);
620                    next_id = next_id.wrapping_add(1);
621                }
622            }
623        }
624        let container = BlockPalette::from_cube(cube);
625
626        let mut written = Vec::new();
627        container
628            .write(&mut written)
629            .expect("writing a global-palette container should not fail");
630
631        let expected_bits = block_state_global_bits();
632        assert_eq!(written[0], expected_bits);
633        let values_per_long = 64 / usize::from(expected_bits);
634        let expected_longs = BlockPalette::VOLUME.div_ceil(values_per_long);
635        assert_eq!(written.len(), 1 + expected_longs * 8);
636    }
637
638    #[test]
639    fn biome_global_palette_uses_registry_derived_bit_width_and_long_count() {
640        init_vanilla_registry();
641        let mut cube = Box::new([[[0u16; 4]; 4]; 4]);
642        let mut next_id = 1u16;
643        for slab in &mut cube {
644            for row in slab.iter_mut() {
645                for value in row.iter_mut() {
646                    *value = next_id;
647                    next_id += 1;
648                }
649            }
650        }
651        let container = BiomePalette::from_cube(cube);
652
653        let mut written = Vec::new();
654        container
655            .write(&mut written)
656            .expect("writing a global-palette container should not fail");
657
658        let expected_bits = biome_global_bits();
659        assert_eq!(written[0], expected_bits);
660        let values_per_long = 64 / usize::from(expected_bits);
661        let expected_longs = BiomePalette::VOLUME.div_ceil(values_per_long);
662        assert_eq!(written.len(), 1 + expected_longs * 8);
663    }
664}
steel_core/chunk/paletted_container.rs

steel_core/chunk/
paletted_container.rs
