ocpi_tariffs/json/

write.rs

//! Tools for writing JSON.

#[cfg(test)]
mod test_tree_writer;

use std::fmt::{self, Write as _};

use super::{Document, 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 its own line.
pub struct Pretty<'a, 'buf> {
    elem: &'a Element<'buf>,
}

impl<'a, 'buf> Pretty<'a, 'buf> {
    /// Pretty print a whole document, starting from its root element.
    pub fn from_doc(doc: &'a Document<'buf>) -> Self {
        Pretty { elem: doc.root() }
    }
}

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`.
                            let Some(depth) = depth.checked_sub(1) else {
                                return Ok(());
                            };
                            write_nl_and_indent(depth, 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`.
                            let Some(depth) = depth.checked_sub(1) else {
                                return Ok(());
                            };
                            write_nl_and_indent(depth, 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;
            };

            if let Value::Array(elements) = elem.value() {
                stack.push(State::Array(elements.iter()));
            } else if let Value::Object(fields) = elem.value() {
                stack.push(State::Object(fields.iter()));
            }
        }
    }
}

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

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

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

/// 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.as_unescaped_str())?;
            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().as_unescaped_str())?;
    write_elem(field.element(), f)
}