1use std::io::{Cursor, Write};
2
3use azalea_block::block_state::BlockStateIntegerRepr;
4use azalea_buf::{AzaleaRead, AzaleaReadVar, AzaleaWrite, AzaleaWriteVar, 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(0);
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) -> BlockStateIntegerRepr {
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) -> BlockStateIntegerRepr {
117 self.get_at_index(self.index_from_coords(x, y, z))
120 }
121
122 pub fn get_and_set(
124 &mut self,
125 x: usize,
126 y: usize,
127 z: usize,
128 value: BlockStateIntegerRepr,
129 ) -> BlockStateIntegerRepr {
130 let paletted_value = self.id_for(value);
131 self.storage
132 .get_and_set(self.index_from_coords(x, y, z), paletted_value as u64)
133 as BlockStateIntegerRepr
134 }
135
136 pub fn set_at_index(&mut self, index: usize, value: BlockStateIntegerRepr) {
139 let paletted_value = self.id_for(value);
140 self.storage.set(index, paletted_value as u64);
141 }
142
143 pub fn set(&mut self, x: usize, y: usize, z: usize, value: BlockStateIntegerRepr) {
145 self.set_at_index(self.index_from_coords(x, y, z), value);
146 }
147
148 fn create_or_reuse_data(&self, bits_per_entry: u8) -> PalettedContainer {
149 let new_palette_type =
150 PaletteKind::from_bits_and_type(bits_per_entry, &self.container_type);
151
152 let old_palette_type = (&self.palette).into();
153 if bits_per_entry == self.bits_per_entry && new_palette_type == old_palette_type {
154 return self.clone();
155 }
156 let storage =
157 BitStorage::new(bits_per_entry as usize, self.container_type.size(), None).unwrap();
158
159 debug_assert_eq!(storage.size(), self.container_type.size());
161
162 let palette = new_palette_type.as_empty_palette();
165 PalettedContainer {
166 bits_per_entry,
167 palette,
168 storage,
169 container_type: self.container_type,
170 }
171 }
172
173 fn on_resize(&mut self, bits_per_entry: u8, value: BlockStateIntegerRepr) -> usize {
174 let mut new_data = self.create_or_reuse_data(bits_per_entry);
177 new_data.copy_from(&self.palette, &self.storage);
178 *self = new_data;
179 self.id_for(value)
180 }
181
182 fn copy_from(&mut self, palette: &Palette, storage: &BitStorage) {
183 for i in 0..storage.size() {
184 let value = palette.value_for(storage.get(i) as usize);
185 let id = self.id_for(value) as u64;
186 self.storage.set(i, id);
187 }
188 }
189
190 pub fn id_for(&mut self, value: BlockStateIntegerRepr) -> usize {
191 match &mut self.palette {
192 Palette::SingleValue(v) => {
193 if *v != value {
194 self.on_resize(1, value)
195 } else {
196 0
197 }
198 }
199 Palette::Linear(palette) => {
200 if let Some(index) = palette.iter().position(|&v| v == value) {
201 return index;
202 }
203 let capacity = 2usize.pow(self.bits_per_entry.into());
204 if capacity > palette.len() {
205 palette.push(value);
206 palette.len() - 1
207 } else {
208 self.on_resize(self.bits_per_entry + 1, value)
209 }
210 }
211 Palette::Hashmap(palette) => {
212 if let Some(index) = palette.iter().position(|v| *v == value) {
214 return index;
215 }
216 let capacity = 2usize.pow(self.bits_per_entry.into());
217 if capacity > palette.len() {
218 palette.push(value);
219 palette.len() - 1
220 } else {
221 self.on_resize(self.bits_per_entry + 1, value)
222 }
223 }
224 Palette::Global => value as usize,
225 }
226 }
227}
228
229impl AzaleaWrite for PalettedContainer {
230 fn azalea_write(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
231 self.bits_per_entry.azalea_write(buf)?;
232 self.palette.azalea_write(buf)?;
233 self.storage.data.azalea_write(buf)?;
234 Ok(())
235 }
236}
237
238#[derive(Clone, Debug, PartialEq, Eq)]
239pub enum PaletteKind {
240 SingleValue,
241 Linear,
242 Hashmap,
243 Global,
244}
245
246#[derive(Clone, Debug)]
248pub enum Palette {
249 SingleValue(BlockStateIntegerRepr),
251 Linear(Vec<BlockStateIntegerRepr>),
254 Hashmap(Vec<BlockStateIntegerRepr>),
255 Global,
256}
257
258impl Palette {
259 pub fn value_for(&self, id: usize) -> BlockStateIntegerRepr {
260 match self {
261 Palette::SingleValue(v) => *v,
262 Palette::Linear(v) => v[id],
263 Palette::Hashmap(v) => v.get(id).copied().unwrap_or_default(),
264 Palette::Global => id as BlockStateIntegerRepr,
265 }
266 }
267}
268
269impl AzaleaWrite for Palette {
270 fn azalea_write(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
271 match self {
272 Palette::SingleValue(value) => {
273 value.azalea_write_var(buf)?;
274 }
275 Palette::Linear(values) => {
276 values.azalea_write_var(buf)?;
277 }
278 Palette::Hashmap(values) => {
279 values.azalea_write_var(buf)?;
280 }
281 Palette::Global => {}
282 }
283 Ok(())
284 }
285}
286
287impl PaletteKind {
288 pub fn from_bits_and_type(bits_per_entry: u8, container_type: &PalettedContainerKind) -> Self {
289 match container_type {
290 PalettedContainerKind::BlockStates => match bits_per_entry {
291 0 => PaletteKind::SingleValue,
292 1..=4 => PaletteKind::Linear,
293 5..=8 => PaletteKind::Hashmap,
294 _ => PaletteKind::Global,
295 },
296 PalettedContainerKind::Biomes => match bits_per_entry {
297 0 => PaletteKind::SingleValue,
298 1..=3 => PaletteKind::Linear,
299 _ => PaletteKind::Global,
300 },
301 }
302 }
303
304 pub fn read(&self, buf: &mut Cursor<&[u8]>) -> Result<Palette, BufReadError> {
305 Ok(match self {
306 PaletteKind::SingleValue => {
309 Palette::SingleValue(BlockStateIntegerRepr::azalea_read_var(buf)?)
310 }
311 PaletteKind::Linear => {
312 Palette::Linear(Vec::<BlockStateIntegerRepr>::azalea_read_var(buf)?)
313 }
314 PaletteKind::Hashmap => {
315 Palette::Hashmap(Vec::<BlockStateIntegerRepr>::azalea_read_var(buf)?)
316 }
317 PaletteKind::Global => Palette::Global,
318 })
319 }
320
321 pub fn as_empty_palette(&self) -> Palette {
322 match self {
323 PaletteKind::SingleValue => Palette::SingleValue(0),
324 PaletteKind::Linear => Palette::Linear(Vec::new()),
325 PaletteKind::Hashmap => Palette::Hashmap(Vec::new()),
326 PaletteKind::Global => Palette::Global,
327 }
328 }
329}
330
331impl From<&Palette> for PaletteKind {
332 fn from(palette: &Palette) -> Self {
333 match palette {
334 Palette::SingleValue(_) => PaletteKind::SingleValue,
335 Palette::Linear(_) => PaletteKind::Linear,
336 Palette::Hashmap(_) => PaletteKind::Hashmap,
337 Palette::Global => PaletteKind::Global,
338 }
339 }
340}
341
342impl PalettedContainerKind {
343 fn size_bits(&self) -> usize {
344 match self {
345 PalettedContainerKind::BlockStates => 4,
346 PalettedContainerKind::Biomes => 2,
347 }
348 }
349
350 fn size(&self) -> usize {
351 1 << (self.size_bits() * 3)
352 }
353}
354
355#[cfg(test)]
356mod tests {
357 use super::*;
358
359 #[test]
360 fn test_resize_0_bits_to_1() {
361 let mut palette_container = PalettedContainer::new(PalettedContainerKind::BlockStates);
362
363 assert_eq!(palette_container.bits_per_entry, 0);
364 assert_eq!(palette_container.get_at_index(0), 0);
365 assert_eq!(
366 PaletteKind::from(&palette_container.palette),
367 PaletteKind::SingleValue
368 );
369 palette_container.set_at_index(0, 1);
370 assert_eq!(palette_container.get_at_index(0), 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 palette_container.set_at_index(0, 0); assert_eq!(palette_container.bits_per_entry, 0);
383
384 palette_container.set_at_index(1, 1); assert_eq!(palette_container.bits_per_entry, 1);
386
387 palette_container.set_at_index(2, 2); assert_eq!(palette_container.bits_per_entry, 2);
389 palette_container.set_at_index(3, 3);
390
391 palette_container.set_at_index(4, 4); assert_eq!(palette_container.bits_per_entry, 3);
393 palette_container.set_at_index(5, 5);
394 palette_container.set_at_index(6, 6);
395 palette_container.set_at_index(7, 7);
396
397 palette_container.set_at_index(8, 8); assert_eq!(palette_container.bits_per_entry, 4);
399 palette_container.set_at_index(9, 9);
400 palette_container.set_at_index(10, 10);
401 palette_container.set_at_index(11, 11);
402 palette_container.set_at_index(12, 12);
403 palette_container.set_at_index(13, 13);
404 palette_container.set_at_index(14, 14);
405 palette_container.set_at_index(15, 15);
406 assert_eq!(palette_container.bits_per_entry, 4);
407
408 palette_container.set_at_index(16, 16); assert_eq!(palette_container.bits_per_entry, 5);
410 }
411
412 #[test]
413 fn test_coords_from_index() {
414 let palette_container = PalettedContainer::new(PalettedContainerKind::BlockStates);
415
416 for x in 0..15 {
417 for y in 0..15 {
418 for z in 0..15 {
419 assert_eq!(
420 palette_container
421 .coords_from_index(palette_container.index_from_coords(x, y, z)),
422 (x, y, z)
423 );
424 }
425 }
426 }
427 }
428}