heightmaps

4 files changed, 225 insertions, 17 deletions

diff --git a/azalea-world/src/bit_storage.rs b/azalea-world/src/bit_storage.rs
index 09b68fae..9f9b7abf 100755
--- a/azalea-world/src/bit_storage.rs
+++ b/azalea-world/src/bit_storage.rs
@@ -106,7 +106,7 @@ impl BitStorage {
             // 0 bit storage
             if data.is_empty() {
                 return Ok(BitStorage {
-                    data: Vec::with_capacity(0),
+                    data: Vec::new(),
                     bits,
                     size,
                     ..Default::default()
diff --git a/azalea-world/src/chunk_storage.rs b/azalea-world/src/chunk_storage.rs
index 133d522b..8d894d9a 100755
--- a/azalea-world/src/chunk_storage.rs
+++ b/azalea-world/src/chunk_storage.rs
@@ -1,10 +1,14 @@
+use crate::heightmap::Heightmap;
+use crate::heightmap::HeightmapKind;
 use crate::palette::PalettedContainer;
 use crate::palette::PalettedContainerKind;
 use azalea_block::BlockState;
 use azalea_buf::{BufReadError, McBufReadable, McBufWritable};
 use azalea_core::{BlockPos, ChunkBlockPos, ChunkPos, ChunkSectionBlockPos};
+use azalea_nbt::NbtCompound;
 use log::{debug, trace, warn};
 use parking_lot::RwLock;
+use std::str::FromStr;
 use std::{
     collections::HashMap,
     fmt::Debug,
@@ -43,6 +47,10 @@ pub struct ChunkStorage {
 #[derive(Debug)]
 pub struct Chunk {
     pub sections: Vec<Section>,
+    /// Heightmaps are used for identifying the surface blocks in a chunk.
+    /// Usually for clients only `WorldSurface` and `MotionBlocking` are
+    /// present.
+    pub heightmaps: HashMap<HeightmapKind, Heightmap>,
 }
 
 /// A section of a chunk, i.e. a 16*16*16 block area.
@@ -73,6 +81,7 @@ impl Default for Chunk {
     fn default() -> Self {
         Chunk {
             sections: vec![Section::default(); (384 / 16) as usize],
+            heightmaps: HashMap::new(),
         }
     }
 }
@@ -119,6 +128,7 @@ impl PartialChunkStorage {
         &mut self,
         pos: &ChunkPos,
         data: &mut Cursor<&[u8]>,
+        heightmaps: &NbtCompound,
         chunk_storage: &mut ChunkStorage,
     ) -> Result<(), BufReadError> {
         debug!("Replacing chunk at {:?}", pos);
@@ -127,7 +137,12 @@ impl PartialChunkStorage {
             return Ok(());
         }
 
-        let chunk = Chunk::read_with_dimension_height(data, chunk_storage.height)?;
+        let chunk = Chunk::read_with_dimension_height(
+            data,
+            chunk_storage.height,
+            chunk_storage.min_y,
+            heightmaps,
+        )?;
 
         trace!("Loaded chunk {:?}", pos);
         self.set(pos, Some(chunk), chunk_storage);
@@ -229,6 +244,8 @@ impl Chunk {
     pub fn read_with_dimension_height(
         buf: &mut Cursor<&[u8]>,
         dimension_height: u32,
+        min_y: i32,
+        heightmaps_nbt: &NbtCompound,
     ) -> Result<Self, BufReadError> {
         let section_count = dimension_height / SECTION_HEIGHT;
         let mut sections = Vec::with_capacity(section_count as usize);
@@ -236,23 +253,30 @@ impl Chunk {
             let section = Section::read_from(buf)?;
             sections.push(section);
         }
-        Ok(Chunk { sections })
+
+        let mut heightmaps = HashMap::new();
+        for (name, heightmap) in heightmaps_nbt.iter() {
+            let Ok(kind) = HeightmapKind::from_str(name) else {
+                warn!("Unknown heightmap kind: {name}");
+                continue;
+            };
+            let Some(data) = heightmap.as_long_array() else {
+                warn!("Heightmap {name} is not a long array");
+                continue;
+            };
+            let data: Vec<u64> = data.iter().map(|x| *x as u64).collect();
+            let heightmap = Heightmap::new(kind, dimension_height, min_y, data);
+            heightmaps.insert(kind, heightmap);
+        }
+
+        Ok(Chunk {
+            sections,
+            heightmaps,
+        })
     }
 
     pub fn get(&self, pos: &ChunkBlockPos, min_y: i32) -> Option<BlockState> {
-        if pos.y < min_y {
-            // y position is out of bounds
-            return None;
-        }
-        let section_index = section_index(pos.y, min_y) as usize;
-        if section_index >= self.sections.len() {
-            // y position is out of bounds
-            return None;
-        };
-        // TODO: make sure the section exists
-        let section = &self.sections[section_index];
-        let chunk_section_pos = ChunkSectionBlockPos::from(pos);
-        Some(section.get(chunk_section_pos))
+        get_block_state_from_sections(&self.sections, pos, min_y)
     }
 
     pub fn get_and_set(
@@ -265,7 +289,13 @@ impl Chunk {
         // TODO: make sure the section exists
         let section = &mut self.sections[section_index as usize];
         let chunk_section_pos = ChunkSectionBlockPos::from(pos);
-        section.get_and_set(chunk_section_pos, state)
+        let previous_state = section.get_and_set(chunk_section_pos, state);
+
+        for heightmap in self.heightmaps.values_mut() {
+            heightmap.update(pos, state, &self.sections);
+        }
+
+        previous_state
     }
 
     pub fn set(&mut self, pos: &ChunkBlockPos, state: BlockState, min_y: i32) {
@@ -274,7 +304,33 @@ impl Chunk {
         let section = &mut self.sections[section_index as usize];
         let chunk_section_pos = ChunkSectionBlockPos::from(pos);
         section.set(chunk_section_pos, state);
+
+        for heightmap in self.heightmaps.values_mut() {
+            heightmap.update(pos, state, &self.sections);
+        }
+    }
+}
+
+/// Get the block state at the given position from a list of sections. Returns
+/// `None` if the position is out of bounds.
+pub fn get_block_state_from_sections(
+    sections: &[Section],
+    pos: &ChunkBlockPos,
+    min_y: i32,
+) -> Option<BlockState> {
+    if pos.y < min_y {
+        // y position is out of bounds
+        return None;
     }
+    let section_index = section_index(pos.y, min_y) as usize;
+    if section_index >= sections.len() {
+        // y position is out of bounds
+        return None;
+    };
+    // TODO: make sure the section exists
+    let section = &sections[section_index];
+    let chunk_section_pos = ChunkSectionBlockPos::from(pos);
+    Some(section.get(chunk_section_pos))
 }
 
 impl McBufWritable for Chunk {
diff --git a/azalea-world/src/heightmap.rs b/azalea-world/src/heightmap.rs
new file mode 100644
index 00000000..ec73adf9
--- /dev/null
+++ b/azalea-world/src/heightmap.rs
@@ -0,0 +1,151 @@
+use std::{fmt::Display, str::FromStr};
+
+use azalea_block::BlockState;
+use azalea_core::{math, ChunkBlockPos};
+use azalea_registry::tags::blocks::LEAVES;
+
+use crate::{chunk_storage::get_block_state_from_sections, BitStorage, Section};
+
+// (wg stands for worldgen)
+
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
+pub enum HeightmapKind {
+    WorldSurfaceWg,
+    WorldSurface,
+    OceanFloorWg,
+    OceanFloor,
+    MotionBlocking,
+    MotionBlockingNoLeaves,
+}
+
+#[derive(Clone, Debug)]
+pub struct Heightmap {
+    pub data: BitStorage,
+    pub min_y: i32,
+    pub kind: HeightmapKind,
+}
+
+fn blocks_motion(block_state: BlockState) -> bool {
+    // TODO
+    !block_state.is_air()
+}
+
+fn motion_blocking(block_state: BlockState) -> bool {
+    // TODO
+    !block_state.is_air() || block_state.waterlogged()
+}
+
+impl HeightmapKind {
+    pub fn is_opaque(self, block_state: BlockState) -> bool {
+        let block = Box::<dyn azalea_block::Block>::from(block_state);
+        let registry_block = block.as_registry_block();
+        match self {
+            HeightmapKind::WorldSurfaceWg => !block_state.is_air(),
+            HeightmapKind::WorldSurface => !block_state.is_air(),
+            HeightmapKind::OceanFloorWg => blocks_motion(block_state),
+            HeightmapKind::OceanFloor => blocks_motion(block_state),
+            HeightmapKind::MotionBlocking => motion_blocking(block_state),
+            HeightmapKind::MotionBlockingNoLeaves => {
+                motion_blocking(block_state) && !LEAVES.contains(&registry_block)
+            }
+        }
+    }
+}
+
+impl Heightmap {
+    pub fn new(kind: HeightmapKind, dimension_height: u32, min_y: i32, data: Vec<u64>) -> Self {
+        let bits = math::ceil_log2(dimension_height + 1);
+        let data = BitStorage::new(bits as usize, 16 * 16, Some(data)).unwrap();
+        Self { kind, data, min_y }
+    }
+
+    pub fn get_index(x: u8, z: u8) -> usize {
+        (x as usize) + (z as usize) * 16
+    }
+
+    pub fn get_first_available_at_index(&self, index: usize) -> i32 {
+        self.data.get(index) as i32 + self.min_y
+    }
+
+    pub fn get_first_available(&self, x: u8, z: u8) -> i32 {
+        self.get_first_available_at_index(Self::get_index(x, z))
+    }
+
+    pub fn get_highest_taken(&self, x: u8, z: u8) -> i32 {
+        self.get_first_available(x, z) - 1
+    }
+
+    pub fn set_height(&mut self, x: u8, z: u8, height: i32) {
+        self.data
+            .set(Self::get_index(x, z), (height - self.min_y) as u64);
+    }
+
+    /// Updates the heightmap with the given block state at the given position.
+    pub fn update(
+        &mut self,
+        pos: &ChunkBlockPos,
+        block_state: BlockState,
+        sections: &[Section],
+    ) -> bool {
+        let first_available_y = self.get_first_available(pos.x, pos.z);
+        if pos.y <= first_available_y - 2 {
+            return false;
+        }
+        if self.kind.is_opaque(block_state) {
+            // increase y
+            if pos.y >= first_available_y {
+                self.set_height(pos.x, pos.z, pos.y + 1);
+                return true;
+            }
+        } else if first_available_y - 1 == pos.y {
+            // decrease y
+            for y in (self.min_y..pos.y).rev() {
+                if self.kind.is_opaque(
+                    get_block_state_from_sections(
+                        sections,
+                        &ChunkBlockPos::new(pos.x, y, pos.z),
+                        self.min_y,
+                    )
+                    .unwrap_or_default(),
+                ) {
+                    self.set_height(pos.x, pos.z, y + 1);
+                    return true;
+                }
+            }
+
+            self.set_height(pos.x, pos.z, self.min_y);
+            return true;
+        }
+
+        false
+    }
+}
+
+impl FromStr for HeightmapKind {
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "WORLD_SURFACE_WG" => Ok(HeightmapKind::WorldSurfaceWg),
+            "WORLD_SURFACE" => Ok(HeightmapKind::WorldSurface),
+            "OCEAN_FLOOR_WG" => Ok(HeightmapKind::OceanFloorWg),
+            "OCEAN_FLOOR" => Ok(HeightmapKind::OceanFloor),
+            "MOTION_BLOCKING" => Ok(HeightmapKind::MotionBlocking),
+            "MOTION_BLOCKING_NO_LEAVES" => Ok(HeightmapKind::MotionBlockingNoLeaves),
+            _ => Err(()),
+        }
+    }
+}
+
+impl Display for HeightmapKind {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            HeightmapKind::WorldSurfaceWg => write!(f, "WORLD_SURFACE_WG"),
+            HeightmapKind::WorldSurface => write!(f, "WORLD_SURFACE"),
+            HeightmapKind::OceanFloorWg => write!(f, "OCEAN_FLOOR_WG"),
+            HeightmapKind::OceanFloor => write!(f, "OCEAN_FLOOR"),
+            HeightmapKind::MotionBlocking => write!(f, "MOTION_BLOCKING"),
+            HeightmapKind::MotionBlockingNoLeaves => write!(f, "MOTION_BLOCKING_NO_LEAVES"),
+        }
+    }
+}
diff --git a/azalea-world/src/lib.rs b/azalea-world/src/lib.rs
index adaafa1f..0d80f75d 100644
--- a/azalea-world/src/lib.rs
+++ b/azalea-world/src/lib.rs
@@ -5,6 +5,7 @@
 mod bit_storage;
 mod chunk_storage;
 mod container;
+pub mod heightmap;
 pub mod iterators;
 pub mod palette;
 mod world;