schala/schala-lang/language/src/ast/visualize.rs

use super::{
    Block, Declaration, Expression, ExpressionKind, ImportSpecifier, InvocationArgument, ModuleSpecifier,
    Signature, Statement, StatementKind, AST,
};

const LEVEL: usize = 2;

fn do_indent(n: usize, buf: &mut String) {
    for _ in 0..n {
        buf.push(' ');
    }
}

fn newline(buf: &mut String) {
    buf.push('\n');
}

pub(super) fn render_ast(ast: &AST) -> String {
    let AST { statements, .. } = ast;

    let mut buf = format!("(AST\n");
    render_block(&statements, LEVEL, &mut buf);

    buf.push(')');

    buf
}

fn render_statement(stmt: &Statement, indent: usize, buf: &mut String) {
    use StatementKind::*;
    do_indent(indent, buf);
    match stmt.kind {
        Expression(ref expr) => render_expression(expr, indent, buf),
        Declaration(ref decl) => render_declaration(decl, indent, buf),
        Import(ref spec) => render_import(spec, indent, buf),
        Module(ref spec) => render_module(spec, indent, buf),
    }
}

fn render_expression(expr: &Expression, indent: usize, buf: &mut String) {
    use ExpressionKind::*;

    buf.push_str("(Expr ");
    match &expr.kind {
        NatLiteral(n) => buf.push_str(&format!("(NatLiteral {})", n)),
        FloatLiteral(f) => buf.push_str(&format!("(FloatLiteral {})", f)),
        StringLiteral(s) => buf.push_str(&format!("(StringLiteral {})", s)),
        BoolLiteral(b) => buf.push_str(&format!("(BoolLiteral {})", b)),
        BinExp(binop, lhs, rhs) => {
            let new_indent = indent + LEVEL;
            buf.push_str(&format!("Binop {}\n", binop.sigil()));

            do_indent(new_indent, buf);
            render_expression(lhs, new_indent, buf);
            newline(buf);

            do_indent(new_indent, buf);
            render_expression(rhs, new_indent, buf);
            newline(buf);

            do_indent(indent, buf);
        }
        PrefixExp(prefix, expr) => {
            let new_indent = indent + LEVEL;
            buf.push_str(&format!("PrefixOp {}\n", prefix.sigil()));

            do_indent(new_indent, buf);
            render_expression(expr, new_indent, buf);
            newline(buf);

            do_indent(indent, buf);
        }
        TupleLiteral(..) => (),
        Value(name) => {
            buf.push_str(&format!("Value {})", name));
        }
        NamedStruct { name: _, fields: _ } => (),
        Call { f, arguments } => {
            let new_indent = indent + LEVEL;
            buf.push_str("Call ");
            render_expression(f, new_indent, buf);
            newline(buf);

            for arg in arguments {
                do_indent(new_indent, buf);
                match arg {
                    InvocationArgument::Positional(expr) => render_expression(expr, new_indent, buf),
                    InvocationArgument::Keyword { .. } => buf.push_str("<keyword>"),
                    InvocationArgument::Ignored => buf.push_str("<ignored>"),
                }
                newline(buf);
                do_indent(indent, buf);
            }
        }
        Index { .. } => buf.push_str("<index>"),
        IfExpression { .. } => buf.push_str("<if-expr>"),
        WhileExpression { .. } => buf.push_str("<while-expr>"),
        ForExpression { .. } => buf.push_str("<for-expr>"),
        Lambda { params, type_anno: _, body } => {
            let new_indent = indent + LEVEL;
            buf.push_str("Lambda ");
            newline(buf);

            do_indent(new_indent, buf);
            buf.push_str("(Args ");
            for p in params {
                buf.push_str(&format!("{} ", p.name));
            }
            buf.push(')');
            newline(buf);

            do_indent(new_indent, buf);
            buf.push_str("(Body ");
            newline(buf);
            render_block(body, new_indent + LEVEL, buf);
            do_indent(new_indent, buf);
            buf.push(')');

            newline(buf);
            do_indent(indent, buf);
        }
        ListLiteral(..) => buf.push_str("<list-literal>"),
    }
    buf.push(')');
}
fn render_declaration(decl: &Declaration, indent: usize, buf: &mut String) {
    use Declaration::*;

    buf.push_str("(Decl ");
    match decl {
        FuncSig(ref sig) => render_signature(sig, indent, buf),
        FuncDecl(ref sig, ref block) => {
            let indent = indent + LEVEL;
            buf.push_str("Function");
            newline(buf);

            do_indent(indent, buf);
            render_signature(sig, indent, buf);
            newline(buf);

            do_indent(indent, buf);
            buf.push_str("(Body");
            newline(buf);

            render_block(block, indent + LEVEL, buf);
            do_indent(indent, buf);
            buf.push_str(")");
            newline(buf);
        }
        TypeDecl { name: _, body: _, .. } => {
            buf.push_str("<type-decl>");
        }
        TypeAlias { alias: _, original: _ } => {
            buf.push_str("<type-alias>");
        }
        Binding { name, constant: _, type_anno: _, expr } => {
            let new_indent = indent + LEVEL;
            buf.push_str(&format!("Binding {}", name));
            newline(buf);
            do_indent(new_indent, buf);
            render_expression(expr, new_indent, buf);
            newline(buf);
        }
        _ => (), /*
                 Impl { type_name: TypeIdentifier, interface_name: Option<TypeSingletonName>, block: Vec<Declaration> },
                 Interface { name: Rc<String>, signatures: Vec<Signature> },
                 Annotation { name: Rc<String>, arguments: Vec<Expression> },
                     */
    }
    do_indent(indent, buf);
    buf.push_str(")");
}

fn render_block(block: &Block, indent: usize, buf: &mut String) {
    for stmt in block.statements.iter() {
        render_statement(&stmt, indent, buf);
        newline(buf);
    }
}

fn render_signature(sig: &Signature, _indent: usize, buf: &mut String) {
    buf.push_str(&format!("(Signature {} )", sig.name));
}

fn render_import(_import: &ImportSpecifier, _indent: usize, buf: &mut String) {
    buf.push_str("(Import <some import>)");
}
fn render_module(_expr: &ModuleSpecifier, _indent: usize, buf: &mut String) {
    buf.push_str("(Module <some mod>)");
}

#[cfg(test)]
mod test {
    use super::render_ast;
    use crate::util::quick_ast;

    #[test]
    fn test_visualization() {
        let ast = quick_ast(
            r#"
    fn test(x) {
        let m = 9
        1 * 4 <> m |> somemod::output(x)
    }

    let quincy = \(no, yes, maybe) {
        let a = 10
        yes * no + a
    }

    let b = 54
    test(b) == 3
    "#,
        );

        let expected_output = r#"(AST
  (Decl Function
    (Signature test )
    (Body
      (Decl Binding m
        (Expr (NatLiteral 9))
      )
      (Expr Binop *
        (Expr (NatLiteral 1))
        (Expr Binop |>
          (Expr Binop <>
            (Expr (NatLiteral 4))
            (Expr Value m))
          )
          (Expr Call (Expr Value somemod::output))
            (Expr Value x))
          )
        )
      )
    )
  )
  (Decl Binding quincy
    (Expr Lambda 
      (Args no yes maybe )
      (Body 
        (Decl Binding a
          (Expr (NatLiteral 10))
        )
        (Expr Binop +
          (Expr Binop *
            (Expr Value yes))
            (Expr Value no))
          )
          (Expr Value a))
        )
      )
    )
  )
  (Decl Binding b
    (Expr (NatLiteral 54))
  )
  (Expr Binop ==
    (Expr Call (Expr Value test))
      (Expr Value b))
    )
    (Expr (NatLiteral 3))
  )
)"#;

        let rendered = render_ast(&ast);
        assert_eq!(rendered, expected_output);
    }
}