azalea_world/palette/
container.rs

1use std::{
2    fmt::Debug,
3    io::{self, Cursor, Write},
4};
5
6use azalea_block::BlockState;
7use azalea_buf::{AzBuf, BufReadError};
8use azalea_core::position::{ChunkSectionBiomePos, ChunkSectionBlockPos};
9use azalea_registry::data::Biome;
10use tracing::{debug, warn};
11
12use super::{Palette, PaletteKind};
13use crate::BitStorage;
14
15#[derive(Clone, Debug, PartialEq)]
16pub struct PalettedContainer<S: PalletedContainerKind> {
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<S>,
24    /// Compacted list of indices pointing to entry IDs in the Palette.
25    pub storage: BitStorage,
26}
27
28pub trait PalletedContainerKind:
29    Copy + Clone + Debug + Default + PartialEq + TryFrom<u32> + Into<u32>
30{
31    type SectionPos: SectionPos;
32
33    fn size_bits() -> usize;
34
35    fn size() -> usize {
36        1 << (Self::size_bits() * 3)
37    }
38
39    fn bits_per_entry_to_palette_kind(bits_per_entry: u8) -> PaletteKind;
40}
41impl PalletedContainerKind for BlockState {
42    type SectionPos = ChunkSectionBlockPos;
43
44    fn size_bits() -> usize {
45        4
46    }
47
48    fn bits_per_entry_to_palette_kind(bits_per_entry: u8) -> PaletteKind {
49        match bits_per_entry {
50            0 => PaletteKind::SingleValue,
51            1..=4 => PaletteKind::Linear,
52            5..=8 => PaletteKind::Hashmap,
53            _ => PaletteKind::Global,
54        }
55    }
56}
57impl PalletedContainerKind for Biome {
58    type SectionPos = ChunkSectionBiomePos;
59
60    fn size_bits() -> usize {
61        2
62    }
63
64    fn bits_per_entry_to_palette_kind(bits_per_entry: u8) -> PaletteKind {
65        match bits_per_entry {
66            0 => PaletteKind::SingleValue,
67            1..=3 => PaletteKind::Linear,
68            _ => PaletteKind::Global,
69        }
70    }
71}
72
73/// A trait for position types that are sometimes valid ways to index into a
74/// chunk section.
75pub trait SectionPos {
76    fn coords(&self) -> (usize, usize, usize);
77    fn new(x: usize, y: usize, z: usize) -> Self;
78}
79impl SectionPos for ChunkSectionBlockPos {
80    fn coords(&self) -> (usize, usize, usize) {
81        (self.x as usize, self.y as usize, self.z as usize)
82    }
83
84    fn new(x: usize, y: usize, z: usize) -> Self {
85        ChunkSectionBlockPos {
86            x: x as u8,
87            y: y as u8,
88            z: z as u8,
89        }
90    }
91}
92impl SectionPos for ChunkSectionBiomePos {
93    fn coords(&self) -> (usize, usize, usize) {
94        (self.x as usize, self.y as usize, self.z as usize)
95    }
96
97    fn new(x: usize, y: usize, z: usize) -> Self {
98        ChunkSectionBiomePos {
99            x: x as u8,
100            y: y as u8,
101            z: z as u8,
102        }
103    }
104}
105
106impl<S: PalletedContainerKind> PalettedContainer<S> {
107    pub fn new() -> Self {
108        let palette = Palette::SingleValue(S::default());
109        let size = S::size();
110        let storage = BitStorage::new(0, size, Some(Box::new([]))).unwrap();
111
112        PalettedContainer {
113            bits_per_entry: 0,
114            palette,
115            storage,
116        }
117    }
118
119    pub fn read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
120        let bits_per_entry = u8::azalea_read(buf)?;
121        let palette_type = S::bits_per_entry_to_palette_kind(bits_per_entry);
122        let palette = palette_type.read(buf)?;
123        let size = S::size();
124
125        let mut storage = match BitStorage::new(
126            bits_per_entry as usize,
127            size,
128            if bits_per_entry == 0 {
129                Some(Box::new([]))
130            } else {
131                // we're going to update the data after creating the bitstorage
132                None
133            },
134        ) {
135            Ok(storage) => storage,
136            Err(e) => {
137                warn!("Failed to create bit storage: {:?}", e);
138                return Err(BufReadError::Custom(
139                    "Failed to create bit storage".to_owned(),
140                ));
141            }
142        };
143
144        // now read the data
145        for i in 0..storage.data.len() {
146            storage.data[i] = u64::azalea_read(buf)?;
147        }
148
149        Ok(PalettedContainer {
150            bits_per_entry,
151            palette,
152            storage,
153        })
154    }
155
156    /// Calculates the index of the given position.
157    pub fn index_from_pos(&self, pos: S::SectionPos) -> usize {
158        let size_bits = S::size_bits();
159        let (x, y, z) = pos.coords();
160        (((y << size_bits) | z) << size_bits) | x
161    }
162
163    pub fn pos_from_index(&self, index: usize) -> S::SectionPos {
164        let size_bits = S::size_bits();
165        let mask = (1 << size_bits) - 1;
166        S::SectionPos::new(
167            index & mask,
168            (index >> size_bits >> size_bits) & mask,
169            (index >> size_bits) & mask,
170        )
171    }
172
173    /// Returns the value at the given index.
174    ///
175    /// # Panics
176    ///
177    /// This function panics if the index is greater than or equal to the number
178    /// of things in the storage. For example, for block states, it must be less
179    /// than 4096.
180    #[inline]
181    pub fn get_at_index(&self, index: usize) -> S {
182        // first get the palette id
183        let paletted_value = self.storage.get(index);
184        // and then get the value from that id
185        self.palette.value_for(paletted_value as usize)
186    }
187
188    /// Returns the value at the given position.
189    pub fn get(&self, pos: S::SectionPos) -> S {
190        self.get_at_index(self.index_from_pos(pos))
191    }
192
193    /// Sets the ID at the given position and return the previous ID.
194    pub fn get_and_set(&mut self, pos: S::SectionPos, value: S) -> S {
195        let paletted_value = self.id_for(value);
196        let old_paletted_value = self
197            .storage
198            .get_and_set(self.index_from_pos(pos), paletted_value as u64);
199        self.palette.value_for(old_paletted_value as usize)
200    }
201
202    /// Sets the ID at the given index and return the previous ID. You probably
203    /// want `.set` instead.
204    pub fn set_at_index(&mut self, index: usize, value: S) {
205        let paletted_value = self.id_for(value);
206        self.storage.set(index, paletted_value as u64);
207    }
208
209    /// Sets the ID at the given position and return the previous ID.
210    pub fn set(&mut self, pos: S::SectionPos, value: S) {
211        self.set_at_index(self.index_from_pos(pos), value);
212    }
213
214    fn create_or_reuse_data(&self, bits_per_entry: u8) -> PalettedContainer<S> {
215        let new_palette_type = S::bits_per_entry_to_palette_kind(bits_per_entry);
216
217        let old_palette_type = (&self.palette).into();
218        if bits_per_entry == self.bits_per_entry && new_palette_type == old_palette_type {
219            return self.clone();
220        }
221        let storage = BitStorage::new(bits_per_entry as usize, S::size(), None).unwrap();
222
223        // sanity check
224        debug_assert_eq!(storage.size(), S::size());
225
226        // let palette = new_palette_type.as_empty_palette(1usize << (bits_per_entry as
227        // usize));
228        let palette = new_palette_type.as_empty_palette();
229        PalettedContainer {
230            bits_per_entry,
231            palette,
232            storage,
233        }
234    }
235
236    fn on_resize(&mut self, bits_per_entry: u8, value: S) -> usize {
237        debug!(
238            "Resizing PalettedContainer from {} bpe to {bits_per_entry} for {value:?} with palette={:?}",
239            self.bits_per_entry, self.palette
240        );
241        // in vanilla this is always true, but it's sometimes false in purpur servers
242        // assert!(bits_per_entry <= 5, "bits_per_entry must be <= 5");
243        let mut new_data = self.create_or_reuse_data(bits_per_entry);
244        new_data.copy_from(&self.palette, &self.storage);
245        *self = new_data;
246        self.id_for(value)
247    }
248
249    fn copy_from(&mut self, palette: &Palette<S>, storage: &BitStorage) {
250        for i in 0..storage.size() {
251            let value = palette.value_for(storage.get(i) as usize);
252            let id = self.id_for(value) as u64;
253            self.storage.set(i, id);
254        }
255    }
256
257    pub fn id_for(&mut self, value: S) -> usize {
258        match &mut self.palette {
259            Palette::SingleValue(v) => {
260                if (*v).into() != value.into() {
261                    self.on_resize(1, value)
262                } else {
263                    0
264                }
265            }
266            Palette::Linear(palette) => {
267                if let Some(index) = palette.iter().position(|&v| v.into() == value.into()) {
268                    return index;
269                }
270                let capacity = 2usize.pow(self.bits_per_entry.into());
271                if capacity > palette.len() {
272                    palette.push(value);
273                    palette.len() - 1
274                } else {
275                    self.on_resize(self.bits_per_entry + 1, value)
276                }
277            }
278            Palette::Hashmap(palette) => {
279                // TODO? vanilla keeps this in memory as a hashmap, but it should be benchmarked
280                // before changing it
281                if let Some(index) = palette.iter().position(|v| (*v).into() == value.into()) {
282                    return index;
283                }
284                let capacity = 2usize.pow(self.bits_per_entry.into());
285                if capacity > palette.len() {
286                    palette.push(value);
287                    palette.len() - 1
288                } else {
289                    self.on_resize(self.bits_per_entry + 1, value)
290                }
291            }
292            Palette::Global => value.into() as usize,
293        }
294    }
295}
296
297impl<S: PalletedContainerKind> PalettedContainer<S> {
298    pub fn write(&self, buf: &mut impl Write) -> io::Result<()> {
299        self.bits_per_entry.azalea_write(buf)?;
300        self.palette.write(buf)?;
301        for word in &self.storage.data {
302            word.azalea_write(buf)?;
303        }
304
305        Ok(())
306    }
307}
308
309impl<S: PalletedContainerKind> Default for PalettedContainer<S> {
310    fn default() -> Self {
311        Self::new()
312    }
313}
azalea_world/palette/container.rs

azalea_world/palette/
container.rs
