terrain_codec/

terrain.rs

//! One-shot heightmap → quantized-mesh (`.terrain`) encoding.
//!
//! This module ties together the three crates that otherwise have to be
//! wired up by hand:
//!
//! 1. [`martini`] generates an adaptive RTIN mesh from the elevation grid.
//! 2. The mesh's `(u, v, height)` are quantised to the 0..=32767 range.
//! 3. [`quantized_mesh`] encodes the header, vertices, edge indices and
//!    optional extensions into the quantized-mesh-1.0 byte stream.
//!
//! The fiddly bits it handles for you:
//!
//! - Re-sampling each mesh vertex's height (martini discards heights once
//!   the error pyramid is built, so the transform has to recover the grid
//!   coordinate from `(u, v)` and read the DEM again).
//! - Computing the encoded height range from the *mesh* vertices (what is
//!   actually stored), not the full grid.
//! - Streaming the mesh vertices through
//!   [`QuantizedMeshHeader::from_bounds_with_vertices_iter`] for a tight
//!   horizon-occlusion point.
//! - Vertex normals via the [`NormalMode`] of your choice.
//! - Folding the ellipsoid [`curvature_bulge`] into the error pyramid so
//!   nearly-flat tiles keep enough triangles to track the globe at low zoom
//!   (the bulge feeds error estimation only, never the emitted heights).
//!
//! # Grid orientation
//!
//! `elevations` (and the `get_height` closure's `y`) are **row-major,
//! north → south**: row `0` is the northern edge, row `grid_size - 1` the
//! southern edge. This matches [`crate::normals::BufferedElevations`], so a
//! buffered grid can be reused directly for [`NormalMode::BufferedGradient`].
//!
//! # Seamless tiling — the caller supplies the halo
//!
//! These functions encode **one tile in isolation**; they never fetch
//! neighbouring tiles. For gap-free, seam-free output the *caller* must
//! widen the input to overlap the neighbours — fetch the halo cells along
//! with the tile and stitch them in before calling:
//!
//! - **Geometry seam.** martini needs a `2^n + 1` grid, so an `N`-post DEM
//!   tile needs one extra post on its east and south edges. That `+1` post
//!   is the neighbour tile's *first* post for the shared edge — read it
//!   from the neighbour, don't edge-replicate, or adjacent tiles won't
//!   agree on the boundary and the globe cracks along tile seams.
//! - **Normal seam.** [`NormalMode::BufferedGradient`] needs a
//!   `buffer`-cell halo of neighbour samples on **every** side (see
//!   [`BufferedElevations`]). Edge vertices read their `±1` neighbours out
//!   of that halo, so the same physical edge gets identical normals from
//!   either tile and lighting stays continuous.
//!
//! Gathering that neighbour data (over HTTP, from disk, from a cache, …) is
//! deliberately left to the caller — hence this module takes an
//! already-assembled grid rather than fetching tiles itself, which also
//! keeps it free of any async/runtime assumptions.
//!
//! # Example
//!
//! ```
//! use terrain_codec::quantized_mesh::TileBounds;
//! use terrain_codec::terrain::{encode_terrain, TerrainOptions};
//!
//! let grid_size = 65; // 2^6 + 1
//! let elevations = vec![0.0f32; (grid_size * grid_size) as usize];
//! let bounds = TileBounds::new(139.0, 35.0, 139.01, 35.01);
//!
//! let terrain: Vec<u8> = encode_terrain(
//!     &elevations,
//!     grid_size,
//!     &bounds,
//!     &TerrainOptions {
//!         max_error: 1.0,
//!         ..Default::default()
//!     },
//! );
//! assert!(terrain.starts_with(&[0x1f, 0x8b])); // gzip magic (default level 6)
//! ```

use std::io::{self, Write};

use martini::Martini;
use quantized_mesh::{
    EdgeIndices, EncodeOptions, QUANTIZED_MAX, QuantizedMeshEncoder, QuantizedMeshHeader,
    QuantizedVertices, TileBounds, TileMetadata, WaterMask,
};

use crate::normals::{BufferedElevations, buffered_gradient_normals, face_normals};

/// How (and whether) per-vertex normals are computed for the oct-encoded
/// vertex-normals extension.
#[derive(Debug, Clone, Default)]
pub enum NormalMode {
    /// No vertex-normals extension.
    #[default]
    None,
    /// Per-tile face normals ([`crate::normals::face_normals`]). Simple, but
    /// produces a visible shading seam at tile boundaries.
    FaceNormals,
    /// Seam-free DEM-gradient normals
    /// ([`crate::normals::buffered_gradient_normals`]) sampled from a
    /// buffer-extended grid. Its `tile_grid_size` must equal the encode
    /// `grid_size`.
    ///
    /// The caller is responsible for filling the `buffer`-cell halo around
    /// the tile with the **neighbour tiles'** elevations — that overlap is
    /// what makes edge normals match across the seam. A halo filled by
    /// edge-replication still encodes fine, but won't be seam-free.
    BufferedGradient(BufferedElevations),
}

/// Options controlling [`encode_terrain`] and the other encode functions in
/// this module.
#[derive(Debug, Clone)]
pub struct TerrainOptions {
    /// Maximum RTIN error threshold in metres. Lower values keep more
    /// triangles (higher fidelity, larger output).
    pub max_error: f64,
    /// Gzip compression level: `0` emits uncompressed bytes, `1..=9` gzip at
    /// that level. Defaults to `6`.
    pub compression_level: u32,
    /// Vertex-normal strategy.
    pub normals: NormalMode,
    /// Optional water-mask extension.
    pub water_mask: Option<WaterMask>,
    /// Optional metadata (child-tile availability) extension.
    pub metadata: Option<TileMetadata>,
}

impl Default for TerrainOptions {
    fn default() -> Self {
        Self {
            max_error: 1.0,
            compression_level: 6,
            normals: NormalMode::None,
            water_mask: None,
            metadata: None,
        }
    }
}

/// Encode a heightmap to a quantized-mesh `.terrain` byte vector, sampling
/// elevations through a closure.
///
/// `get_height(x, y)` returns the elevation in metres at grid column `x`
/// (`0..grid_size`, west → east) and row `y` (`0..grid_size`, north →
/// south).
///
/// This is the primitive form; [`encode_terrain`] wraps it for a flat
/// `&[f32]` grid. `get_height` is called twice per grid vertex that ends up
/// in the mesh (once while building the error pyramid, once to recover the
/// stored height), so keep it cheap or memoised.
///
/// # Panics
///
/// Panics if `grid_size` is not `2^n + 1`, or — for
/// [`NormalMode::BufferedGradient`] — if the buffered grid's
/// `tile_grid_size` does not equal `grid_size`.
pub fn encode_terrain_from_fn<F>(
    grid_size: u32,
    bounds: &TileBounds,
    get_height: F,
    options: &TerrainOptions,
) -> Vec<u8>
where
    F: Fn(u32, u32) -> f64,
{
    let (encoder, encode_opts) = build(grid_size, bounds, get_height, options);
    encoder.encode_with_options(&encode_opts)
}

/// Like [`encode_terrain_from_fn`], but streams the encoded bytes to a
/// writer instead of allocating a `Vec`.
///
/// # Panics
///
/// Same panics as [`encode_terrain_from_fn`].
pub fn encode_terrain_from_fn_to<F, W>(
    grid_size: u32,
    bounds: &TileBounds,
    get_height: F,
    options: &TerrainOptions,
    writer: W,
) -> io::Result<()>
where
    F: Fn(u32, u32) -> f64,
    W: Write,
{
    let (encoder, encode_opts) = build(grid_size, bounds, get_height, options);
    encoder.encode_to_with_options(writer, &encode_opts)
}

/// Encode a flat row-major (north → south) `f32` elevation grid to a
/// quantized-mesh `.terrain` byte vector.
///
/// `elevations.len()` must equal `grid_size * grid_size`.
///
/// # Panics
///
/// Panics if the length check fails, or for the panics listed on
/// [`encode_terrain_from_fn`].
pub fn encode_terrain(
    elevations: &[f32],
    grid_size: u32,
    bounds: &TileBounds,
    options: &TerrainOptions,
) -> Vec<u8> {
    assert_grid_len(elevations.len(), grid_size);
    let gs = grid_size as usize;
    encode_terrain_from_fn(
        grid_size,
        bounds,
        |x, y| elevations[y as usize * gs + x as usize] as f64,
        options,
    )
}

/// Like [`encode_terrain`], but streams the encoded bytes to a writer.
///
/// # Panics
///
/// Same panics as [`encode_terrain`].
pub fn encode_terrain_to<W: Write>(
    elevations: &[f32],
    grid_size: u32,
    bounds: &TileBounds,
    options: &TerrainOptions,
    writer: W,
) -> io::Result<()> {
    assert_grid_len(elevations.len(), grid_size);
    let gs = grid_size as usize;
    encode_terrain_from_fn_to(
        grid_size,
        bounds,
        |x, y| elevations[y as usize * gs + x as usize] as f64,
        options,
        writer,
    )
}

fn assert_grid_len(len: usize, grid_size: u32) {
    let expected = (grid_size as usize) * (grid_size as usize);
    assert_eq!(
        len, expected,
        "elevations length mismatch: expected {expected} ({grid_size}×{grid_size}), got {len}"
    );
}

/// Radial deviation (metres) of the WGS84 ellipsoid surface above the flat
/// bilinear interpolation of the tile's four corners, evaluated at grid cell
/// `(x, y)`. The encode functions in this module always add this to the height
/// field that drives martini's error pyramid (never to the stored heights), so
/// nearly-flat tiles still tessellate enough to track the globe's curvature
/// instead of collapsing to a flat quad that cuts under the ellipsoid at low
/// zoom. Exposed so callers can reason about / reproduce the subdivision.
///
/// Derivation: along one geodesic edge spanning angle `Δ`, the arc rises above
/// its chord by `R·(cos((t−½)Δ) − cos(Δ/2))`, which for the small `Δ` of a
/// tile is `≈ R·(Δ²/2)·t·(1−t)` — zero at the corners, peaking at the centre.
/// Because martini's interpolation is exact for affine fields, only this
/// non-linear `t·(1−t)` term contributes error, so it is precisely the signal
/// that controls subdivision. Longitude and latitude separate; longitude span
/// is scaled by `cos(lat)` for meridian convergence. The term shrinks with the
/// square of the tile span, so it forces dense meshes at low zoom (few, large
/// tiles) and fades to nothing at high zoom (negligible curvature per tile).
///
/// `x` / `y` are grid coordinates in `0..grid_size`. The result is zero on the
/// four tile corners.
pub fn curvature_bulge(x: u32, y: u32, grid_size: u32, bounds: &TileBounds) -> f64 {
    // WGS84 mean radius. Sub-metre accuracy here is irrelevant — this only
    // scales an error threshold, never an emitted height.
    const EARTH_RADIUS_M: f64 = 6_371_008.8;
    let n = (grid_size.saturating_sub(1)).max(1) as f64;
    let u = x as f64 / n;
    let v = y as f64 / n;
    let dlon = (bounds.east - bounds.west).to_radians();
    let dlat = (bounds.north - bounds.south).to_radians();
    let mid_lat = ((bounds.south + bounds.north) * 0.5).to_radians();
    let dlon_eff = dlon * mid_lat.cos();
    0.5 * EARTH_RADIUS_M * (dlat * dlat * v * (1.0 - v) + dlon_eff * dlon_eff * u * (1.0 - u))
}

/// Run martini, quantise the mesh, build the header + extensions, and return
/// a ready-to-encode [`QuantizedMeshEncoder`] alongside its [`EncodeOptions`].
fn build<F>(
    grid_size: u32,
    bounds: &TileBounds,
    get_height: F,
    options: &TerrainOptions,
) -> (QuantizedMeshEncoder, EncodeOptions)
where
    F: Fn(u32, u32) -> f64,
{
    if let NormalMode::BufferedGradient(buf) = &options.normals {
        assert_eq!(
            buf.tile_grid_size, grid_size,
            "BufferedGradient tile_grid_size ({}) must equal encode grid_size ({grid_size})",
            buf.tile_grid_size
        );
    }

    let mut martini = Martini::new(grid_size);
    let max = (grid_size - 1) as f64;
    // The error pyramid always sees the ellipsoid bulge so nearly-flat tiles
    // still subdivide enough to track the globe's curvature; the stored
    // heights below stay the true `get_height` values (the bulge never leaks
    // into emitted heights).
    let tile = martini.create_terrain(|x, y| {
        get_height(x as u32, y as u32) + curvature_bulge(x as u32, y as u32, grid_size, bounds)
    });

    // Hijack the UV transform to keep martini's `(u, v)` and re-sample the
    // height at the grid vertex. Martini computes `u = x/max` and
    // `v = 1 - y/max`, both exact for grid points, so the inverse recovers
    // the integer grid coordinate without drift.
    let (positions, indices, _uvs) =
        tile.construct_mesh(&mut martini, options.max_error, &mut |(u, v)| {
            let gx = (u * max).round();
            let gy = ((1.0 - v) * max).round();
            (u, v, get_height(gx as u32, gy as u32))
        });

    let vertex_count = positions.len() / 3;

    // Height range over the mesh vertices — i.e. exactly the heights we
    // quantise and store. A flat tile collapses to a zero span.
    let mut min_h = f64::INFINITY;
    let mut max_h = f64::NEG_INFINITY;
    for i in 0..vertex_count {
        let h = positions[i * 3 + 2] as f64;
        min_h = min_h.min(h);
        max_h = max_h.max(h);
    }
    if vertex_count == 0 {
        min_h = 0.0;
        max_h = 0.0;
    }
    let height_span = max_h - min_h;

    // Quantise (u, v, height) → 0..=32767.
    let quant_max = QUANTIZED_MAX as f64;
    let mut vertices = QuantizedVertices::with_capacity(vertex_count);
    for i in 0..vertex_count {
        let u = positions[i * 3] as f64;
        let v = positions[i * 3 + 1] as f64;
        let h = positions[i * 3 + 2] as f64;
        let uq = (u * quant_max).round().clamp(0.0, quant_max) as u16;
        let vq = (v * quant_max).round().clamp(0.0, quant_max) as u16;
        let hq = if height_span > 0.0 {
            (((h - min_h) / height_span) * quant_max)
                .round()
                .clamp(0.0, quant_max) as u16
        } else {
            0
        };
        vertices.push(uq, vq, hq);
    }

    let edge_indices = EdgeIndices::from_vertices(&vertices);

    // Feed the mesh vertices (geodetic) to the header so the horizon
    // occlusion point is as tight as possible.
    let lon_span = bounds.east - bounds.west;
    let lat_span = bounds.north - bounds.south;
    let geodetic = (0..vertex_count).map(|i| {
        let u = positions[i * 3] as f64;
        let v = positions[i * 3 + 1] as f64;
        let h = positions[i * 3 + 2] as f64;
        [bounds.west + u * lon_span, bounds.south + v * lat_span, h]
    });
    let header = QuantizedMeshHeader::from_bounds_with_vertices_iter(
        bounds,
        min_h as f32,
        max_h as f32,
        geodetic,
    );

    let normals = match &options.normals {
        NormalMode::None => None,
        NormalMode::FaceNormals => Some(face_normals(&vertices, &indices, bounds, min_h, max_h)),
        NormalMode::BufferedGradient(buf) => {
            Some(buffered_gradient_normals(&vertices, bounds, buf))
        }
    };

    let encode_opts = EncodeOptions {
        include_normals: normals.is_some(),
        normals,
        include_water_mask: options.water_mask.is_some(),
        water_mask: options.water_mask.clone(),
        include_metadata: options.metadata.is_some(),
        metadata: options.metadata.clone(),
        compression_level: options.compression_level,
    };

    let encoder = QuantizedMeshEncoder::new(header, vertices, indices, edge_indices);
    (encoder, encode_opts)
}

#[cfg(test)]
mod tests {
    use super::*;
    use quantized_mesh::DecodedMesh;

    fn bumpy(x: u32, y: u32) -> f64 {
        ((x as f64) / 8.0).sin() * 50.0 + ((y as f64) / 8.0).cos() * 30.0
    }

    #[test]
    fn flat_high_zoom_tile_collapses_to_two_triangles() {
        // A tiny (high-zoom) flat tile: the curvature bulge is sub-millimetre
        // here, well under max_error, so it adds no triangles and the tile
        // collapses to the 2 corner triangles.
        let bounds = TileBounds::new(139.0, 35.0, 139.01, 35.01);
        let bytes = encode_terrain_from_fn(
            65,
            &bounds,
            |_, _| 0.0,
            &TerrainOptions {
                max_error: 1.0,
                compression_level: 0,
                ..Default::default()
            },
        );

        let mesh = DecodedMesh::decode(&bytes).expect("decode");
        assert_eq!(mesh.indices.len(), 6);
        assert_eq!(mesh.header.min_height, 0.0);
        assert_eq!(mesh.header.max_height, 0.0);
        // All four corners present, heights all quantise to 0.
        assert!(mesh.vertices.height.iter().all(|&h| h == 0));
    }

    #[test]
    fn flat_low_zoom_tile_subdivides_for_curvature_keeping_heights_flat() {
        // A wide (low-zoom), perfectly flat tile. Without the curvature bulge
        // martini would collapse it to two triangles that cut under the
        // globe; the bulge forces subdivision — but never touches the emitted
        // heights, which stay flat at 0.
        let bounds = TileBounds::new(0.0, 0.0, 90.0, 45.0);
        let mesh = DecodedMesh::decode(&encode_terrain_from_fn(
            65,
            &bounds,
            |_, _| 0.0,
            &TerrainOptions {
                max_error: 1.0,
                compression_level: 0,
                ..Default::default()
            },
        ))
        .expect("decode");

        assert!(
            mesh.indices.len() > 6,
            "curvature must subdivide a wide flat tile, got {} indices",
            mesh.indices.len()
        );
        // The bulge only feeds the error pyramid — emitted heights stay flat.
        assert_eq!(mesh.header.min_height, 0.0);
        assert_eq!(mesh.header.max_height, 0.0);
        assert!(mesh.vertices.height.iter().all(|&h| h == 0));
    }

    #[test]
    fn curvature_bulge_zero_on_corners_and_positive_at_centre() {
        let b = TileBounds::new(0.0, 0.0, 90.0, 45.0);
        assert_eq!(curvature_bulge(0, 0, 65, &b), 0.0);
        assert_eq!(curvature_bulge(64, 0, 65, &b), 0.0);
        assert_eq!(curvature_bulge(0, 64, 65, &b), 0.0);
        assert_eq!(curvature_bulge(64, 64, 65, &b), 0.0);
        assert!(curvature_bulge(32, 32, 65, &b) > 0.0);
    }

    #[test]
    fn default_options_gzip_compress() {
        let bounds = TileBounds::new(139.0, 35.0, 139.01, 35.01);
        let bytes = encode_terrain_from_fn(65, &bounds, bumpy, &TerrainOptions::default());
        assert_eq!(&bytes[0..2], &[0x1f, 0x8b]); // gzip magic
    }

    #[test]
    fn height_range_matches_decoded_extremes() {
        let bounds = TileBounds::new(139.0, 35.0, 139.01, 35.01);
        let bytes = encode_terrain_from_fn(
            129,
            &bounds,
            bumpy,
            &TerrainOptions {
                max_error: 0.5,
                compression_level: 0,
                ..Default::default()
            },
        );
        let mesh = DecodedMesh::decode(&bytes).expect("decode");

        // The lowest mesh vertex must quantise to 0 and the highest to
        // QUANTIZED_MAX (the encoded range is defined by the header extremes).
        assert_eq!(*mesh.vertices.height.iter().min().unwrap(), 0);
        assert_eq!(*mesh.vertices.height.iter().max().unwrap(), QUANTIZED_MAX);
        assert!(mesh.header.max_height > mesh.header.min_height);
    }

    #[test]
    fn slice_and_closure_agree() {
        let grid_size = 65u32;
        let gs = grid_size as usize;
        let elevations: Vec<f32> = (0..gs * gs)
            .map(|i| bumpy((i % gs) as u32, (i / gs) as u32) as f32)
            .collect();
        let bounds = TileBounds::new(139.0, 35.0, 139.01, 35.01);
        let opts = TerrainOptions {
            max_error: 1.0,
            compression_level: 0,
            ..Default::default()
        };

        let from_slice = encode_terrain(&elevations, grid_size, &bounds, &opts);
        let from_fn = encode_terrain_from_fn(
            grid_size,
            &bounds,
            |x, y| elevations[y as usize * gs + x as usize] as f64,
            &opts,
        );
        assert_eq!(from_slice, from_fn);
    }

    #[test]
    fn writer_form_matches_vec_form() {
        let bounds = TileBounds::new(139.0, 35.0, 139.01, 35.01);
        let opts = TerrainOptions {
            max_error: 1.0,
            compression_level: 6,
            ..Default::default()
        };
        let vec_form = encode_terrain_from_fn(129, &bounds, bumpy, &opts);

        let mut writer_form = Vec::new();
        encode_terrain_from_fn_to(129, &bounds, bumpy, &opts, &mut writer_form).unwrap();
        assert_eq!(vec_form, writer_form);
    }

    #[test]
    fn face_normals_are_emitted_and_unit_length() {
        let bounds = TileBounds::new(139.0, 35.0, 139.01, 35.01);
        let bytes = encode_terrain_from_fn(
            65,
            &bounds,
            bumpy,
            &TerrainOptions {
                max_error: 1.0,
                compression_level: 0,
                normals: NormalMode::FaceNormals,
                ..Default::default()
            },
        );
        let mesh = DecodedMesh::decode(&bytes).expect("decode");
        let normals = mesh.extensions.normals.expect("normals present");
        assert_eq!(normals.len(), mesh.vertices.len());
        for n in &normals {
            let len = (n[0] * n[0] + n[1] * n[1] + n[2] * n[2]).sqrt();
            // Oct-encoding is lossy, so allow a little slack around unit length.
            assert!(
                (len - 1.0).abs() < 0.05,
                "normal not ~unit: {n:?} (len {len})"
            );
        }
    }

    #[test]
    fn buffered_gradient_normals_are_emitted() {
        let grid_size = 65u32;
        let buffer = 1u32;
        let full = (grid_size + 2 * buffer) as usize;
        // Buffered grid sampling the same bumpy field, including the halo.
        let mut buffered = Vec::with_capacity(full * full);
        for j in 0..full {
            for i in 0..full {
                let x = i as i64 - buffer as i64;
                let y = j as i64 - buffer as i64;
                buffered.push(bumpy(x.max(0) as u32, y.max(0) as u32));
            }
        }
        let buffered = BufferedElevations::new(buffered, grid_size, buffer);

        let bounds = TileBounds::new(139.0, 35.0, 139.01, 35.01);
        let bytes = encode_terrain_from_fn(
            grid_size,
            &bounds,
            bumpy,
            &TerrainOptions {
                max_error: 1.0,
                compression_level: 0,
                normals: NormalMode::BufferedGradient(buffered),
                ..Default::default()
            },
        );
        let mesh = DecodedMesh::decode(&bytes).expect("decode");
        let normals = mesh.extensions.normals.expect("normals present");
        assert_eq!(normals.len(), mesh.vertices.len());
    }

    #[test]
    #[should_panic(expected = "elevations length mismatch")]
    fn slice_length_mismatch_panics() {
        let bounds = TileBounds::new(139.0, 35.0, 139.01, 35.01);
        encode_terrain(&[0.0f32; 10], 65, &bounds, &TerrainOptions::default());
    }
}