azalea_world/
chunk_storage.rs

1use std::{
2    collections::{HashMap, hash_map::Entry},
3    fmt,
4    fmt::Debug,
5    io,
6    io::{Cursor, Write},
7    sync::{Arc, Weak},
8};
9
10use azalea_block::{
11    block_state::{BlockState, BlockStateIntegerRepr},
12    fluid_state::FluidState,
13};
14use azalea_buf::{AzaleaRead, AzaleaWrite, BufReadError};
15use azalea_core::position::{
16    BlockPos, ChunkBiomePos, ChunkBlockPos, ChunkPos, ChunkSectionBiomePos, ChunkSectionBlockPos,
17};
18use azalea_registry::Biome;
19use nohash_hasher::IntMap;
20use parking_lot::RwLock;
21use tracing::{debug, trace, warn};
22
23use crate::{
24    heightmap::{Heightmap, HeightmapKind},
25    palette::PalettedContainer,
26};
27
28const SECTION_HEIGHT: u32 = 16;
29
30/// An efficient storage of chunks for a client that has a limited render
31/// distance.
32///
33/// This has support for using a shared [`ChunkStorage`].
34pub struct PartialChunkStorage {
35    /// The center of the view, i.e. the chunk the player is currently in.
36    view_center: ChunkPos,
37    pub(crate) chunk_radius: u32,
38    view_range: u32,
39    // chunks is a list of size chunk_radius * chunk_radius
40    chunks: Box<[Option<Arc<RwLock<Chunk>>>]>,
41}
42
43/// A storage for chunks where they're only stored weakly, so if they're not
44/// actively being used somewhere else they'll be forgotten.
45///
46/// This is used for shared worlds.
47///
48/// This is relatively cheap to clone since it's just an `IntMap` with `Weak`
49/// pointers.
50#[derive(Debug, Clone)]
51pub struct ChunkStorage {
52    pub height: u32,
53    pub min_y: i32,
54    pub map: IntMap<ChunkPos, Weak<RwLock<Chunk>>>,
55}
56
57/// A single chunk in a world (16*?*16 blocks).
58///
59/// This only contains blocks and biomes. You can derive the height of the chunk
60/// from the number of sections, but you need a [`ChunkStorage`] to get the
61/// minimum Y coordinate.
62#[derive(Debug)]
63pub struct Chunk {
64    pub sections: Box<[Section]>,
65    /// Heightmaps are used for identifying the surface blocks in a chunk.
66    /// Usually for clients only `WorldSurface` and `MotionBlocking` are
67    /// present.
68    pub heightmaps: HashMap<HeightmapKind, Heightmap>,
69}
70
71/// A section of a chunk, i.e. a 16*16*16 block area.
72#[derive(Clone, Debug, Default)]
73pub struct Section {
74    pub block_count: u16,
75    pub states: PalettedContainer<BlockState>,
76    pub biomes: PalettedContainer<Biome>,
77}
78
79/// Get the actual stored view distance for the selected view distance.
80///
81/// For some reason, Minecraft stores an extra 3 chunks.
82pub fn calculate_chunk_storage_range(view_distance: u32) -> u32 {
83    u32::max(view_distance, 2) + 3
84}
85
86impl Default for Chunk {
87    fn default() -> Self {
88        Chunk {
89            sections: vec![Section::default(); (384 / 16) as usize].into(),
90            heightmaps: HashMap::new(),
91        }
92    }
93}
94
95impl PartialChunkStorage {
96    pub fn new(chunk_radius: u32) -> Self {
97        let view_range = chunk_radius * 2 + 1;
98        PartialChunkStorage {
99            view_center: ChunkPos::new(0, 0),
100            chunk_radius,
101            view_range,
102            chunks: vec![None; (view_range * view_range) as usize].into(),
103        }
104    }
105
106    /// Update the chunk to center the view on.
107    ///
108    /// This should be called when the client receives a `SetChunkCacheCenter`
109    /// packet.
110    pub fn update_view_center(&mut self, view_center: ChunkPos) {
111        // this code block makes it force unload the chunks that are out of range after
112        // updating the view center. it's usually fine without it but the commented code
113        // is there in case you want to temporarily uncomment to test something
114
115        // ```
116        // for index in 0..self.chunks.len() {
117        //     let chunk_pos = self.chunk_pos_from_index(index);
118        //     if !in_range_for_view_center_and_radius(&chunk_pos, view_center, self.chunk_radius) {
119        //         self.chunks[index] = None;
120        //     }
121        // }
122        // ```
123
124        self.view_center = view_center;
125    }
126
127    /// Get the center of the view. This is usually the chunk that the player is
128    /// in.
129    pub fn view_center(&self) -> ChunkPos {
130        self.view_center
131    }
132
133    pub fn view_range(&self) -> u32 {
134        self.view_range
135    }
136
137    pub fn index_from_chunk_pos(&self, chunk_pos: &ChunkPos) -> usize {
138        let view_range = self.view_range as i32;
139
140        let x = i32::rem_euclid(chunk_pos.x, view_range) * view_range;
141        let z = i32::rem_euclid(chunk_pos.z, view_range);
142        (x + z) as usize
143    }
144
145    pub fn chunk_pos_from_index(&self, index: usize) -> ChunkPos {
146        let view_range = self.view_range as i32;
147
148        // find the base from the view center
149        let base_x = self.view_center.x.div_euclid(view_range) * view_range;
150        let base_z = self.view_center.z.div_euclid(view_range) * view_range;
151
152        // add the offset from the base
153        let offset_x = index as i32 / view_range;
154        let offset_z = index as i32 % view_range;
155
156        ChunkPos::new(base_x + offset_x, base_z + offset_z)
157    }
158
159    pub fn in_range(&self, chunk_pos: &ChunkPos) -> bool {
160        in_range_for_view_center_and_radius(chunk_pos, self.view_center, self.chunk_radius)
161    }
162
163    pub fn set_block_state(
164        &self,
165        pos: BlockPos,
166        state: BlockState,
167        chunk_storage: &ChunkStorage,
168    ) -> Option<BlockState> {
169        if pos.y < chunk_storage.min_y
170            || pos.y >= (chunk_storage.min_y + chunk_storage.height as i32)
171        {
172            return None;
173        }
174        let chunk_pos = ChunkPos::from(pos);
175        let chunk_lock = chunk_storage.get(&chunk_pos)?;
176        let mut chunk = chunk_lock.write();
177        Some(chunk.get_and_set_block_state(&ChunkBlockPos::from(pos), state, chunk_storage.min_y))
178    }
179
180    pub fn replace_with_packet_data(
181        &mut self,
182        pos: &ChunkPos,
183        data: &mut Cursor<&[u8]>,
184        heightmaps: &[(HeightmapKind, Box<[u64]>)],
185        chunk_storage: &mut ChunkStorage,
186    ) -> Result<(), BufReadError> {
187        debug!("Replacing chunk at {:?}", pos);
188        if !self.in_range(pos) {
189            warn!("Ignoring chunk since it's not in the view range: {pos:?}");
190            return Ok(());
191        }
192
193        let chunk = Chunk::read_with_dimension_height(
194            data,
195            chunk_storage.height,
196            chunk_storage.min_y,
197            heightmaps,
198        )?;
199
200        self.set(pos, Some(chunk), chunk_storage);
201        trace!("Loaded chunk {pos:?}");
202
203        Ok(())
204    }
205
206    /// Get a [`Chunk`] within render distance, or `None` if it's not loaded.
207    /// Use [`ChunkStorage::get`] to get a chunk from the shared storage.
208    pub fn limited_get(&self, pos: &ChunkPos) -> Option<&Arc<RwLock<Chunk>>> {
209        if !self.in_range(pos) {
210            warn!(
211                "Chunk at {:?} is not in the render distance (center: {:?}, {} chunks)",
212                pos, self.view_center, self.chunk_radius,
213            );
214            return None;
215        }
216
217        let index = self.index_from_chunk_pos(pos);
218        self.chunks[index].as_ref()
219    }
220    /// Get a mutable reference to a [`Chunk`] within render distance, or
221    /// `None` if it's not loaded.
222    ///
223    /// Use [`ChunkStorage::get`] to get a chunk from the shared storage.
224    pub fn limited_get_mut(&mut self, pos: &ChunkPos) -> Option<&mut Option<Arc<RwLock<Chunk>>>> {
225        if !self.in_range(pos) {
226            return None;
227        }
228
229        let index = self.index_from_chunk_pos(pos);
230
231        Some(&mut self.chunks[index])
232    }
233
234    /// Set a chunk in the shared storage and reference it from the limited
235    /// storage.
236    ///
237    /// Use [`Self::limited_set`] if you already have an `Arc<RwLock<Chunk>>`.
238    ///
239    /// # Panics
240    /// If the chunk is not in the render distance.
241    pub fn set(&mut self, pos: &ChunkPos, chunk: Option<Chunk>, chunk_storage: &mut ChunkStorage) {
242        let new_chunk;
243
244        // add the chunk to the shared storage
245        if let Some(chunk) = chunk {
246            match chunk_storage.map.entry(*pos) {
247                Entry::Occupied(mut e) => {
248                    if let Some(old_chunk) = e.get_mut().upgrade() {
249                        *old_chunk.write() = chunk;
250                        new_chunk = Some(old_chunk);
251                    } else {
252                        let chunk_lock = Arc::new(RwLock::new(chunk));
253                        e.insert(Arc::downgrade(&chunk_lock));
254                        new_chunk = Some(chunk_lock);
255                    }
256                }
257                Entry::Vacant(e) => {
258                    let chunk_lock = Arc::new(RwLock::new(chunk));
259                    e.insert(Arc::downgrade(&chunk_lock));
260                    new_chunk = Some(chunk_lock);
261                }
262            }
263        } else {
264            // don't remove it from the shared storage, since it'll be removed
265            // automatically if this was the last reference
266
267            new_chunk = None;
268        }
269
270        self.limited_set(pos, new_chunk);
271    }
272
273    /// Set a chunk in our limited storage, useful if your chunk is already
274    /// referenced somewhere else and you want to make it also be referenced by
275    /// this storage.
276    ///
277    /// Use [`Self::set`] if you don't already have an `Arc<RwLock<Chunk>>`.
278    ///
279    /// # Panics
280    /// If the chunk is not in the render distance.
281    pub fn limited_set(&mut self, pos: &ChunkPos, chunk: Option<Arc<RwLock<Chunk>>>) {
282        if let Some(chunk_mut) = self.limited_get_mut(pos) {
283            *chunk_mut = chunk;
284        }
285    }
286
287    /// Get an iterator over all the chunks in the storage.
288    pub fn chunks(&self) -> impl Iterator<Item = &Option<Arc<RwLock<Chunk>>>> {
289        self.chunks.iter()
290    }
291}
292impl ChunkStorage {
293    pub fn new(height: u32, min_y: i32) -> Self {
294        ChunkStorage {
295            height,
296            min_y,
297            map: IntMap::default(),
298        }
299    }
300
301    pub fn get(&self, pos: &ChunkPos) -> Option<Arc<RwLock<Chunk>>> {
302        self.map.get(pos).and_then(|chunk| chunk.upgrade())
303    }
304
305    pub fn get_block_state(&self, pos: BlockPos) -> Option<BlockState> {
306        let chunk_pos = ChunkPos::from(pos);
307        let chunk = self.get(&chunk_pos)?;
308        let chunk = chunk.read();
309        chunk.get_block_state(&ChunkBlockPos::from(pos), self.min_y)
310    }
311
312    pub fn get_fluid_state(&self, pos: BlockPos) -> Option<FluidState> {
313        let block_state = self.get_block_state(pos)?;
314        Some(FluidState::from(block_state))
315    }
316
317    pub fn get_biome(&self, pos: BlockPos) -> Option<Biome> {
318        let chunk_pos = ChunkPos::from(pos);
319        let chunk = self.get(&chunk_pos)?;
320        let chunk = chunk.read();
321        chunk.get_biome(ChunkBiomePos::from(pos), self.min_y)
322    }
323
324    pub fn set_block_state(&self, pos: BlockPos, state: BlockState) -> Option<BlockState> {
325        if pos.y < self.min_y || pos.y >= (self.min_y + self.height as i32) {
326            return None;
327        }
328        let chunk_pos = ChunkPos::from(pos);
329        let chunk = self.get(&chunk_pos)?;
330        let mut chunk = chunk.write();
331        Some(chunk.get_and_set_block_state(&ChunkBlockPos::from(pos), state, self.min_y))
332    }
333}
334
335pub fn in_range_for_view_center_and_radius(
336    chunk_pos: &ChunkPos,
337    view_center: ChunkPos,
338    chunk_radius: u32,
339) -> bool {
340    (chunk_pos.x - view_center.x).unsigned_abs() <= chunk_radius
341        && (chunk_pos.z - view_center.z).unsigned_abs() <= chunk_radius
342}
343
344impl Chunk {
345    pub fn read_with_dimension_height(
346        buf: &mut Cursor<&[u8]>,
347        dimension_height: u32,
348        min_y: i32,
349        heightmaps_data: &[(HeightmapKind, Box<[u64]>)],
350    ) -> Result<Self, BufReadError> {
351        let section_count = dimension_height / SECTION_HEIGHT;
352        let mut sections = Vec::with_capacity(section_count as usize);
353        for _ in 0..section_count {
354            let section = Section::azalea_read(buf)?;
355            sections.push(section);
356        }
357        let sections = sections.into_boxed_slice();
358
359        let mut heightmaps = HashMap::new();
360        for (kind, data) in heightmaps_data {
361            let data: Box<[u64]> = data.clone();
362            let heightmap = Heightmap::new(*kind, dimension_height, min_y, data);
363            heightmaps.insert(*kind, heightmap);
364        }
365
366        Ok(Chunk {
367            sections,
368            heightmaps,
369        })
370    }
371
372    pub fn get_block_state(&self, pos: &ChunkBlockPos, min_y: i32) -> Option<BlockState> {
373        get_block_state_from_sections(&self.sections, pos, min_y)
374    }
375
376    #[must_use = "Use Chunk::set_block_state instead if you don't need the previous state"]
377    pub fn get_and_set_block_state(
378        &mut self,
379        pos: &ChunkBlockPos,
380        state: BlockState,
381        min_y: i32,
382    ) -> BlockState {
383        let section_index = section_index(pos.y, min_y);
384        let Some(section) = self.sections.get_mut(section_index as usize) else {
385            warn!(
386                "Tried to get and set block state {state:?} at out-of-bounds relative chunk position {pos:?}",
387            );
388            return BlockState::AIR;
389        };
390        let chunk_section_pos = ChunkSectionBlockPos::from(pos);
391        let previous_state = section.get_and_set_block_state(chunk_section_pos, state);
392
393        for heightmap in self.heightmaps.values_mut() {
394            heightmap.update(pos, state, &self.sections);
395        }
396
397        previous_state
398    }
399
400    pub fn set_block_state(&mut self, pos: &ChunkBlockPos, state: BlockState, min_y: i32) {
401        let section_index = section_index(pos.y, min_y);
402        let Some(section) = self.sections.get_mut(section_index as usize) else {
403            warn!(
404                "Tried to set block state {state:?} at out-of-bounds relative chunk position {pos:?}",
405            );
406            return;
407        };
408        let chunk_section_pos = ChunkSectionBlockPos::from(pos);
409        section.set_block_state(chunk_section_pos, state);
410
411        for heightmap in self.heightmaps.values_mut() {
412            heightmap.update(pos, state, &self.sections);
413        }
414    }
415
416    /// Get the biome at the given position, or `None` if it's out of bounds.
417    pub fn get_biome(&self, pos: ChunkBiomePos, min_y: i32) -> Option<Biome> {
418        if pos.y < min_y {
419            // y position is out of bounds
420            return None;
421        }
422        let section_index = section_index(pos.y, min_y);
423        let Some(section) = self.sections.get(section_index as usize) else {
424            warn!("Tried to get biome at out-of-bounds relative chunk position {pos:?}",);
425            return None;
426        };
427        let chunk_section_pos = ChunkSectionBiomePos::from(pos);
428        Some(section.get_biome(chunk_section_pos))
429    }
430}
431
432/// Get the block state at the given position from a list of sections. Returns
433/// `None` if the position is out of bounds.
434#[inline]
435pub fn get_block_state_from_sections(
436    sections: &[Section],
437    pos: &ChunkBlockPos,
438    min_y: i32,
439) -> Option<BlockState> {
440    if pos.y < min_y {
441        // y position is out of bounds
442        return None;
443    }
444    let section_index = section_index(pos.y, min_y) as usize;
445    if section_index >= sections.len() {
446        // y position is out of bounds
447        return None;
448    };
449    let section = &sections[section_index];
450    let chunk_section_pos = ChunkSectionBlockPos::from(pos);
451    Some(section.get_block_state(chunk_section_pos))
452}
453
454impl AzaleaWrite for Chunk {
455    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
456        for section in &self.sections {
457            section.azalea_write(buf)?;
458        }
459        Ok(())
460    }
461}
462
463impl Debug for PartialChunkStorage {
464    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
465        f.debug_struct("PartialChunkStorage")
466            .field("view_center", &self.view_center)
467            .field("chunk_radius", &self.chunk_radius)
468            .field("view_range", &self.view_range)
469            // .field("chunks", &self.chunks)
470            .field("chunks", &format_args!("{} items", self.chunks.len()))
471            .finish()
472    }
473}
474
475impl AzaleaRead for Section {
476    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
477        let block_count = u16::azalea_read(buf)?;
478
479        // this is commented out because the vanilla server is wrong
480        // TODO: ^ this comment was written ages ago. needs more investigation.
481        // assert!(
482        //     block_count <= 16 * 16 * 16,
483        //     "A section has more blocks than what should be possible. This is a bug!"
484        // );
485
486        let states = PalettedContainer::<BlockState>::read(buf)?;
487
488        for i in 0..states.storage.size() {
489            if !BlockState::is_valid_state(states.storage.get(i) as BlockStateIntegerRepr) {
490                return Err(BufReadError::Custom(format!(
491                    "Invalid block state {} (index {i}) found in section.",
492                    states.storage.get(i)
493                )));
494            }
495        }
496
497        let biomes = PalettedContainer::<Biome>::read(buf)?;
498        Ok(Section {
499            block_count,
500            states,
501            biomes,
502        })
503    }
504}
505
506impl AzaleaWrite for Section {
507    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
508        self.block_count.azalea_write(buf)?;
509        self.states.azalea_write(buf)?;
510        self.biomes.azalea_write(buf)?;
511        Ok(())
512    }
513}
514
515impl Section {
516    pub fn get_block_state(&self, pos: ChunkSectionBlockPos) -> BlockState {
517        self.states.get(pos)
518    }
519    pub fn get_and_set_block_state(
520        &mut self,
521        pos: ChunkSectionBlockPos,
522        state: BlockState,
523    ) -> BlockState {
524        self.states.get_and_set(pos, state)
525    }
526    pub fn set_block_state(&mut self, pos: ChunkSectionBlockPos, state: BlockState) {
527        self.states.set(pos, state);
528    }
529
530    pub fn get_biome(&self, pos: ChunkSectionBiomePos) -> Biome {
531        self.biomes.get(pos)
532    }
533    pub fn set_biome(&mut self, pos: ChunkSectionBiomePos, biome: Biome) {
534        self.biomes.set(pos, biome);
535    }
536    pub fn get_and_set_biome(&mut self, pos: ChunkSectionBiomePos, biome: Biome) -> Biome {
537        self.biomes.get_and_set(pos, biome)
538    }
539}
540
541impl Default for PartialChunkStorage {
542    fn default() -> Self {
543        Self::new(8)
544    }
545}
546impl Default for ChunkStorage {
547    fn default() -> Self {
548        Self::new(384, -64)
549    }
550}
551
552/// Get the index of where a section is in a chunk based on its y coordinate
553/// and the minimum y coordinate of the world.
554#[inline]
555pub fn section_index(y: i32, min_y: i32) -> u32 {
556    if y < min_y {
557        #[cfg(debug_assertions)]
558        warn!("y ({y}) must be at least {min_y}");
559        #[cfg(not(debug_assertions))]
560        trace!("y ({y}) must be at least {min_y}")
561    };
562    let min_section_index = min_y >> 4;
563    ((y >> 4) - min_section_index) as u32
564}
565
566#[cfg(test)]
567mod tests {
568    use super::*;
569
570    #[test]
571    fn test_section_index() {
572        assert_eq!(section_index(0, 0), 0);
573        assert_eq!(section_index(128, 0), 8);
574        assert_eq!(section_index(127, 0), 7);
575        assert_eq!(section_index(0, -64), 4);
576        assert_eq!(section_index(-64, -64), 0);
577        assert_eq!(section_index(-49, -64), 0);
578        assert_eq!(section_index(-48, -64), 1);
579        assert_eq!(section_index(128, -64), 12);
580    }
581
582    #[test]
583    fn test_out_of_bounds_y() {
584        let mut chunk_storage = ChunkStorage::default();
585        let mut partial_chunk_storage = PartialChunkStorage::default();
586        partial_chunk_storage.set(
587            &ChunkPos { x: 0, z: 0 },
588            Some(Chunk::default()),
589            &mut chunk_storage,
590        );
591        assert!(
592            chunk_storage
593                .get_block_state(BlockPos { x: 0, y: 319, z: 0 })
594                .is_some()
595        );
596        assert!(
597            chunk_storage
598                .get_block_state(BlockPos { x: 0, y: 320, z: 0 })
599                .is_none()
600        );
601        assert!(
602            chunk_storage
603                .get_block_state(BlockPos { x: 0, y: 338, z: 0 })
604                .is_none()
605        );
606        assert!(
607            chunk_storage
608                .get_block_state(BlockPos { x: 0, y: -64, z: 0 })
609                .is_some()
610        );
611        assert!(
612            chunk_storage
613                .get_block_state(BlockPos { x: 0, y: -65, z: 0 })
614                .is_none()
615        );
616    }
617
618    #[test]
619    fn test_chunk_pos_from_index() {
620        let mut partial_chunk_storage = PartialChunkStorage::new(5);
621        partial_chunk_storage.update_view_center(ChunkPos::new(0, -1));
622        assert_eq!(
623            partial_chunk_storage.chunk_pos_from_index(
624                partial_chunk_storage.index_from_chunk_pos(&ChunkPos::new(2, -1))
625            ),
626            ChunkPos::new(2, -1),
627        );
628    }
629}