core/test/rust/performance/json/writer.rs

//! 冻结原写出算法供前后对照；生产运行库不依赖此通用节点适配。
use dever_runtime::wire::document::{MAX_BYTES, MAX_VALUES, WriteError};
use dever_runtime::wire::serde_json;

pub enum Node<Text, Children, Members> {
    Null,
    Boolean(bool),
    Number(Text),
    String(Text),
    Array(Children),
    Object(Members),
}

#[derive(Clone, Copy)]
struct Size {
    bytes: usize,
    values: usize,
    depth: usize,
}

/// 节点直接按下标借用，编号和文档结构在 measure 中检查；不再给取节点包装错误通道。
pub fn stringify_nodes<'node, Children, Members>(
    root: usize,
    count: usize,
    node: impl Fn(usize) -> Node<&'node str, Children, Members>,
    max_depth: usize,
) -> Result<String, WriteError>
where
    Children: Iterator<Item = i64>,
    Members: Iterator<Item = (&'node str, i64)>,
{
    let bytes = measure(root, count, &node, max_depth)?;
    let mut output = String::with_capacity(bytes);
    write_node(&node, root, &mut output);
    Ok(output)
}

fn measure<'node, Children, Members>(
    root: usize,
    count: usize,
    node: &impl Fn(usize) -> Node<&'node str, Children, Members>,
    max_depth: usize,
) -> Result<usize, WriteError>
where
    Children: Iterator<Item = i64>,
    Members: Iterator<Item = (&'node str, i64)>,
{
    let fail = |node, message: &str| WriteError {
        node,
        message: message.into(),
    };
    if count > MAX_VALUES {
        return Err(fail(root, "JSON node count exceeds 65536 values"));
    }
    if root >= count {
        return Err(fail(root, "JSON root does not refer to a document node"));
    }
    let mut sizes = Vec::<Size>::with_capacity(count);
    for index in 0..count {
        let mut size = Size {
            bytes: 0,
            values: 1,
            depth: 0,
        };
        match node(index) {
            Node::Null => size.bytes = 4,
            Node::Boolean(value) => size.bytes = if value { 4 } else { 5 },
            Node::Number(raw) => {
                if !matches!(raw.as_bytes().first(), Some(b'-' | b'0'..=b'9'))
                    || !serde_json::from_str::<&serde_json::value::RawValue>(raw)
                        .is_ok_and(|value| value.get() == raw)
                {
                    return Err(fail(index, "invalid JSON number"));
                }
                size.bytes = raw.len();
            }
            Node::String(text) => size.bytes = string_size(text),
            Node::Array(children) => {
                size.bytes = 2;
                size.depth = 1;
                for (position, child) in children.enumerate() {
                    size.bytes += usize::from(position != 0);
                    add_child(&mut size, child, &sizes, index, max_depth)?;
                }
            }
            Node::Object(fields) => {
                size.bytes = 2;
                size.depth = 1;
                for (position, (name, child)) in fields.enumerate() {
                    size.bytes += string_size(name) + 1 + usize::from(position != 0);
                    add_child(&mut size, child, &sizes, index, max_depth)?;
                }
            }
        }
        check(size, index, max_depth)?;
        sizes.push(size);
    }
    Ok(sizes[root].bytes)
}

fn add_child(
    size: &mut Size,
    child: i64,
    sizes: &[Size],
    node: usize,
    max_depth: usize,
) -> Result<(), WriteError> {
    let child = usize::try_from(child).map_err(|_| WriteError {
        node,
        message: "negative JSON node ID".into(),
    })?;
    let child = sizes.get(child).ok_or_else(|| WriteError {
        node,
        message: "JSON child must precede parent".into(),
    })?;
    size.bytes += child.bytes;
    size.values += child.values;
    size.depth = size.depth.max(child.depth + 1);
    check(*size, node, max_depth)
}

fn check(size: Size, node: usize, max_depth: usize) -> Result<(), WriteError> {
    let message = if size.bytes > MAX_BYTES {
        "JSON exceeds 16 MiB"
    } else if size.values > MAX_VALUES {
        "JSON contains more than 65536 values"
    } else if size.depth > max_depth {
        "JSON nesting exceeds configured limit"
    } else {
        return Ok(());
    };
    Err(WriteError {
        node,
        message: message.into(),
    })
}

fn string_size(text: &str) -> usize {
    text.len()
        + 2
        + text
            .bytes()
            .map(|byte| match byte {
                b'"' | b'\\' | 8 | 12 | b'\n' | b'\r' | b'\t' => 1,
                0..=31 => 5,
                _ => 0,
            })
            .sum::<usize>()
}

fn write_node<'node, Children, Members>(
    node: &impl Fn(usize) -> Node<&'node str, Children, Members>,
    index: usize,
    output: &mut String,
) where
    Children: Iterator<Item = i64>,
    Members: Iterator<Item = (&'node str, i64)>,
{
    match node(index) {
        Node::Null => output.push_str("null"),
        Node::Boolean(value) => output.push_str(if value { "true" } else { "false" }),
        Node::Number(raw) => output.push_str(raw),
        Node::String(text) => write_string(text, output),
        Node::Array(children) => {
            output.push('[');
            for (position, child) in children.enumerate() {
                if position != 0 {
                    output.push(',');
                }
                write_node(node, child as usize, output);
            }
            output.push(']');
        }
        Node::Object(fields) => {
            output.push('{');
            for (position, (name, child)) in fields.enumerate() {
                if position != 0 {
                    output.push(',');
                }
                write_string(name, output);
                output.push(':');
                write_node(node, child as usize, output);
            }
            output.push('}');
        }
    }
}

fn write_string(text: &str, output: &mut String) {
    output.push('"');
    for character in text.chars() {
        match character {
            '"' => output.push_str("\\\""),
            '\\' => output.push_str("\\\\"),
            '\u{08}' => output.push_str("\\b"),
            '\u{0c}' => output.push_str("\\f"),
            '\n' => output.push_str("\\n"),
            '\r' => output.push_str("\\r"),
            '\t' => output.push_str("\\t"),
            character if character < '\u{20}' => {
                use std::fmt::Write;
                write!(output, "\\u{:04x}", character as u32).unwrap();
            }
            character => output.push(character),
        }
    }
    output.push('"');
}