azalea_world/palette/
container.rs

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

azalea_world/palette/
container.rs
