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

steel_core/chunk/
paletted_container.rs
