ocpi_tariffs/json/

write.rs

//! Tools for writing JSON.
use std::fmt::{self, Write};

use super::{Element, Field, Value};

const TAB: &str = "    ";
const ARRAY_OPEN: char = '[';
const ARRAY_CLOSE: char = ']';
const OBJECT_OPEN: char = '{';
const OBJECT_CLOSE: char = '}';
const COMMA: char = ',';
const NEWLINE: char = '\n';

/// Write a parsed and potentially modified `json::Element` tree to a buffer formatted
/// for human readability.
///
/// The JSON is formatted so that each element is indented and put on it's own line.
pub struct Pretty<'a, 'bin> {
    elem: &'a Element<'bin>,
}

impl<'a, 'bin> Pretty<'a, 'bin> {
    pub fn new(elem: &'a Element<'bin>) -> Self {
        Pretty { elem }
    }
}

impl fmt::Display for Pretty<'_, '_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut stack = vec![State::Root(self.elem)];
        let mut write_comma = WriteComma::Skip;

        // This outer loop drives the pushing of elements on to the stack.
        loop {
            // This inner loop drives the popping of states off the stack.
            let elem = loop {
                let depth = stack.len();

                let Some(mut state) = stack.pop() else {
                    // If the stack is empty, we're done writing.
                    return Ok(());
                };

                let elem = match &mut state {
                    State::Root(elem) => {
                        // The root `Element` never needs a comma written after it.
                        // It's either an opening compound object such as an array or object; or it's
                        // a single Value.
                        write_elem(elem, f)?;
                        elem
                    }
                    State::Array(iter) => {
                        let Some(elem) = iter.next() else {
                            // If there is no next `Element` then we need to move up the stack
                            // and check there for a next `Element`.
                            write_nl_and_indent(depth - 1, f)?;
                            f.write_char(ARRAY_CLOSE)?;
                            continue;
                        };

                        if let WriteComma::Write = write_comma {
                            f.write_char(COMMA)?;
                        }
                        write_nl_and_indent(depth, f)?;
                        write_comma = write_elem(elem, f)?;
                        elem
                    }
                    State::Object(iter) => {
                        let Some(field) = iter.next() else {
                            // If there is no next `Element` then we need to move up the stack
                            // and check there for a next `Element`.
                            write_nl_and_indent(depth - 1, f)?;
                            f.write_char(OBJECT_CLOSE)?;
                            continue;
                        };

                        if let WriteComma::Write = write_comma {
                            f.write_char(COMMA)?;
                        }
                        write_nl_and_indent(depth, f)?;
                        write_comma = write_field(field, f)?;
                        field.element()
                    }
                };

                match &state {
                    State::Array(_) | State::Object(_) => stack.push(state),
                    State::Root(_) => (),
                }

                break elem;
            };

            match elem.value() {
                Value::Array(elements) => stack.push(State::Array(elements.iter())),
                Value::Object(fields) => stack.push(State::Object(fields.iter())),
                _ => (),
            }
        }
    }
}

/// The single stack level.
#[derive(Debug)]
enum State<'a, 'bin> {
    /// The root element.
    Root(&'a Element<'bin>),

    /// A collection of Array `Element`s.
    Array(std::slice::Iter<'a, Element<'bin>>),

    /// A collection of Object `Field`s.
    Object(std::slice::Iter<'a, Field<'bin>>),
}

/// Write a newline and indent ready for the next `Element`'s `Value` to be written.
fn write_nl_and_indent(depth: usize, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    f.write_char(NEWLINE)?;

    for _ in 0..depth {
        f.write_str(TAB)?;
    }

    Ok(())
}

/// Whether or not a comma should be interspersed between values.
#[derive(Copy, Clone)]
enum WriteComma {
    /// Write a comma on the next iteration.
    Write,

    /// Skip writing a comma on the next iteration.
    Skip,
}

/// Shallow write an `Element` and return whether a comma should be written on the next iteration.
///
/// If the `Element` is a compound object like an `Array` or `Object`, only the opening brace is written.
/// This is done to avoid implementing a recursive write fn.
fn write_elem(elem: &Element<'_>, f: &mut fmt::Formatter<'_>) -> Result<WriteComma, fmt::Error> {
    match elem.value() {
        Value::Null => {
            f.write_str("null")?;
            Ok(WriteComma::Write)
        }
        Value::True => {
            f.write_str("true")?;
            Ok(WriteComma::Write)
        }
        Value::False => {
            f.write_str("false")?;
            Ok(WriteComma::Write)
        }
        Value::String(s) => {
            write!(f, "\"{s}\"")?;
            Ok(WriteComma::Write)
        }
        Value::Number(n) => {
            write!(f, "{n}")?;
            Ok(WriteComma::Write)
        }
        Value::Array(_) => {
            f.write_char(ARRAY_OPEN)?;
            Ok(WriteComma::Skip)
        }
        Value::Object(_) => {
            f.write_char(OBJECT_OPEN)?;
            Ok(WriteComma::Skip)
        }
    }
}

/// Write a `Field`'s key and value and return whether a comma should be written on the next iteration.
fn write_field(field: &Field<'_>, f: &mut fmt::Formatter<'_>) -> Result<WriteComma, fmt::Error> {
    write!(f, "\"{}\": ", field.key())?;
    write_elem(field.element(), f)
}

#[cfg(test)]
mod test_tree_writer {
    use crate::{json, test};

    use super::Pretty;

    #[test]
    fn should_pretty_print_root_null() {
        const JSON_IN: &str = "null";

        test::setup();
        let elem = json::parse(JSON_IN).unwrap();
        let pretty = Pretty::new(&elem);

        let s = format!("{pretty}");
        assert_eq!(s, JSON_IN);
        let _ignored = serde_json::from_str::<serde_json::Value>(&s).unwrap();
    }

    #[test]
    fn should_pretty_print_root_bool() {
        const JSON_IN: &str = "true";

        test::setup();
        let elem = json::parse(JSON_IN).unwrap();
        let pretty = Pretty::new(&elem);

        let s = format!("{pretty}");
        assert_eq!(s, JSON_IN);
        let _ignored = serde_json::from_str::<serde_json::Value>(&s).unwrap();
    }

    #[test]
    fn should_pretty_print_root_string() {
        const JSON_IN: &str = r#""one""#;

        test::setup();
        let elem = json::parse(JSON_IN).unwrap();
        let pretty = Pretty::new(&elem);

        let s = format!("{pretty}");
        assert_eq!(s, JSON_IN);
        let _ignored = serde_json::from_str::<serde_json::Value>(&s).unwrap();
    }

    #[test]
    fn should_pretty_print_escaped_str() {
        const JSON_IN: &str = r#""one\ntwo\u0021three\"four\"five""#;

        test::setup();
        let elem = json::parse(JSON_IN).unwrap();
        let pretty = Pretty::new(&elem);

        let s = format!("{pretty}");
        assert_eq!(s, JSON_IN);
        let _ignored = serde_json::from_str::<serde_json::Value>(&s).unwrap();
    }

    #[test]
    fn should_pretty_print_empty_array() {
        const JSON_IN: &str = "[]";
        const JSON_OUT: &str = "[
]";

        test::setup();
        let elem = json::parse(JSON_IN).unwrap();
        let pretty = Pretty::new(&elem);

        let s = format!("{pretty}");
        assert_eq!(s, JSON_OUT);
        let _ignored = serde_json::from_str::<serde_json::Value>(&s).unwrap();
    }

    #[test]
    fn should_pretty_print_string_array() {
        const JSON_IN: &str = r#"["a", "b", "c"]"#;
        const JSON_OUT: &str = r#"[
    "a",
    "b",
    "c"
]"#;

        test::setup();
        let elem = json::parse(JSON_IN).unwrap();
        let pretty = Pretty::new(&elem);

        let s = format!("{pretty}");
        assert_eq!(s, JSON_OUT);
        let _ignored = serde_json::from_str::<serde_json::Value>(&s).unwrap();
    }

    #[test]
    fn should_pretty_print_float_array() {
        const JSON_IN: &str = "[0.000564, 1.0, 3.14159, 2.71828182845904523536028747135266249775724709369995957496696]";
        const JSON_OUT: &str = "[
    0.000564,
    1.0,
    3.14159,
    2.71828182845904523536028747135266249775724709369995957496696
]";

        test::setup();
        let elem = json::parse(JSON_IN).unwrap();
        let pretty = Pretty::new(&elem);

        let s = format!("{pretty}");
        assert_eq!(s, JSON_OUT);
        let _ignored = serde_json::from_str::<serde_json::Value>(&s).unwrap();
    }

    #[test]
    fn should_pretty_print_empty_object() {
        const JSON_IN: &str = "{}";
        const JSON_OUT: &str = "{
}";

        test::setup();
        let elem = json::parse(JSON_IN).unwrap();
        let pretty = Pretty::new(&elem);

        let s = format!("{pretty}");
        assert_eq!(s, JSON_OUT);
        let _ignored = serde_json::from_str::<serde_json::Value>(&s).unwrap();
    }

    #[test]
    fn should_pretty_print_object_with_bools() {
        const JSON_IN: &str = r#"{"one":true,"two":false}"#;

        const JSON_OUT: &str = r#"{
    "one": true,
    "two": false
}"#;

        test::setup();
        let elem = json::parse(JSON_IN).unwrap();
        let pretty = Pretty::new(&elem);

        let s = format!("{pretty}");
        assert_eq!(s, JSON_OUT);
        let _ignored = serde_json::from_str::<serde_json::Value>(&s).unwrap();
    }

    #[test]
    fn should_pretty_print_object_with_nulls() {
        const JSON_IN: &str = r#"{"one":null,"two":null}"#;

        const JSON_OUT: &str = r#"{
    "one": null,
    "two": null
}"#;

        test::setup();
        let elem = json::parse(JSON_IN).unwrap();
        let pretty = Pretty::new(&elem);

        let s = format!("{pretty}");
        assert_eq!(s, JSON_OUT);
        let _ignored = serde_json::from_str::<serde_json::Value>(&s).unwrap();
    }

    #[test]
    fn should_pretty_print_root_object() {
        const JSON_IN: &str =
            r#"{"one": 1,"two": ["a", "b", "c"],"three": {"d": 4,"e": 5,"f": 6},"four":null}"#;

        const JSON_OUT: &str = r#"{
    "one": 1,
    "two": [
        "a",
        "b",
        "c"
    ],
    "three": {
        "d": 4,
        "e": 5,
        "f": 6
    },
    "four": null
}"#;

        test::setup();
        let elem = json::parse(JSON_IN).unwrap();
        let pretty = Pretty::new(&elem);

        let s = format!("{pretty}");
        assert_eq!(s, JSON_OUT);
        let _ignored = serde_json::from_str::<serde_json::Value>(&s).unwrap();
    }
}