use crate::{
ast::{
ApplicationExpr, Ast, BinOpExpr, CondExpr, ConstructExpr, Expr, FunExpr, LetInExpr,
LiteralExpr, Pattern, PatternMatchExpr, TupleExpr, VarExpr,
},
lexer::Lexer,
terminal_colors::{BLUE, END, YELLOW},
};
pub struct AstPrinter<'a> {
pub lexer: &'a Lexer,
indent: String,
}
impl<'a> AstPrinter<'a> {
pub fn new(lexer: &'a Lexer) -> Self {
Self {
lexer,
indent: String::new(),
}
}
pub fn pretty_print(mut self, ast: &Ast) {
for binding in &ast.binds {
let name = if let Some(name) = &binding.name {
self.lexer.str_from_span(name)
} else {
"()"
};
println!("{YELLOW}{name}{END} = ");
self.pretty_print_expr(&binding.expr.borrow());
println!();
self.indent = String::new();
}
}
pub fn pretty_print_expr(&mut self, expr: &Expr) {
match expr {
Expr::Literal(literal_expr) => self.pretty_print_literal_expr(literal_expr),
Expr::Var(var_expr) => self.pretty_print_var_expr(var_expr),
Expr::Fun(fun_expr) => self.pretty_print_fun_expr(fun_expr),
Expr::Tuple(tuple_expr) => self.pretty_print_tuple_expr(tuple_expr),
Expr::Construction(construct_expr) => self.pretty_print_cons_expr(construct_expr),
Expr::Application(application_expr) => {
self.pretty_print_application_expr(application_expr)
}
Expr::LetIn(let_in_expr) => self.pretty_print_let_in_expr(let_in_expr),
Expr::BinOp(bin_op_expr) => self.pretty_print_bin_op_expr(bin_op_expr),
Expr::Conditional(cond_expr) => self.pretty_print_cond_expr(cond_expr),
Expr::PatternMatch(pattern_match_expr) => {
self.pretty_print_pattern_match_expr(pattern_match_expr)
}
}
}
fn pretty_print_cons_expr(&mut self, construct_expr: &ConstructExpr) {
println!("{}{BLUE}ConstructExpr{END}", self.indent);
let cons = self.lexer.str_from_span(&construct_expr.cons);
if let Some(arg) = &construct_expr.arg {
println!("{}├─▸ constructor: {}", self.indent, cons);
println!("{}└─▸ arg:", self.indent);
self.indent += " ";
self.pretty_print_expr(&arg.borrow());
} else {
println!("{}└─▸ constructor: {}", self.indent, cons);
}
}
fn pretty_print_fun_expr(&mut self, fun_expr: &FunExpr) {
println!("{}{BLUE}FunExpr{END}", self.indent);
let parameters = fun_expr
.params
.iter()
.map(|p| self.lexer.str_from_span(p))
.collect::<Vec<&str>>()
.join(", ");
let captures = fun_expr
.captures
.iter()
.map(|s| &s[..])
.collect::<Vec<&str>>()
.join(", ");
let recursive = match fun_expr.recursive_bind {
Some(_) => "yes",
None => "no",
};
println!("{}├─▸ parameters: [{parameters}]", self.indent);
println!("{}├─▸ captures: [{captures}]", self.indent);
println!("{}├─▸ recursive: {recursive}", self.indent);
println!("{}└─▸ body:", self.indent);
self.indent += " ";
self.pretty_print_expr(&fun_expr.body.borrow());
}
pub fn pretty_print_tuple_expr(&mut self, tuple_expr: &TupleExpr) {
println!("{}{BLUE}TupleExpr{END}", self.indent);
println!("{}└─▸ elements:", self.indent);
let orig_indent = self.indent.clone();
for (i, element) in tuple_expr.elements.iter().enumerate() {
self.indent = orig_indent.clone();
if i < tuple_expr.elements.len() - 1 {
println!("{} ├─▸ ({i})", self.indent);
self.indent += " │ ";
} else {
println!("{} └─▸ ({i})", self.indent);
self.indent += " ";
}
self.pretty_print_expr(&element.borrow());
}
}
pub fn pretty_print_application_expr(&mut self, application_expr: &ApplicationExpr) {
println!("{}{BLUE}ApplicationExpr{END}", self.indent);
println!("{}├─▸ function:", self.indent);
let orig_indent = self.indent.clone();
self.indent += "│ ";
self.pretty_print_expr(&application_expr.fun.borrow());
self.indent = orig_indent.clone();
println!("{}└─▸ binds:", self.indent);
for (i, bind) in application_expr.binds.iter().enumerate() {
self.indent = orig_indent.clone();
if i < application_expr.binds.len() - 1 {
println!("{} ├─▸ ({i})", self.indent);
self.indent += " │ ";
} else {
println!("{} └─▸ ({i})", self.indent);
self.indent += " ";
}
self.pretty_print_expr(&bind.borrow());
}
}
fn pretty_print_bin_op_expr(&mut self, bin_op_expr: &BinOpExpr) {
println!("{}{BLUE}BinOpExpr{END}", self.indent);
println!("{}├─▸ operator: {}", self.indent, bin_op_expr.op);
println!("{}├─▸ lhs:", self.indent);
let orig_indent = self.indent.clone();
self.indent += "│ ";
self.pretty_print_expr(&bin_op_expr.lhs.borrow());
self.indent = orig_indent;
println!("{}└─▸ rhs:", self.indent);
self.indent += " ";
self.pretty_print_expr(&bin_op_expr.rhs.borrow());
}
fn pretty_print_cond_expr(&mut self, cond_expr: &CondExpr) {
println!("{}{BLUE}CondExpr{END}", self.indent);
println!("{}├─▸ condition:", self.indent);
let orig_indent = self.indent.clone();
self.indent += "│ ";
self.pretty_print_expr(&cond_expr.cond.borrow());
self.indent = orig_indent.clone();
println!("{}├─▸ then expr:", self.indent);
self.indent += "│ ";
self.pretty_print_expr(&cond_expr.yes.borrow());
self.indent = orig_indent;
println!("{}└─▸ else expr:", self.indent);
self.indent += " ";
self.pretty_print_expr(&cond_expr.no.borrow());
}
fn pretty_print_pattern_match_expr(&mut self, pattern_match_expr: &PatternMatchExpr) {
println!("{}{BLUE}PatternMatchExpr{END}", self.indent);
println!("{}├─▸ matched:", self.indent);
let orig_indent = self.indent.clone();
self.indent += "│ ";
self.pretty_print_expr(&pattern_match_expr.matched.borrow());
self.indent = orig_indent.clone();
println!("{}└─▸ branches:", self.indent);
for (i, branch) in pattern_match_expr.branches.iter().enumerate() {
self.indent = orig_indent.clone();
if i < pattern_match_expr.branches.len() - 1 {
println!(
"{} ├─▸ with {}:",
self.indent,
self.pattern_to_string(&branch.0)
);
self.indent += " │ ";
} else {
println!(
"{} └─▸ with {}:",
self.indent,
self.pattern_to_string(&branch.0)
);
self.indent += " ";
}
self.pretty_print_expr(&branch.1.borrow());
}
}
fn pretty_print_let_in_expr(&mut self, let_in_expr: &LetInExpr) {
let bind_name = self.lexer.str_from_span(&let_in_expr.bind.0);
println!("{}{BLUE}LetInExpr{END}", self.indent);
println!("{}├─▸ bind name: {bind_name}", self.indent);
println!("{}├─▸ bind value:", self.indent);
let orig_indent = self.indent.clone();
self.indent += "│ ";
self.pretty_print_expr(&let_in_expr.bind.1.borrow());
self.indent = orig_indent;
println!("{}└─▸ expr:", self.indent);
self.indent += " ";
self.pretty_print_expr(&let_in_expr.expr.borrow());
}
fn pretty_print_var_expr(&mut self, var_expr: &VarExpr) {
let var = self.lexer.str_from_span(&var_expr.id);
println!("{}{BLUE}VarExpr{END} (\"{var}\")", self.indent);
}
fn pretty_print_literal_expr(&mut self, literal_expr: &LiteralExpr) {
println!(
"{}{BLUE}LiteralExpr{END} ({})",
self.indent,
self.literal_expr_to_string(literal_expr)
);
}
fn literal_expr_to_string(&self, literal_expr: &LiteralExpr) -> String {
match literal_expr {
LiteralExpr::Integer(val, _) => val.to_string(),
LiteralExpr::Unit(_) => "unit".to_string(),
}
}
fn pattern_to_string(&self, pattern: &Pattern) -> String {
match pattern {
Pattern::Tuple(elements) => {
let elements = elements
.iter()
.map(|p| self.pattern_to_string(p))
.collect::<Vec<String>>()
.join(", ");
format!("({elements})")
}
Pattern::Constructor(span, Some(arg)) => format!(
"{} {}",
self.lexer.str_from_span(span),
self.pattern_to_string(arg)
),
Pattern::Constructor(span, None) | Pattern::Identifier(span) => {
self.lexer.str_from_span(span).to_string()
}
Pattern::Literal(literal_expr) => self.literal_expr_to_string(literal_expr),
Pattern::None => "_".to_string(),
}
}
}