mapterm 0.1.2

//! Mapbox Vector Tile (MVT 2.x) decoding.

use std::collections::HashMap;

use anyhow::{Result, bail};

use crate::pbf::{Reader, packed_u32, zigzag32, zigzag64};

#[derive(Clone, Debug, PartialEq)]
/// A feature property value.
pub enum Value {
    /// Text.
    Str(String),
    /// A number (every MVT numeric type, as f64).
    Num(f64),
    /// true or false.
    Bool(bool),
}

impl std::fmt::Display for Value {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Value::Str(s) => f.write_str(s),
            Value::Num(n) if n.fract() == 0.0 && n.abs() < 1e15 => write!(f, "{}", *n as i64),
            Value::Num(n) => write!(f, "{n}"),
            Value::Bool(b) => write!(f, "{b}"),
        }
    }
}

impl Value {
    /// The text, if it is text.
    pub fn as_str(&self) -> Option<&str> {
        match self {
            Value::Str(s) => Some(s),
            _ => None,
        }
    }

    /// The number, if it is a number.
    pub fn as_f64(&self) -> Option<f64> {
        match self {
            Value::Num(n) => Some(*n),
            _ => None,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
/// A feature's geometry type.
pub enum GeomType {
    /// Not given.
    Unknown,
    /// Points.
    Point,
    /// Lines.
    Line,
    /// Polygons (rings).
    Polygon,
}

#[derive(Clone, Debug)]
/// A decoded feature.
pub struct Feature {
    /// The feature's id, when the tile gives one.
    pub id: Option<u64>,
    /// Points, lines or polygons.
    pub geom_type: GeomType,
    /// Its properties.
    pub props: HashMap<String, Value>,
    /// Parts (points, lines or rings) in tile coordinates, 0..extent.
    pub geometry: Vec<Vec<[i32; 2]>>,
}

impl Feature {
    /// A text property.
    pub fn str(&self, k: &str) -> Option<&str> {
        self.props.get(k).and_then(Value::as_str)
    }

    /// A numeric property.
    pub fn num(&self, k: &str) -> Option<f64> {
        self.props.get(k).and_then(Value::as_f64)
    }
}

#[derive(Clone, Debug)]
/// A layer of a tile (`water`, `transportation`, ...).
pub struct Layer {
    /// The layer's name.
    pub name: String,
    /// Coordinates run from 0 to `extent` across the tile.
    pub extent: u32,
    /// Its features.
    pub features: Vec<Feature>,
}

#[derive(Clone, Debug, Default)]
/// A decoded vector tile.
pub struct Tile {
    /// Its layers.
    pub layers: Vec<Layer>,
}

/// The most memory a tile's feature properties may take once decoded. Each feature gets its
/// own copy of the keys and values it references from the layer's dictionaries, so a small
/// tile (one long string shared by many features) could otherwise decode to gigabytes.
/// Each entry counts its text plus 64 bytes of overhead; real tiles stay far below it.
pub const MAX_PROP_BYTES: usize = 64 << 20;
/// Bytes counted per property entry on top of its text (map entry and string headers).
const PROP_OVERHEAD: usize = 64;

/// Decode MVT bytes (already decompressed). Fails if the features' properties would take
/// more than [`MAX_PROP_BYTES`].
pub fn decode(buf: &[u8]) -> Result<Tile> {
    let mut r = Reader::new(buf);
    let mut tile = Tile::default();
    let mut budget = MAX_PROP_BYTES;
    while !r.done() {
        let (field, wire) = r.key()?;
        if field == 3 && wire == 2 {
            tile.layers.push(layer(r.bytes()?, &mut budget)?);
        } else {
            r.skip(wire)?;
        }
    }
    Ok(tile)
}

fn layer(b: &[u8], budget: &mut usize) -> Result<Layer> {
    let mut r = Reader::new(b);
    let mut name = String::new();
    let mut extent = 4096u32;
    let mut keys = Vec::new();
    let mut values = Vec::new();
    let mut raw = Vec::new();
    while !r.done() {
        match r.key()? {
            (1, 2) => name = String::from_utf8_lossy(r.bytes()?).into_owned(),
            (2, 2) => raw.push(r.bytes()?),
            (3, 2) => keys.push(String::from_utf8_lossy(r.bytes()?).into_owned()),
            (4, 2) => values.push(value(r.bytes()?)?),
            (5, 0) => extent = (r.varint()? as u32).max(1),
            (_, w) => r.skip(w)?,
        }
    }
    let mut features = Vec::with_capacity(raw.len());
    for f in raw {
        let f = feature(f, &keys, &values)?;
        let cost: usize = f.props.iter().map(|(k, v)| k.len() + v.as_str().map_or(0, str::len) + PROP_OVERHEAD).sum();
        *budget = match budget.checked_sub(cost) {
            Some(left) => left,
            None => bail!("tile properties decode past {} MB", MAX_PROP_BYTES >> 20),
        };
        features.push(f);
    }
    Ok(Layer { name, extent, features })
}

fn value(b: &[u8]) -> Result<Value> {
    let mut r = Reader::new(b);
    let mut v = Value::Bool(false);
    while !r.done() {
        v = match r.key()? {
            (1, 2) => Value::Str(String::from_utf8_lossy(r.bytes()?).into_owned()),
            (2, 5) => Value::Num(f32::from_le_bytes(r.fixed32()?) as f64),
            (3, 1) => Value::Num(f64::from_le_bytes(r.fixed64()?)),
            (4, 0) => Value::Num(r.varint()? as i64 as f64),
            (5, 0) => Value::Num(r.varint()? as f64),
            (6, 0) => Value::Num(zigzag64(r.varint()?) as f64),
            (7, 0) => Value::Bool(r.varint()? != 0),
            (_, w) => {
                r.skip(w)?;
                continue;
            }
        };
    }
    Ok(v)
}

fn feature(b: &[u8], keys: &[String], values: &[Value]) -> Result<Feature> {
    let mut r = Reader::new(b);
    let mut tags = Vec::new();
    let mut geom_type = GeomType::Unknown;
    let mut cmds = Vec::new();
    let mut id = None;
    while !r.done() {
        match r.key()? {
            (1, 0) => id = Some(r.varint()?),
            (2, 2) => tags = packed_u32(r.bytes()?)?,
            (3, 0) => {
                geom_type = match r.varint()? {
                    1 => GeomType::Point,
                    2 => GeomType::Line,
                    3 => GeomType::Polygon,
                    _ => GeomType::Unknown,
                }
            }
            (4, 2) => cmds = packed_u32(r.bytes()?)?,
            (_, w) => r.skip(w)?,
        }
    }
    let mut props = HashMap::new();
    for pair in tags.as_chunks::<2>().0 {
        if let (Some(k), Some(v)) = (keys.get(pair[0] as usize), values.get(pair[1] as usize)) {
            props.insert(k.clone(), v.clone());
        }
    }
    Ok(Feature { id, geom_type, props, geometry: geometry(&cmds) })
}

fn geometry(cmds: &[u32]) -> Vec<Vec<[i32; 2]>> {
    let mut parts = Vec::new();
    let mut cur: Vec<[i32; 2]> = Vec::new();
    let (mut x, mut y) = (0i32, 0i32);
    let mut i = 0;
    while i < cmds.len() {
        let (cmd, count) = (cmds[i] & 7, (cmds[i] >> 3) as usize);
        i += 1;
        match cmd {
            1 | 2 => {
                for _ in 0..count {
                    if i + 1 >= cmds.len() {
                        break;
                    }
                    let (dx, dy) = (zigzag32(cmds[i]), zigzag32(cmds[i + 1]));
                    i += 2;
                    x = x.wrapping_add(dx);
                    y = y.wrapping_add(dy);
                    if cmd == 1 && !cur.is_empty() {
                        parts.push(std::mem::take(&mut cur));
                    }
                    cur.push([x, y]);
                }
            }
            7 => {
                if let Some(&first) = cur.first() {
                    cur.push(first);
                }
            }
            _ => break,
        }
    }
    if !cur.is_empty() {
        parts.push(cur);
    }
    parts
}

#[cfg(test)]
mod tests {

    /// One long string in the layer's dictionary, referenced by many tiny features: the
    /// raw tile is small, but each feature would get its own copy of the string.
    fn amplifying_tile(features: usize, string_bytes: usize) -> Vec<u8> {
        fn varint(out: &mut Vec<u8>, mut v: u64) {
            while v >= 0x80 {
                out.push(v as u8 | 0x80);
                v >>= 7;
            }
            out.push(v as u8);
        }
        fn field(out: &mut Vec<u8>, n: u64, b: &[u8]) {
            varint(out, n << 3 | 2);
            varint(out, b.len() as u64);
            out.extend_from_slice(b);
        }
        let mut layer = Vec::new();
        field(&mut layer, 1, b"poi");
        for _ in 0..features {
            // tags [0, 0]; type POINT; MoveTo(1) to (100, 100)
            let mut f = Vec::new();
            field(&mut f, 2, &[0, 0]);
            f.extend_from_slice(&[3 << 3, 1]);
            field(&mut f, 4, &[9, 200, 1, 200, 1]);
            field(&mut layer, 2, &f);
        }
        field(&mut layer, 3, b"name");
        let mut value = Vec::new();
        field(&mut value, 1, &vec![b'A'; string_bytes]);
        field(&mut layer, 4, &value);
        let mut tile = Vec::new();
        field(&mut tile, 3, &layer);
        tile
    }

    #[test]
    fn a_small_tile_cannot_decode_to_gigabytes() {
        // 1 MiB raw; 1000 copies of the string would be about 1 GiB.
        let tile = amplifying_tile(1000, 1 << 20);
        assert!(tile.len() < 2 << 20);
        let err = decode(&tile).expect_err("should be refused").to_string();
        assert!(err.contains("decode past"), "{err}");
        // A few copies are fine.
        let ok = decode(&amplifying_tile(10, 1 << 20)).unwrap();
        assert_eq!(ok.layers[0].features.len(), 10);
    }

    use super::*;

    #[test]
    fn decodes_spec_examples() {
        // MVT spec examples: a point at (25,17); a line (2,2)->(2,10)->(10,10).
        assert_eq!(geometry(&[9, 50, 34]), vec![vec![[25, 17]]]);
        assert_eq!(geometry(&[9, 4, 4, 18, 0, 16, 16, 0]), vec![vec![[2, 2], [2, 10], [10, 10]]]);
        // Polygon (3,6)->(8,12)->(20,34), closed.
        assert_eq!(geometry(&[9, 6, 12, 18, 10, 12, 24, 44, 15]), vec![vec![[3, 6], [8, 12], [20, 34], [3, 6]]]);
    }

    #[test]
    fn truncated_geometry_is_safe() {
        assert_eq!(geometry(&[9, 4]), Vec::<Vec<[i32; 2]>>::new());
        assert_eq!(geometry(&[(1 << 3) | 2, 2, 2]), vec![vec![[1, 1]]]);
    }

    #[test]
    fn feature_ids() {
        let t = decode(&tests_support::square_tile_with_id("water", 4096, Some(42))).unwrap();
        assert_eq!(t.layers[0].features[0].id, Some(42));
        let t = decode(&tests_support::square_tile("water", 4096)).unwrap();
        assert_eq!(t.layers[0].features[0].id, None);
    }

    #[test]
    fn garbage_does_not_panic() {
        for seed in 0u32..200 {
            let bytes: Vec<u8> = (0..64).map(|i| (seed.wrapping_mul(2654435761).wrapping_add(i * 97) >> 7) as u8).collect();
            let _ = decode(&bytes);
        }
    }
}

/// A tiny MVT encoder for tests.
#[cfg(test)]
pub(crate) mod tests_support {
    fn varint(out: &mut Vec<u8>, mut v: u64) {
        while v >= 0x80 {
            out.push((v as u8) | 0x80);
            v >>= 7;
        }
        out.push(v as u8);
    }

    fn field_bytes(out: &mut Vec<u8>, field: u32, b: &[u8]) {
        varint(out, ((field as u64) << 3) | 2);
        varint(out, b.len() as u64);
        out.extend_from_slice(b);
    }

    fn zz(v: i32) -> u32 {
        ((v << 1) ^ (v >> 31)) as u32
    }

    /// One tile with one layer holding one polygon feature that covers the whole tile.
    pub fn square_tile(layer: &str, extent: u32) -> Vec<u8> {
        square_tile_with_id(layer, extent, None)
    }

    /// Like [`square_tile`], with a feature id.
    pub fn square_tile_with_id(layer: &str, extent: u32, id: Option<u64>) -> Vec<u8> {
        let e = extent as i32;
        let cmds: Vec<u32> = vec![9, zz(0), zz(0), 26, zz(e), zz(0), zz(0), zz(e), zz(-e), zz(0), 15];
        let mut geom = Vec::new();
        for c in cmds {
            varint(&mut geom, c as u64);
        }
        let mut feat = Vec::new();
        if let Some(id) = id {
            varint(&mut feat, 1 << 3);
            varint(&mut feat, id);
        }
        varint(&mut feat, 3 << 3); // type
        varint(&mut feat, 3); // polygon
        field_bytes(&mut feat, 4, &geom);
        let mut l = Vec::new();
        varint(&mut l, 15 << 3);
        varint(&mut l, 2);
        field_bytes(&mut l, 1, layer.as_bytes());
        field_bytes(&mut l, 2, &feat);
        varint(&mut l, 5 << 3);
        varint(&mut l, extent as u64);
        let mut tile = Vec::new();
        field_bytes(&mut tile, 3, &l);
        tile
    }
}