Prettier Layout Algorithm

layout.rs

use std::fmt::Write;

// The layout algorithm works as follows.
//
// 1. You walk through the AST and build up a LayoutNode tree (see test for example).
// 2. We walk through the LayoutNode tree. For each Group we check if can be laid
//    out over a single line. If so do that, otherwise splay it over multiple lines.
// 3. Each time we get to a Nested node, increase the indentation level.
//
// This works pretty well, but it suffers from one issue: splayed lines can
// actually be longer in some cases! May not be an issue though.

#[derive(Debug)]
pub enum LayoutNode {
    // A group of nodes - if their flattened length is short enough then
    // output the entire tree with `Newline` replaced with a single space.
    // Otherwise recurse.
    Group(LayoutGroup),
    // Just a list of nodes. Literally nothing happens, but newlines in here are
    // always just literal newlines.
    List(Vec<LayoutNode>),
    // Increase the nesting level for the nodes below this.
    Nest(Box<LayoutNode>),
    // Some literal text with no newlines in it. This will never be split into multiple lines.
    Text(String),
    // Reference to some text in a SystemVerilog source file.
    Locate(sv_parser::Locate),
    // A place where we can add a newline if there isn't enough space.
    // If we do, the first option is inserted followed by a newline,
    // otherwise the second one is.
    // So for example you might have (",", ", ") for a place between arguments
    // or (",", "") for afer the last argument to a function.
    MaybeNewline(&'static str, &'static str),
}

#[derive(Debug)]
pub struct LayoutGroup {
    children: Vec<LayoutNode>,
    flattened_length: usize,
}

pub fn group(children: Vec<LayoutNode>) -> LayoutNode {
    LayoutNode::Group(LayoutGroup {
        children,
        flattened_length: 0,
    })
}

pub fn list(children: Vec<LayoutNode>) -> LayoutNode {
    LayoutNode::List(children)
}

pub fn nest(child: LayoutNode) -> LayoutNode {
    LayoutNode::Nest(Box::new(child))
}

pub fn maybe_newline(if_new_line: &'static str, if_same_line: &'static str) -> LayoutNode {
    LayoutNode::MaybeNewline(if_new_line, if_same_line)
}

pub fn locate(loc: sv_parser::Locate) -> LayoutNode {
    LayoutNode::Locate(loc)
}

#[macro_export]
macro_rules! text {
    ($($arg:tt)*) => {{
        let res = std::fmt::format(std::format_args!($($arg)*));
        LayoutNode::Text(res)
    }}
}

pub fn pretty(root: &LayoutNode, sv_source: &str, max_width: usize) -> String {
    // The general algorithm is to:

    // 1. Walk the tree, starting from the leaves and calculate
    //    the length of each Group if it were entirely flattened
    //    onto one line.
    // 2. Walk the tree a second time. Every time we get to a Group we split
    //    it over multiple lines iff the remaining AST nodes below that point
    //    wouldn't fit on the rest of the line.

    let mut output = String::new();
    layout(root, max_width, 0, false, sv_source, &mut output);
    output
}

/// Calculate the flattened length of a node and return it. Save it in the node
/// for Group nodes.
pub fn calculate_flattened_lengths(node: &mut LayoutNode) -> usize {
    // TODO: This is recursive; change to iterative.
    match node {
        LayoutNode::Group(g) => {
            let len = g.children.iter_mut().map(|c| calculate_flattened_lengths(c)).sum::<usize>();
            g.flattened_length = len;
            len
        }
        LayoutNode::List(l) => {
            l.iter_mut().map(|c| calculate_flattened_lengths(c)).sum::<usize>()
        }
        LayoutNode::Nest(n) => calculate_flattened_lengths(n),
        LayoutNode::Text(t) => t.len(),
        LayoutNode::Locate(l) => l.len,
        LayoutNode::MaybeNewline(_, if_same_line) => if_same_line.len(),
    }
}

// The current state of layout. It always starts out indented by a certain
// amount but at some point we figure out that the rest of the AST below
// a given node can fit on the current line flattened. Then we switch to
// flattened and don't need to know the indentation.


fn layout(
    node: &LayoutNode,
    max_width: usize,
    indent: usize,
    flattened: bool,
    sv_source: &str,
    output: &mut String) {
    match node {
        LayoutNode::Group(g) => {
            let flattened = indent + g.flattened_length <= max_width;

            for child in g.children.iter() {
                layout(child, max_width, indent, flattened, sv_source, output);
            }
        }
        LayoutNode::List(l) => {
            for child in l.iter() {
                layout(child, max_width, indent, flattened, sv_source, output);
            }
        }
        LayoutNode::Nest(n) => {
            layout(n, max_width, indent + 1, flattened, sv_source, output);
        }
        LayoutNode::Text(t) => {
            output.push_str(t);
        }
        LayoutNode::Locate(l) => {
            output.push_str(l.str(sv_source));
        }
        LayoutNode::MaybeNewline(if_new_line, if_same_line) => {
            if flattened {
                output.push_str(if_same_line);
            } else {
                write!(output, "{if_new_line}\n{:indent$}", "", indent=indent * 4).unwrap();
            }
        }
    };
}



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

    #[derive(Clone)]
    enum AstNode {
        Element(Element),
        Text(String),
    }

    #[derive(Clone)]
    struct Element {
        name: String,
        attributes: Vec<Attribute>,
        children: Vec<AstNode>,
    }

    #[derive(Clone)]
    struct Attribute {
        key: String,
        value: String,
    }

    // Render an AST to a LayoutNode tree.
    fn ast_to_layout(node: &AstNode) -> LayoutNode {
        match node {
            AstNode::Element(e) => {
                let mut ns: Vec<LayoutNode> = Vec::new();
                if e.attributes.is_empty() {
                    ns.push(text!("<{}>", e.name));
                } else {
                    let mut attr_nest: Vec<LayoutNode> = Vec::new();
                    for attribute in e.attributes.iter() {
                        attr_nest.push(maybe_newline("", " "));
                        attr_nest.push(text!("\"{}\"=\"{}\"", attribute.key, attribute.value));
                    }
                    ns.push(group(vec![
                        text!("<{}", e.name),
                        nest(list(attr_nest)),
                        maybe_newline("", ""),
                        text!(">"),
                    ]));
                }

                let mut child_nest: Vec<LayoutNode> = Vec::new();
                for child in e.children.iter() {
                    child_nest.push(maybe_newline("", ""));
                    child_nest.push(ast_to_layout(child));
                }

                ns.push(nest(list(child_nest)));
                ns.push(maybe_newline("", ""));
                ns.push(text!("</{}>", e.name));

                group(ns)
            }
            AstNode::Text(t) => LayoutNode::Text(t.clone()),
        }
    }


    fn ast_el(name: &str, attributes: &[Attribute], children: &[AstNode]) -> AstNode {
        AstNode::Element(Element{
            name: name.to_string(),
            attributes: attributes.to_vec(),
            children: children.to_vec(),
        })
    }

    fn ast_attr(key: &str, value: &str) -> Attribute {
        Attribute{ key: key.to_string(), value: value.to_string() }
    }

    fn ast_text(text: &str) -> AstNode {
        AstNode::Text(text.to_string())
    }

    #[test]
    fn xml() {

        let ast = ast_el(
            "p",
            &[
                ast_attr("color", "red"),
                ast_attr("font", "Times"),
                ast_attr("size", "10"),
            ],
            &[
                ast_text("Here is some"),
                ast_el(
                    "em",
                    &[],
                    &[
                        ast_text("emphasized"),
                        ast_el(
                            "b",
                            &[
                                ast_attr("aaaaa", "bbbbb"),
                            ],
                            &[
                                ast_text("and bold"),
                                ast_el("img", &[], &[]),
                                ast_el("div", &[], &[]),
                                ast_el("span", &[], &[]),
                            ],
                        ),
                    ],
                ),
                ast_text("TEXT."),
                ast_text("Here is a"),
                ast_el(
                    "a",
                    &[
                        ast_attr("href", "http://www.eg.com/"),
                    ],
                    &[
                        ast_text("link"),
                    ],
                ),
                ast_text("elsewhere."),
            ],
        );

        let mut layout = ast_to_layout(&ast);

        calculate_flattened_lengths(&mut layout);

        dbg!(&layout);

        for i in [10usize, 20, 40, 60, 80].iter() {

            println!("{}\n\n", pretty(
                &layout,
                "",
                *i,
            ));
        }
    }
}