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 pub palette: Palette,
24 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 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 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 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 pub fn get_at_index(&self, index: usize) -> BlockState {
109 let paletted_value = self.storage.get(index);
111 self.palette.value_for(paletted_value as usize)
113 }
114
115 pub fn get(&self, x: usize, y: usize, z: usize) -> BlockState {
117 self.get_at_index(self.index_from_coords(x, y, z))
120 }
121
122 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 #[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 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 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 debug_assert_eq!(storage.size(), self.container_type.size());
165
166 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 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 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#[derive(Clone, Debug)]
253pub enum Palette {
254 SingleValue(BlockState),
256 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 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); assert_eq!(palette_container.bits_per_entry, 0);
387
388 set(&mut palette_container, 1, 1); assert_eq!(palette_container.bits_per_entry, 1);
390
391 set(&mut palette_container, 2, 2); assert_eq!(palette_container.bits_per_entry, 2);
393 set(&mut palette_container, 3, 3);
394
395 set(&mut palette_container, 4, 4); 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); 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); 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}