azalea_world/
palette.rs

1use std::io::{Cursor, Write};
2
3use azalea_block::BlockState;
4use azalea_buf::{AzaleaRead, AzaleaWrite, BufReadError};
5use tracing::warn;
6
7use crate::BitStorage;
8
9#[derive(Clone, Debug, Copy)]
10pub enum PalettedContainerKind {
11    Biomes,
12    BlockStates,
13}
14
15#[derive(Clone, Debug)]
16pub struct PalettedContainer {
17    pub bits_per_entry: u8,
18    /// This is usually a list of unique values that appear in the container so
19    /// they can be indexed by the bit storage.
20    ///
21    /// Sometimes it doesn't contain anything if there's too many unique items
22    /// in the bit storage, though.
23    pub palette: Palette,
24    /// Compacted list of indices pointing to entry IDs in the Palette.
25    pub storage: BitStorage,
26    pub container_type: PalettedContainerKind,
27}
28
29impl PalettedContainer {
30    pub fn new(container_type: PalettedContainerKind) -> Self {
31        let palette = Palette::SingleValue(BlockState::AIR);
32        let size = container_type.size();
33        let storage = BitStorage::new(0, size, Some(Box::new([]))).unwrap();
34
35        PalettedContainer {
36            bits_per_entry: 0,
37            palette,
38            storage,
39            container_type,
40        }
41    }
42
43    pub fn read_with_type(
44        buf: &mut Cursor<&[u8]>,
45        container_type: &'static PalettedContainerKind,
46    ) -> Result<Self, BufReadError> {
47        let bits_per_entry = u8::azalea_read(buf)?;
48        let palette_type = PaletteKind::from_bits_and_type(bits_per_entry, container_type);
49        let palette = palette_type.read(buf)?;
50        let size = container_type.size();
51
52        let mut storage = match BitStorage::new(
53            bits_per_entry as usize,
54            size,
55            if bits_per_entry == 0 {
56                Some(Box::new([]))
57            } else {
58                // we're going to update the data after creating the bitstorage
59                None
60            },
61        ) {
62            Ok(storage) => storage,
63            Err(e) => {
64                warn!("Failed to create bit storage: {:?}", e);
65                return Err(BufReadError::Custom(
66                    "Failed to create bit storage".to_string(),
67                ));
68            }
69        };
70
71        // now read the data
72        for i in 0..storage.data.len() {
73            storage.data[i] = u64::azalea_read(buf)?;
74        }
75
76        Ok(PalettedContainer {
77            bits_per_entry,
78            palette,
79            storage,
80            container_type: *container_type,
81        })
82    }
83
84    /// Calculates the index of the given coordinates.
85    pub fn index_from_coords(&self, x: usize, y: usize, z: usize) -> usize {
86        let size_bits = self.container_type.size_bits();
87
88        (((y << size_bits) | z) << size_bits) | x
89    }
90
91    pub fn coords_from_index(&self, index: usize) -> (usize, usize, usize) {
92        let size_bits = self.container_type.size_bits();
93        let mask = (1 << size_bits) - 1;
94        (
95            index & mask,
96            (index >> size_bits >> size_bits) & mask,
97            (index >> size_bits) & mask,
98        )
99    }
100
101    /// Returns the value at the given index.
102    ///
103    /// # Panics
104    ///
105    /// This function panics if the index is greater than or equal to the number
106    /// of things in the storage. (So for block states, it must be less than
107    /// 4096).
108    pub fn get_at_index(&self, index: usize) -> BlockState {
109        // first get the palette id
110        let paletted_value = self.storage.get(index);
111        // and then get the value from that id
112        self.palette.value_for(paletted_value as usize)
113    }
114
115    /// Returns the value at the given coordinates.
116    pub fn get(&self, x: usize, y: usize, z: usize) -> BlockState {
117        // let paletted_value = self.storage.get(self.get_index(x, y, z));
118        // self.palette.value_for(paletted_value as usize)
119        self.get_at_index(self.index_from_coords(x, y, z))
120    }
121
122    /// Sets the id at the given coordinates and return the previous id
123    pub fn get_and_set(&mut self, x: usize, y: usize, z: usize, value: BlockState) -> BlockState {
124        let paletted_value = self.id_for(value);
125        let block_state_id = self
126            .storage
127            .get_and_set(self.index_from_coords(x, y, z), paletted_value as u64);
128        // error in debug mode
129        #[cfg(debug_assertions)]
130        if block_state_id > BlockState::MAX_STATE.into() {
131            warn!(
132                "Old block state from get_and_set {block_state_id} was greater than max state {}",
133                BlockState::MAX_STATE
134            );
135        }
136
137        BlockState::try_from(block_state_id as u32).unwrap_or_default()
138    }
139
140    /// Sets the id at the given index and return the previous id. You probably
141    /// want `.set` instead.
142    pub fn set_at_index(&mut self, index: usize, value: BlockState) {
143        let paletted_value = self.id_for(value);
144        self.storage.set(index, paletted_value as u64);
145    }
146
147    /// Sets the id at the given coordinates and return the previous id
148    pub fn set(&mut self, x: usize, y: usize, z: usize, value: BlockState) {
149        self.set_at_index(self.index_from_coords(x, y, z), value);
150    }
151
152    fn create_or_reuse_data(&self, bits_per_entry: u8) -> PalettedContainer {
153        let new_palette_type =
154            PaletteKind::from_bits_and_type(bits_per_entry, &self.container_type);
155
156        let old_palette_type = (&self.palette).into();
157        if bits_per_entry == self.bits_per_entry && new_palette_type == old_palette_type {
158            return self.clone();
159        }
160        let storage =
161            BitStorage::new(bits_per_entry as usize, self.container_type.size(), None).unwrap();
162
163        // sanity check
164        debug_assert_eq!(storage.size(), self.container_type.size());
165
166        // let palette = new_palette_type.as_empty_palette(1usize << (bits_per_entry as
167        // usize));
168        let palette = new_palette_type.as_empty_palette();
169        PalettedContainer {
170            bits_per_entry,
171            palette,
172            storage,
173            container_type: self.container_type,
174        }
175    }
176
177    fn on_resize(&mut self, bits_per_entry: u8, value: BlockState) -> usize {
178        // in vanilla this is always true, but it's sometimes false in purpur servers
179        // assert!(bits_per_entry <= 5, "bits_per_entry must be <= 5");
180        let mut new_data = self.create_or_reuse_data(bits_per_entry);
181        new_data.copy_from(&self.palette, &self.storage);
182        *self = new_data;
183        self.id_for(value)
184    }
185
186    fn copy_from(&mut self, palette: &Palette, storage: &BitStorage) {
187        for i in 0..storage.size() {
188            let value = palette.value_for(storage.get(i) as usize);
189            let id = self.id_for(value) as u64;
190            self.storage.set(i, id);
191        }
192    }
193
194    pub fn id_for(&mut self, value: BlockState) -> usize {
195        match &mut self.palette {
196            Palette::SingleValue(v) => {
197                if *v != value {
198                    self.on_resize(1, value)
199                } else {
200                    0
201                }
202            }
203            Palette::Linear(palette) => {
204                if let Some(index) = palette.iter().position(|&v| v == value) {
205                    return index;
206                }
207                let capacity = 2usize.pow(self.bits_per_entry.into());
208                if capacity > palette.len() {
209                    palette.push(value);
210                    palette.len() - 1
211                } else {
212                    self.on_resize(self.bits_per_entry + 1, value)
213                }
214            }
215            Palette::Hashmap(palette) => {
216                // TODO? vanilla keeps this in memory as a hashmap, but it should be benchmarked
217                // before changing it
218                if let Some(index) = palette.iter().position(|v| *v == value) {
219                    return index;
220                }
221                let capacity = 2usize.pow(self.bits_per_entry.into());
222                if capacity > palette.len() {
223                    palette.push(value);
224                    palette.len() - 1
225                } else {
226                    self.on_resize(self.bits_per_entry + 1, value)
227                }
228            }
229            Palette::Global => value.id() as usize,
230        }
231    }
232}
233
234impl AzaleaWrite for PalettedContainer {
235    fn azalea_write(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
236        self.bits_per_entry.azalea_write(buf)?;
237        self.palette.azalea_write(buf)?;
238        self.storage.data.azalea_write(buf)?;
239        Ok(())
240    }
241}
242
243#[derive(Clone, Debug, PartialEq, Eq)]
244pub enum PaletteKind {
245    SingleValue,
246    Linear,
247    Hashmap,
248    Global,
249}
250
251/// A representation of the different types of chunk palettes Minecraft uses.
252#[derive(Clone, Debug)]
253pub enum Palette {
254    /// ID of the corresponding entry in its global palette
255    SingleValue(BlockState),
256    // in vanilla this keeps a `size` field that might be less than the length, but i'm not sure
257    // it's actually needed?
258    Linear(Vec<BlockState>),
259    Hashmap(Vec<BlockState>),
260    Global,
261}
262
263impl Palette {
264    pub fn value_for(&self, id: usize) -> BlockState {
265        match self {
266            Palette::SingleValue(v) => *v,
267            Palette::Linear(v) => v.get(id).copied().unwrap_or_default(),
268            Palette::Hashmap(v) => v.get(id).copied().unwrap_or_default(),
269            Palette::Global => BlockState::try_from(id as u32).unwrap_or_default(),
270        }
271    }
272}
273
274impl AzaleaWrite for Palette {
275    fn azalea_write(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
276        match self {
277            Palette::SingleValue(value) => {
278                value.azalea_write(buf)?;
279            }
280            Palette::Linear(values) => {
281                values.azalea_write(buf)?;
282            }
283            Palette::Hashmap(values) => {
284                values.azalea_write(buf)?;
285            }
286            Palette::Global => {}
287        }
288        Ok(())
289    }
290}
291
292impl PaletteKind {
293    pub fn from_bits_and_type(bits_per_entry: u8, container_type: &PalettedContainerKind) -> Self {
294        match container_type {
295            PalettedContainerKind::BlockStates => match bits_per_entry {
296                0 => PaletteKind::SingleValue,
297                1..=4 => PaletteKind::Linear,
298                5..=8 => PaletteKind::Hashmap,
299                _ => PaletteKind::Global,
300            },
301            PalettedContainerKind::Biomes => match bits_per_entry {
302                0 => PaletteKind::SingleValue,
303                1..=3 => PaletteKind::Linear,
304                _ => PaletteKind::Global,
305            },
306        }
307    }
308
309    pub fn read(&self, buf: &mut Cursor<&[u8]>) -> Result<Palette, BufReadError> {
310        Ok(match self {
311            // since they're read as varints it's actually fine to just use BlockStateIntegerRepr
312            // instead of the correct type (u32)
313            PaletteKind::SingleValue => Palette::SingleValue(BlockState::azalea_read(buf)?),
314            PaletteKind::Linear => Palette::Linear(Vec::<BlockState>::azalea_read(buf)?),
315            PaletteKind::Hashmap => Palette::Hashmap(Vec::<BlockState>::azalea_read(buf)?),
316            PaletteKind::Global => Palette::Global,
317        })
318    }
319
320    pub fn as_empty_palette(&self) -> Palette {
321        match self {
322            PaletteKind::SingleValue => Palette::SingleValue(BlockState::AIR),
323            PaletteKind::Linear => Palette::Linear(Vec::new()),
324            PaletteKind::Hashmap => Palette::Hashmap(Vec::new()),
325            PaletteKind::Global => Palette::Global,
326        }
327    }
328}
329
330impl From<&Palette> for PaletteKind {
331    fn from(palette: &Palette) -> Self {
332        match palette {
333            Palette::SingleValue(_) => PaletteKind::SingleValue,
334            Palette::Linear(_) => PaletteKind::Linear,
335            Palette::Hashmap(_) => PaletteKind::Hashmap,
336            Palette::Global => PaletteKind::Global,
337        }
338    }
339}
340
341impl PalettedContainerKind {
342    fn size_bits(&self) -> usize {
343        match self {
344            PalettedContainerKind::BlockStates => 4,
345            PalettedContainerKind::Biomes => 2,
346        }
347    }
348
349    fn size(&self) -> usize {
350        1 << (self.size_bits() * 3)
351    }
352}
353
354#[cfg(test)]
355mod tests {
356    use super::*;
357
358    #[test]
359    fn test_resize_0_bits_to_1() {
360        let mut palette_container = PalettedContainer::new(PalettedContainerKind::BlockStates);
361
362        assert_eq!(palette_container.bits_per_entry, 0);
363        assert_eq!(palette_container.get_at_index(0), BlockState::AIR);
364        assert_eq!(
365            PaletteKind::from(&palette_container.palette),
366            PaletteKind::SingleValue
367        );
368        let block_state_1 = BlockState::try_from(1_u32).unwrap();
369        palette_container.set_at_index(0, block_state_1);
370        assert_eq!(palette_container.get_at_index(0), block_state_1);
371        assert_eq!(
372            PaletteKind::from(&palette_container.palette),
373            PaletteKind::Linear
374        );
375    }
376
377    #[test]
378    fn test_resize_0_bits_to_5() {
379        let mut palette_container = PalettedContainer::new(PalettedContainerKind::BlockStates);
380
381        let set = |pc: &mut PalettedContainer, i, v: u32| {
382            pc.set_at_index(i, BlockState::try_from(v).unwrap());
383        };
384
385        set(&mut palette_container, 0, 0); // 0 bits
386        assert_eq!(palette_container.bits_per_entry, 0);
387
388        set(&mut palette_container, 1, 1); // 1 bit
389        assert_eq!(palette_container.bits_per_entry, 1);
390
391        set(&mut palette_container, 2, 2); // 2 bits
392        assert_eq!(palette_container.bits_per_entry, 2);
393        set(&mut palette_container, 3, 3);
394
395        set(&mut palette_container, 4, 4); // 3 bits
396        assert_eq!(palette_container.bits_per_entry, 3);
397        set(&mut palette_container, 5, 5);
398        set(&mut palette_container, 6, 6);
399        set(&mut palette_container, 7, 7);
400
401        set(&mut palette_container, 8, 8); // 4 bits
402        assert_eq!(palette_container.bits_per_entry, 4);
403        set(&mut palette_container, 9, 9);
404        set(&mut palette_container, 10, 10);
405        set(&mut palette_container, 11, 11);
406        set(&mut palette_container, 12, 12);
407        set(&mut palette_container, 13, 13);
408        set(&mut palette_container, 14, 14);
409        set(&mut palette_container, 15, 15);
410        assert_eq!(palette_container.bits_per_entry, 4);
411
412        set(&mut palette_container, 16, 16); // 5 bits
413        assert_eq!(palette_container.bits_per_entry, 5);
414    }
415
416    #[test]
417    fn test_coords_from_index() {
418        let palette_container = PalettedContainer::new(PalettedContainerKind::BlockStates);
419
420        for x in 0..15 {
421            for y in 0..15 {
422                for z in 0..15 {
423                    assert_eq!(
424                        palette_container
425                            .coords_from_index(palette_container.index_from_coords(x, y, z)),
426                        (x, y, z)
427                    );
428                }
429            }
430        }
431    }
432}