symbios_ground/

heightmap.rs

use std::sync::OnceLock;

use serde::{Deserialize, Serialize};

/// A 2D heightmap stored as a flat row-major `Vec<f32>` buffer.
///
/// Covers world space `[0, width * scale) × [0, height * scale)`.
/// `width` and `height` are grid-cell counts; `scale` is world units per cell.
///
/// The invariant `data.len() == width * height` is enforced by all constructors
/// and mutation methods; fields are private to prevent external corruption.
///
/// A grid-cell normal cache is computed lazily on first read and reused across
/// repeated queries (e.g. by [`SplatMapper`](crate::SplatMapper)). It is
/// invalidated whenever the heights change via [`HeightMap::set`],
/// [`HeightMap::get_mut`], [`HeightMap::data_mut`], or [`HeightMap::normalize`].
#[derive(Debug, Serialize, Deserialize)]
pub struct HeightMap {
    data: Vec<f32>,
    width: usize,
    height: usize,
    scale: f32,
    /// Pooled lakes detected by hydraulic erosion. Empty until erosion writes
    /// to it; persisted across serialisation so renderers can visualise water
    /// without re-running the simulation.
    #[serde(default)]
    lakes: Vec<Lake>,
    /// Lazy per-grid-cell central-difference normals. Skipped from serde so the
    /// cache stays consistent with the height data after deserialisation.
    #[serde(skip)]
    normals: OnceLock<Vec<[f32; 3]>>,
}

/// A pooled body of standing water produced by [`HydraulicErosion`](crate::HydraulicErosion).
///
/// Built up at points where droplet velocity drops below the configured
/// threshold while still above `water_level`, marking realistic basins and
/// depressions in the heightmap.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct Lake {
    /// Row-major index into the heightmap (`z * width + x`).
    pub index: usize,
    /// Accumulated water depth at this cell.
    pub depth: f32,
    /// World-space area attributed to this lake cell — `scale * scale`.
    pub area: f32,
}

impl Clone for HeightMap {
    fn clone(&self) -> Self {
        Self {
            data: self.data.clone(),
            width: self.width,
            height: self.height,
            scale: self.scale,
            lakes: self.lakes.clone(),
            normals: OnceLock::new(),
        }
    }
}

impl HeightMap {
    /// Create a heightmap of `width × height` cells, all initialised to `0.0`.
    ///
    /// `scale` is the world-unit size of each cell (must be `> 0`).
    ///
    /// # Panics
    ///
    /// Panics if `width == 0`, `height == 0`, or `scale <= 0.0`.
    pub fn new(width: usize, height: usize, scale: f32) -> Self {
        assert!(width > 0 && height > 0, "dimensions must be positive");
        assert!(scale > 0.0, "scale must be positive");
        Self {
            data: vec![0.0; width * height],
            width,
            height,
            scale,
            lakes: Vec::new(),
            normals: OnceLock::new(),
        }
    }

    /// Grid-cell width.
    #[inline]
    pub fn width(&self) -> usize {
        self.width
    }

    /// Grid-cell depth.
    #[inline]
    pub fn height(&self) -> usize {
        self.height
    }

    /// World units per grid cell.
    #[inline]
    pub fn scale(&self) -> f32 {
        self.scale
    }

    /// Read-only view of the flat row-major height buffer.
    #[inline]
    pub fn data(&self) -> &[f32] {
        &self.data
    }

    /// Mutable view of the flat row-major height buffer.
    ///
    /// The caller must not change the slice length; only element values may be
    /// modified.  Dimensions (`width`, `height`) are unaffected. Invalidates
    /// the normal cache.
    #[inline]
    pub fn data_mut(&mut self) -> &mut [f32] {
        self.invalidate_caches();
        &mut self.data
    }

    /// Return the height at grid cell `(x, z)`.
    ///
    /// # Panics
    ///
    /// Panics if `x >= width` or `z >= height`.
    #[inline]
    pub fn get(&self, x: usize, z: usize) -> f32 {
        self.data[z * self.width + x]
    }

    /// Return a mutable reference to the height at grid cell `(x, z)`.
    /// Invalidates the normal cache.
    ///
    /// # Panics
    ///
    /// Panics if `x >= width` or `z >= height`.
    #[inline]
    pub fn get_mut(&mut self, x: usize, z: usize) -> &mut f32 {
        self.invalidate_caches();
        &mut self.data[z * self.width + x]
    }

    /// Set the height at grid cell `(x, z)` to `val`. Invalidates the normal
    /// cache.
    ///
    /// # Panics
    ///
    /// Panics if `x >= width` or `z >= height`.
    #[inline]
    pub fn set(&mut self, x: usize, z: usize, val: f32) {
        self.invalidate_caches();
        self.data[z * self.width + x] = val;
    }

    /// Return the height at grid cell `(x, z)`, clamping coordinates to the
    /// valid range instead of panicking on out-of-bounds indices.
    #[inline]
    pub fn get_clamped(&self, x: i32, z: i32) -> f32 {
        let cx = x.clamp(0, self.width as i32 - 1) as usize;
        let cz = z.clamp(0, self.height as i32 - 1) as usize;
        self.get(cx, cz)
    }

    /// Sample height at world position using bilinear interpolation.
    /// Clamps to heightmap boundaries.
    pub fn get_height_at(&self, world_x: f32, world_z: f32) -> f32 {
        let gx = world_x / self.scale;
        let gz = world_z / self.scale;

        let x0 = gx.floor() as i32;
        let z0 = gz.floor() as i32;
        let fx = gx - x0 as f32;
        let fz = gz - z0 as f32;

        let h00 = self.get_clamped(x0, z0);
        let h10 = self.get_clamped(x0 + 1, z0);
        let h01 = self.get_clamped(x0, z0 + 1);
        let h11 = self.get_clamped(x0 + 1, z0 + 1);

        let h0 = h00 + (h10 - h00) * fx;
        let h1 = h01 + (h11 - h01) * fx;
        h0 + (h1 - h0) * fz
    }

    /// Compute surface normal at world position using central differences.
    /// Returns a normalized `[x, y, z]` vector where `y` is up.
    ///
    /// `scale` is always > 0 (enforced by the constructor), so the `2*scale`
    /// divisor and the `len` (≥ 1.0 since ny = 1) are both safe.
    pub fn get_normal_at(&self, world_x: f32, world_z: f32) -> [f32; 3] {
        let step = self.scale;
        let hl = self.get_height_at(world_x - step, world_z);
        let hr = self.get_height_at(world_x + step, world_z);
        let hd = self.get_height_at(world_x, world_z - step);
        let hu = self.get_height_at(world_x, world_z + step);

        let dhdx = (hr - hl) / (2.0 * step);
        let dhdz = (hu - hd) / (2.0 * step);

        let nx = -dhdx;
        let ny = 1.0_f32;
        let nz = -dhdz;
        let len = (nx * nx + ny * ny + nz * nz).sqrt();
        // If scale is a denormal (< ~1e-38), dhdx/dhdz can overflow to ±INF,
        // making len = INF and the division NaN. Return a flat up-normal instead.
        if !len.is_finite() {
            return [0.0, 1.0, 0.0];
        }
        [nx / len, ny / len, nz / len]
    }

    /// Cached central-difference normal at grid cell `(x, z)`.
    ///
    /// Equivalent to `get_normal_at(x as f32 * scale, z as f32 * scale)` for
    /// in-bounds cells, but reads from a lazily populated cache so callers that
    /// scan the whole grid (e.g. [`SplatMapper`](crate::SplatMapper)) avoid
    /// recomputing four central differences per pixel.
    ///
    /// # Panics
    ///
    /// Panics if `x >= width` or `z >= height`.
    pub fn normal_at_grid(&self, x: usize, z: usize) -> [f32; 3] {
        self.normals_grid()[z * self.width + x]
    }

    /// Lazily populated per-grid-cell normal table; row-major, length
    /// `width * height`.
    pub fn normals_grid(&self) -> &[[f32; 3]] {
        self.normals.get_or_init(|| self.compute_normals())
    }

    fn compute_normals(&self) -> Vec<[f32; 3]> {
        let w = self.width;
        let h = self.height;
        let mut out = Vec::with_capacity(w * h);
        let step = self.scale;
        let inv_2step = 1.0_f32 / (2.0 * step);
        for z in 0..h {
            for x in 0..w {
                let hl = self.get_clamped(x as i32 - 1, z as i32);
                let hr = self.get_clamped(x as i32 + 1, z as i32);
                let hd = self.get_clamped(x as i32, z as i32 - 1);
                let hu = self.get_clamped(x as i32, z as i32 + 1);

                let dhdx = (hr - hl) * inv_2step;
                let dhdz = (hu - hd) * inv_2step;

                let nx = -dhdx;
                let ny = 1.0_f32;
                let nz = -dhdz;
                let len = (nx * nx + ny * ny + nz * nz).sqrt();
                if !len.is_finite() {
                    out.push([0.0, 1.0, 0.0]);
                } else {
                    out.push([nx / len, ny / len, nz / len]);
                }
            }
        }
        out
    }

    /// Lakes detected by the most recent hydraulic erosion pass. Empty until
    /// [`HydraulicErosion::erode`](crate::HydraulicErosion::erode) populates
    /// it.
    pub fn lakes(&self) -> &[Lake] {
        &self.lakes
    }

    /// Internal: replace the lake list (used by erosion).
    pub(crate) fn set_lakes(&mut self, lakes: Vec<Lake>) {
        self.lakes = lakes;
    }

    /// Internal: clear all derived caches that depend on `data`.
    fn invalidate_caches(&mut self) {
        self.normals.take();
    }

    /// Normalize all height values to `[0.0, 1.0]`.
    pub fn normalize(&mut self) {
        let min = self.data.iter().cloned().fold(f32::INFINITY, f32::min);
        let max = self.data.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let range = max - min;
        if range > f32::EPSILON {
            for v in &mut self.data {
                *v = (*v - min) / range;
            }
        }
        self.invalidate_caches();
    }

    /// World-space width of the heightmap.
    pub fn world_width(&self) -> f32 {
        self.width as f32 * self.scale
    }

    /// World-space depth of the heightmap.
    pub fn world_depth(&self) -> f32 {
        self.height as f32 * self.scale
    }
}