secs-project 0.15.3

//! Python-compatible JSON values, parser, and pretty-printer.

use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::fmt;
use std::mem::size_of;

const DEFAULT_BUDGET_EXTRA: usize = 16 * 1024 * 1024;

/// A JSON value exactly as Python's `json` module reads and writes it.
#[derive(Clone, Debug, PartialEq)]
pub enum PyJson {
    Null,
    Bool(bool),
    /// Any integer, as canonical decimal text.
    Int(String),
    /// Any `f64`, including NaN, infinities, and negative zero.
    Float(f64),
    Str(String),
    Array(Vec<PyJson>),
    /// Object members in insertion order.
    Object(Vec<(String, PyJson)>),
}

impl PyJson {
    pub(crate) fn object(&self) -> Option<&[(String, PyJson)]> {
        match self {
            Self::Object(value) => Some(value),
            _ => None,
        }
    }

    pub(crate) fn array(&self) -> Option<&[PyJson]> {
        match self {
            Self::Array(value) => Some(value),
            _ => None,
        }
    }

    pub(crate) fn str(&self) -> Option<&str> {
        match self {
            Self::Str(value) => Some(value),
            _ => None,
        }
    }

    pub(crate) fn bool(&self) -> Option<bool> {
        match self {
            Self::Bool(value) => Some(*value),
            _ => None,
        }
    }

    pub(crate) fn get(&self, key: &str) -> Option<&PyJson> {
        self.object()?
            .iter()
            .find(|(name, _)| name == key)
            .map(|(_, value)| value)
    }

    pub(crate) fn object_mut(&mut self) -> Option<&mut Vec<(String, PyJson)>> {
        match self {
            Self::Object(value) => Some(value),
            _ => None,
        }
    }

    /// Python's `bool(value)`: `None`, `False`, zero and empty containers are false.
    pub(crate) fn is_truthy(&self) -> bool {
        match self {
            Self::Null => false,
            Self::Bool(value) => *value,
            Self::Int(value) => value != "0",
            Self::Float(value) => *value != 0.0,
            Self::Str(value) => !value.is_empty(),
            Self::Array(value) => !value.is_empty(),
            Self::Object(value) => !value.is_empty(),
        }
    }
}

/// How a value the writer rebuilds from a typed field meets the value the
/// file held, so saving keeps keys this version does not know (ADR 0004,
/// docs/DEVIATIONS.md F16).
#[derive(Debug)]
pub(crate) enum Shape {
    /// An object with fixed field names: the writer's fields in its order,
    /// then the file's other fields in theirs. `optional` names fields the
    /// writer leaves out on purpose, which therefore do not come back from the
    /// file. `nested` gives the shape of fields merged in turn; the writer's
    /// value replaces every other field.
    Record {
        optional: &'static [&'static str],
        nested: &'static [(&'static str, Shape)],
    },
    /// An object whose keys are data (roles, bundles): only the writer's
    /// entries, each merged with the file's entry of the same key.
    Map(&'static Shape),
    /// An array merged element by element when both sides have the same
    /// length; otherwise the writer's array.
    List(&'static Shape),
}

/// Merges `new`, rebuilt by the writer, with `old`, read from the file.
pub(crate) fn merge_with_file(new: PyJson, old: &PyJson, shape: &Shape) -> PyJson {
    match (shape, new, old) {
        (
            Shape::Record { optional, nested },
            PyJson::Object(fields),
            PyJson::Object(old_fields),
        ) => {
            let mut out = fields
                .into_iter()
                .map(|(key, value)| {
                    let value = match (
                        nested.iter().find(|(name, _)| *name == key),
                        old_fields.iter().find(|(name, _)| *name == key),
                    ) {
                        (Some((_, shape)), Some((_, old))) => merge_with_file(value, old, shape),
                        _ => value,
                    };
                    (key, value)
                })
                .collect::<Vec<_>>();
            let unknown = old_fields
                .iter()
                .filter(|(key, _)| {
                    !optional.contains(&key.as_str()) && !out.iter().any(|(name, _)| name == key)
                })
                .cloned()
                .collect::<Vec<_>>();
            out.extend(unknown);
            PyJson::Object(out)
        }
        (Shape::Map(entry), PyJson::Object(fields), PyJson::Object(old_fields)) => PyJson::Object(
            fields
                .into_iter()
                .map(|(key, value)| {
                    let value = match old_fields.iter().find(|(name, _)| *name == key) {
                        Some((_, old)) => merge_with_file(value, old, entry),
                        None => value,
                    };
                    (key, value)
                })
                .collect(),
        ),
        (Shape::List(item), PyJson::Array(items), PyJson::Array(old_items))
            if items.len() == old_items.len() =>
        {
            PyJson::Array(
                items
                    .into_iter()
                    .zip(old_items)
                    .map(|(value, old)| merge_with_file(value, old, item))
                    .collect(),
            )
        }
        (_, new, _) => new,
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PyJsonError {
    offset: usize,
    message: String,
}

impl PyJsonError {
    fn new(offset: usize, message: impl Into<String>) -> Self {
        Self {
            offset,
            message: message.into(),
        }
    }
}

impl fmt::Display for PyJsonError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} at byte {}", self.message, self.offset)
    }
}

impl std::error::Error for PyJsonError {}

/// Parses the same value language as CPython's `json.loads`.
pub fn loads(text: &str) -> Result<PyJson, PyJsonError> {
    let budget = text
        .len()
        .saturating_mul(4)
        .saturating_add(DEFAULT_BUDGET_EXTRA);
    loads_with_budget(text, budget)
}

/// Parses with an explicit allocation budget in bytes.
pub fn loads_with_budget(text: &str, budget: usize) -> Result<PyJson, PyJsonError> {
    let mut parser = Parser {
        text,
        bytes: text.as_bytes(),
        pos: 0,
        remaining: budget,
    };
    parser.whitespace();
    let value = parser.value(0)?;
    parser.whitespace();
    if parser.pos != parser.bytes.len() {
        return Err(parser.error("extra data"));
    }
    Ok(value)
}

struct Parser<'a> {
    text: &'a str,
    bytes: &'a [u8],
    pos: usize,
    remaining: usize,
}

impl Parser<'_> {
    fn error(&self, message: impl Into<String>) -> PyJsonError {
        PyJsonError::new(self.pos, message)
    }

    fn charge(&mut self, bytes: usize) -> Result<(), PyJsonError> {
        self.remaining = self
            .remaining
            .checked_sub(bytes)
            .ok_or_else(|| self.error("JSON allocation budget exceeded"))?;
        Ok(())
    }

    fn node(&mut self, value: PyJson, string_bytes: usize) -> Result<PyJson, PyJsonError> {
        self.charge(size_of::<PyJson>().saturating_add(string_bytes))?;
        Ok(value)
    }

    fn reserve_for_push<T>(&mut self, values: &mut Vec<T>) -> Result<(), PyJsonError> {
        if values.len() < values.capacity() {
            return Ok(());
        }
        let old_capacity = values.capacity();
        let new_capacity = old_capacity.saturating_mul(2).max(4);
        let additional = new_capacity.saturating_sub(old_capacity);
        self.charge(additional.saturating_mul(size_of::<T>()))?;
        values
            .try_reserve_exact(additional)
            .map_err(|_| self.error("JSON allocation failed"))?;
        let extra = values.capacity().saturating_sub(new_capacity);
        self.charge(extra.saturating_mul(size_of::<T>()))
    }

    fn whitespace(&mut self) {
        while matches!(self.bytes.get(self.pos), Some(b' ' | b'\n' | b'\r' | b'\t')) {
            self.pos += 1;
        }
    }

    fn value(&mut self, depth: usize) -> Result<PyJson, PyJsonError> {
        match self.bytes.get(self.pos).copied() {
            Some(b'n') => self.keyword(b"null", PyJson::Null),
            Some(b't') => self.keyword(b"true", PyJson::Bool(true)),
            Some(b'f') => self.keyword(b"false", PyJson::Bool(false)),
            Some(b'N') => self.keyword(b"NaN", PyJson::Float(f64::NAN)),
            Some(b'I') => self.keyword(b"Infinity", PyJson::Float(f64::INFINITY)),
            Some(b'-')
                if self
                    .bytes
                    .get(self.pos..)
                    .is_some_and(|s| s.starts_with(b"-Infinity")) =>
            {
                self.keyword(b"-Infinity", PyJson::Float(f64::NEG_INFINITY))
            }
            Some(b'"') => {
                let string = self.string()?;
                self.node(PyJson::Str(string), 0)
            }
            Some(b'[') => self.array_value(depth),
            Some(b'{') => self.object_value(depth),
            Some(b'-' | b'0'..=b'9') => self.number(),
            _ => Err(self.error("expecting value")),
        }
    }

    fn keyword(&mut self, word: &[u8], value: PyJson) -> Result<PyJson, PyJsonError> {
        if self
            .bytes
            .get(self.pos..)
            .is_some_and(|rest| rest.starts_with(word))
        {
            self.pos += word.len();
            self.node(value, 0)
        } else {
            Err(self.error("expecting value"))
        }
    }

    fn nested(&self, depth: usize) -> Result<usize, PyJsonError> {
        let next = depth + 1;
        if next > 1000 {
            Err(self.error("maximum JSON nesting depth exceeded"))
        } else {
            Ok(next)
        }
    }

    fn array_value(&mut self, depth: usize) -> Result<PyJson, PyJsonError> {
        let depth = self.nested(depth)?;
        self.pos += 1;
        self.whitespace();
        let mut values = Vec::new();
        if self.bytes.get(self.pos) == Some(&b']') {
            self.pos += 1;
            return self.node(PyJson::Array(values), 0);
        }
        loop {
            let value = self.value(depth)?;
            self.reserve_for_push(&mut values)?;
            values.push(value);
            self.whitespace();
            match self.bytes.get(self.pos) {
                Some(b',') => {
                    self.pos += 1;
                    self.whitespace();
                }
                Some(b']') => {
                    self.pos += 1;
                    return self.node(PyJson::Array(values), 0);
                }
                _ => return Err(self.error("expecting ',' or ']'")),
            }
        }
    }

    fn object_value(&mut self, depth: usize) -> Result<PyJson, PyJsonError> {
        let depth = self.nested(depth)?;
        self.pos += 1;
        self.whitespace();
        let mut values: Vec<(String, PyJson)> = Vec::new();
        if self.bytes.get(self.pos) == Some(&b'}') {
            self.pos += 1;
            return self.node(PyJson::Object(values), 0);
        }
        // Key -> position in `values`, only while this object is parsed.
        let mut index: HashMap<String, usize> = HashMap::new();
        loop {
            if self.bytes.get(self.pos) != Some(&b'"') {
                return Err(self.error("expecting property name"));
            }
            let key = self.string()?;
            self.whitespace();
            if self.bytes.get(self.pos) != Some(&b':') {
                return Err(self.error("expecting ':'"));
            }
            self.pos += 1;
            self.whitespace();
            let value = self.value(depth)?;
            // Like a Python dict: a repeated key keeps its first position and
            // takes the last value.
            match index.entry(key) {
                Entry::Occupied(slot) => values[*slot.get()].1 = value,
                Entry::Vacant(slot) => {
                    self.charge(
                        slot.key()
                            .len()
                            .saturating_add(size_of::<(String, usize)>()),
                    )?;
                    let cloned_key = slot.key().clone();
                    self.charge(cloned_key.capacity().saturating_sub(slot.key().len()))?;
                    self.reserve_for_push(&mut values)?;
                    values.push((cloned_key, value));
                    slot.insert(values.len() - 1);
                }
            }
            self.whitespace();
            match self.bytes.get(self.pos) {
                Some(b',') => {
                    self.pos += 1;
                    self.whitespace();
                }
                Some(b'}') => {
                    self.pos += 1;
                    return self.node(PyJson::Object(values), 0);
                }
                _ => return Err(self.error("expecting ',' or '}'")),
            }
        }
    }

    fn string_capacity(&self) -> Result<usize, PyJsonError> {
        let start = self.pos.saturating_add(1);
        let mut cursor = start;
        while let Some(&byte) = self.bytes.get(cursor) {
            match byte {
                b'"' => return Ok(cursor - start),
                b'\\' => {
                    cursor += 1;
                    if self.bytes.get(cursor).is_none() {
                        return Err(PyJsonError::new(cursor, "unterminated string"));
                    }
                    cursor += 1;
                }
                0..=0x1f => {
                    return Err(PyJsonError::new(cursor, "invalid control character"));
                }
                _ => cursor += 1,
            }
        }
        Err(PyJsonError::new(cursor, "unterminated string"))
    }

    fn string(&mut self) -> Result<String, PyJsonError> {
        let requested = self.string_capacity()?;
        self.charge(requested)?;
        let mut out = String::new();
        out.try_reserve_exact(requested)
            .map_err(|_| self.error("JSON allocation failed"))?;
        self.charge(out.capacity().saturating_sub(requested))?;
        self.pos += 1;
        let mut raw_start = self.pos;
        while let Some(&byte) = self.bytes.get(self.pos) {
            match byte {
                b'"' => {
                    out.push_str(&self.text[raw_start..self.pos]);
                    self.pos += 1;
                    return Ok(out);
                }
                b'\\' => {
                    out.push_str(&self.text[raw_start..self.pos]);
                    self.pos += 1;
                    let escaped = *self
                        .bytes
                        .get(self.pos)
                        .ok_or_else(|| self.error("unterminated string"))?;
                    self.pos += 1;
                    match escaped {
                        b'"' => out.push('"'),
                        b'\\' => out.push('\\'),
                        b'/' => out.push('/'),
                        b'b' => out.push('\u{8}'),
                        b'f' => out.push('\u{c}'),
                        b'n' => out.push('\n'),
                        b'r' => out.push('\r'),
                        b't' => out.push('\t'),
                        b'u' => {
                            let first = self.hex_unit()?;
                            let scalar = if (0xd800..=0xdbff).contains(&first) {
                                if self.bytes.get(self.pos..self.pos.saturating_add(2))
                                    != Some(b"\\u")
                                {
                                    return Err(self.error("unpaired UTF-16 surrogate"));
                                }
                                self.pos += 2;
                                let second = self.hex_unit()?;
                                if !(0xdc00..=0xdfff).contains(&second) {
                                    return Err(self.error("unpaired UTF-16 surrogate"));
                                }
                                0x10000
                                    + ((u32::from(first) - 0xd800) << 10)
                                    + (u32::from(second) - 0xdc00)
                            } else if (0xdc00..=0xdfff).contains(&first) {
                                return Err(self.error("unpaired UTF-16 surrogate"));
                            } else {
                                u32::from(first)
                            };
                            out.push(
                                char::from_u32(scalar)
                                    .ok_or_else(|| self.error("invalid Unicode escape"))?,
                            );
                        }
                        _ => return Err(self.error("invalid escape")),
                    }
                    raw_start = self.pos;
                }
                0..=0x1f => return Err(self.error("invalid control character")),
                _ => {
                    let ch = self.text[self.pos..]
                        .chars()
                        .next()
                        .ok_or_else(|| self.error("unterminated string"))?;
                    self.pos += ch.len_utf8();
                }
            }
        }
        Err(self.error("unterminated string"))
    }

    fn hex_unit(&mut self) -> Result<u16, PyJsonError> {
        let end = self
            .pos
            .checked_add(4)
            .ok_or_else(|| self.error("invalid Unicode escape"))?;
        let digits = self
            .text
            .get(self.pos..end)
            .ok_or_else(|| self.error("invalid Unicode escape"))?;
        if !digits.bytes().all(|b| b.is_ascii_hexdigit()) {
            return Err(self.error("invalid Unicode escape"));
        }
        self.pos = end;
        u16::from_str_radix(digits, 16).map_err(|_| self.error("invalid Unicode escape"))
    }

    fn number(&mut self) -> Result<PyJson, PyJsonError> {
        let start = self.pos;
        if self.bytes.get(self.pos) == Some(&b'-') {
            self.pos += 1;
        }
        match self.bytes.get(self.pos) {
            Some(b'0') => self.pos += 1,
            Some(b'1'..=b'9') => {
                self.pos += 1;
                while self.bytes.get(self.pos).is_some_and(u8::is_ascii_digit) {
                    self.pos += 1;
                }
            }
            _ => return Err(self.error("invalid number")),
        }
        let mut float = false;
        if self.bytes.get(self.pos) == Some(&b'.') {
            float = true;
            self.pos += 1;
            let before = self.pos;
            while self.bytes.get(self.pos).is_some_and(u8::is_ascii_digit) {
                self.pos += 1;
            }
            if self.pos == before {
                return Err(self.error("invalid number"));
            }
        }
        if matches!(self.bytes.get(self.pos), Some(b'e' | b'E')) {
            float = true;
            self.pos += 1;
            if matches!(self.bytes.get(self.pos), Some(b'+' | b'-')) {
                self.pos += 1;
            }
            let before = self.pos;
            while self.bytes.get(self.pos).is_some_and(u8::is_ascii_digit) {
                self.pos += 1;
            }
            if self.pos == before {
                return Err(self.error("invalid number"));
            }
        }
        let raw = &self.text[start..self.pos];
        if float {
            let value = raw
                .parse::<f64>()
                .map_err(|_| self.error("invalid number"))?;
            self.node(PyJson::Float(value), 0)
        } else {
            let (text_start, bytes) = if raw == "-0" {
                (None, 1)
            } else {
                (Some(start), raw.len())
            };
            self.charge(bytes)?;
            let value = text_start.map_or_else(
                || "0".to_owned(),
                |start| self.text[start..self.pos].to_owned(),
            );
            self.charge(value.capacity().saturating_sub(bytes))?;
            self.node(PyJson::Int(value), 0)
        }
    }
}

/// Formats like `json.dumps(value, ensure_ascii=False, indent=2)`.
pub fn dumps(value: &PyJson) -> String {
    let mut out = String::new();
    write_value(value, 0, &mut out);
    out
}

/// Formats like `json.dumps(value, ensure_ascii=False)`.
pub fn dumps_compact(value: &PyJson) -> String {
    let mut out = String::new();
    write_compact(value, &mut out);
    out
}

fn write_compact(value: &PyJson, out: &mut String) {
    match value {
        PyJson::Null => out.push_str("null"),
        PyJson::Bool(value) => out.push_str(if *value { "true" } else { "false" }),
        PyJson::Int(value) => out.push_str(value),
        PyJson::Float(value) => out.push_str(&python_float(*value)),
        PyJson::Str(value) => write_string(value, out),
        PyJson::Array(values) => {
            out.push('[');
            for (index, value) in values.iter().enumerate() {
                if index != 0 {
                    out.push_str(", ");
                }
                write_compact(value, out);
            }
            out.push(']');
        }
        PyJson::Object(values) => {
            out.push('{');
            for (index, (key, value)) in values.iter().enumerate() {
                if index != 0 {
                    out.push_str(", ");
                }
                write_string(key, out);
                out.push_str(": ");
                write_compact(value, out);
            }
            out.push('}');
        }
    }
}

fn write_value(value: &PyJson, depth: usize, out: &mut String) {
    match value {
        PyJson::Null => out.push_str("null"),
        PyJson::Bool(value) => out.push_str(if *value { "true" } else { "false" }),
        PyJson::Int(value) => out.push_str(value),
        PyJson::Float(value) => out.push_str(&python_float(*value)),
        PyJson::Str(value) => write_string(value, out),
        PyJson::Array(values) => {
            if values.is_empty() {
                out.push_str("[]");
                return;
            }
            out.push_str("[\n");
            for (index, value) in values.iter().enumerate() {
                out.push_str(&" ".repeat((depth + 1) * 2));
                write_value(value, depth + 1, out);
                if index + 1 != values.len() {
                    out.push(',');
                }
                out.push('\n');
            }
            out.push_str(&" ".repeat(depth * 2));
            out.push(']');
        }
        PyJson::Object(values) => {
            if values.is_empty() {
                out.push_str("{}");
                return;
            }
            out.push_str("{\n");
            for (index, (key, value)) in values.iter().enumerate() {
                out.push_str(&" ".repeat((depth + 1) * 2));
                write_string(key, out);
                out.push_str(": ");
                write_value(value, depth + 1, out);
                if index + 1 != values.len() {
                    out.push(',');
                }
                out.push('\n');
            }
            out.push_str(&" ".repeat(depth * 2));
            out.push('}');
        }
    }
}

fn write_string(value: &str, out: &mut String) {
    out.push('"');
    for ch in value.chars() {
        match ch {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\u{8}' => out.push_str("\\b"),
            '\u{c}' => out.push_str("\\f"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            '\0'..='\u{1f}' => {
                use std::fmt::Write as _;
                write!(out, "\\u{:04x}", ch as u32).expect("writing to String cannot fail");
            }
            _ => out.push(ch),
        }
    }
    out.push('"');
}

fn python_float(value: f64) -> String {
    if value.is_nan() {
        return "NaN".into();
    }
    if value == f64::INFINITY {
        return "Infinity".into();
    }
    if value == f64::NEG_INFINITY {
        return "-Infinity".into();
    }
    if value == 0.0 && value.is_sign_negative() {
        return "-0.0".into();
    }
    let mut text = ryu::Buffer::new().format_finite(value).to_owned();
    let absolute = value.abs();
    if !text.contains(['e', 'E']) && absolute != 0.0 && (absolute < 1e-4 || absolute >= 1e16) {
        let negative = text.starts_with('-');
        let unsigned = text.trim_start_matches('-');
        let (digits, exponent) = if absolute < 1.0 {
            let fractional = unsigned.trim_start_matches("0.");
            let zeroes = fractional.chars().take_while(|ch| *ch == '0').count();
            (
                fractional.trim_start_matches('0').to_owned(),
                -(zeroes as i32) - 1,
            )
        } else {
            let whole = unsigned.find('.').unwrap_or(unsigned.len());
            (unsigned.replace('.', ""), whole as i32 - 1)
        };
        let (first, rest) = digits.split_at(1);
        let mantissa = if rest.is_empty() {
            first.into()
        } else {
            format!("{first}.{rest}")
        };
        text = format!("{}{mantissa}e{exponent:+}", if negative { "-" } else { "" });
    }
    if !text.contains(['.', 'e', 'E']) {
        text.push_str(".0");
    }
    if let Some(index) = text.find('e') {
        let exponent = text[index + 1..]
            .parse::<i32>()
            .expect("ryu exponent is an integer");
        text = format!(
            "{}e{}{:02}",
            &text[..index],
            if exponent >= 0 { "+" } else { "-" },
            exponent.abs()
        );
    }
    text
}

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

    #[test]
    fn scanner_and_encoder_edge_cases() {
        let input = r#" {"a":-0,"b":-0.0,"c":1e400,"d":340282366920938463463374607431768211456,"a":2,"s":"한글 😀\n😀"} "#;
        let value = loads(input).unwrap();
        assert_eq!(value.get("a"), Some(&PyJson::Int("2".into())));
        assert_eq!(value.get("b"), Some(&PyJson::Float(-0.0)));
        assert_eq!(value.get("c"), Some(&PyJson::Float(f64::INFINITY)));
        assert!(dumps(&value).contains("340282366920938463463374607431768211456"));
        assert!(dumps(&value).contains("한글 😀\\n😀"));
    }

    #[test]
    fn rejects_controls_surrogates_and_excess_depth() {
        assert!(loads("\"line\n\"").is_err());
        assert!(loads(r#""\ud800""#).is_err());
        assert!(loads(&format!("{}0{}", "[".repeat(1000), "]".repeat(1000))).is_ok());
        assert!(loads(&format!("{}0{}", "[".repeat(1001), "]".repeat(1001))).is_err());
    }

    #[test]
    fn allocation_budget_rejects_a_wide_array_td19() {
        let document = format!("[{}]", vec!["null"; 64].join(","));
        let budget = std::mem::size_of::<PyJson>() * 8;
        let error = loads_with_budget(&document, budget).unwrap_err();
        assert!(error.to_string().contains("allocation budget"));
        assert!(loads(&document).is_ok());
    }

    #[test]
    fn allocation_budget_rejects_a_large_string_before_allocating_td19() {
        let document = format!("\"{}\"", "x".repeat(2 * 1024 * 1024));
        let error = loads_with_budget(&document, 64).unwrap_err();
        assert_eq!(error.offset, 0);
        assert!(error.to_string().contains("allocation budget"));
    }
}