use crate::{
Result,
ast::{
BinaryOp, Block, Const, Expr, ExprKind, Function, Item, MatchArm, Module, Pattern,
Statement, Test, TypeDecl, TypeRef, UnaryOp, Validator,
},
parser::parse_module,
};
pub fn format_source(file: &str, source: &str) -> Result<String> {
let module = parse_module(file, source).map_err(|error| error.with_source(source))?;
Ok(format_module(&module))
}
fn format_module(module: &Module) -> String {
let mut rendered = String::new();
rendered.push_str(&format!("module {};\n", module.name));
for item in &module.items {
rendered.push('\n');
match item {
Item::Type(type_decl) => format_type_decl(type_decl, &mut rendered),
Item::Const(constant) => format_const(constant, &mut rendered),
Item::Function(function) => format_function(function, &mut rendered),
Item::Validator(validator) => format_validator(validator, &mut rendered),
Item::Test(test) => format_test(test, &mut rendered),
}
}
rendered
}
fn format_type_decl(type_decl: &TypeDecl, rendered: &mut String) {
rendered.push_str(&format!("type {} {{\n", type_decl.name));
for (index, variant) in type_decl.variants.iter().enumerate() {
rendered.push_str(" ");
rendered.push_str(&variant.name);
if !variant.fields.is_empty() {
let fields = variant
.fields
.iter()
.map(|field| match &field.name {
Some(name) => format!("{name}: {}", format_type_ref(&field.ty)),
None => format_type_ref(&field.ty),
})
.collect::<Vec<_>>()
.join(", ");
rendered.push_str(&format!("({fields})"));
}
if index + 1 != type_decl.variants.len() {
rendered.push(',');
}
rendered.push('\n');
}
rendered.push('}');
rendered.push('\n');
}
fn format_const(constant: &Const, rendered: &mut String) {
rendered.push_str(&format!(
"const {}: {} = {};\n",
constant.name,
format_type_ref(&constant.ty),
format_expr(&constant.value)
));
}
fn format_function(function: &Function, rendered: &mut String) {
let params = format_params(&function.params);
rendered.push_str(&format!(
"fn {}({params}) -> {} ",
function.name,
format_type_ref(&function.return_type)
));
format_block(&function.body, 0, rendered);
rendered.push('\n');
}
fn format_validator(validator: &Validator, rendered: &mut String) {
let params = format_params(&validator.params);
rendered.push_str(&format!("validator {}({params}) ", validator.name));
format_block(&validator.body, 0, rendered);
rendered.push('\n');
}
fn format_test(test: &Test, rendered: &mut String) {
if test.should_fail {
rendered.push_str(&format!("test {} fails ", test.name));
} else {
rendered.push_str(&format!("test {} ", test.name));
}
format_block(&test.body, 0, rendered);
rendered.push('\n');
}
fn format_params(params: &[crate::ast::Param]) -> String {
params
.iter()
.map(|param| format!("{}: {}", param.name, format_type_ref(¶m.ty)))
.collect::<Vec<_>>()
.join(", ")
}
fn format_type_ref(ty: &TypeRef) -> String {
if ty.args.is_empty() {
ty.name.clone()
} else {
let args = ty
.args
.iter()
.map(format_type_ref)
.collect::<Vec<_>>()
.join(", ");
format!("{}<{args}>", ty.name)
}
}
fn format_block(block: &Block, indent: usize, rendered: &mut String) {
rendered.push_str("{\n");
for statement in &block.statements {
format_statement(statement, indent + 1, rendered);
}
rendered.push_str(&spaces(indent));
rendered.push('}');
}
fn format_statement(statement: &Statement, indent: usize, rendered: &mut String) {
rendered.push_str(&spaces(indent));
match statement {
Statement::Let {
name, ty, value, ..
} => {
if let Some(ty) = ty {
rendered.push_str(&format!(
"let {name}: {} = {};\n",
format_type_ref(ty),
format_expr(value)
));
} else {
rendered.push_str(&format!("let {name} = {};\n", format_expr(value)));
}
}
Statement::Require { condition, .. } => {
rendered.push_str(&format!("require {};\n", format_expr(condition)));
}
Statement::Trace { message, .. } => {
rendered.push_str(&format!("trace {};\n", format_expr(message)));
}
Statement::Return { value, .. } => {
rendered.push_str(&format!(
"return {};\n",
indent_multiline(&format_expr(value), indent)
));
}
Statement::Expr {
value,
trailing_semicolon,
..
} => {
rendered.push_str(&indent_multiline(&format_expr(value), indent));
if *trailing_semicolon {
rendered.push(';');
}
rendered.push('\n');
}
}
}
fn format_expr(expr: &Expr) -> String {
match &expr.kind {
ExprKind::Bool(value) => value.to_string(),
ExprKind::Int(value) => value.to_string(),
ExprKind::String(value) => format!("\"{}\"", escape(value)),
ExprKind::ByteArray(hex) => format!("#{hex}"),
ExprKind::Unit => "()".to_string(),
ExprKind::Fail => "fail".to_string(),
ExprKind::List(items) => {
let items = items.iter().map(format_expr).collect::<Vec<_>>().join(", ");
format!("[{items}]")
}
ExprKind::Variable(name) => name.clone(),
ExprKind::Unary { op, expr } => match op {
UnaryOp::Not => format!("!{}", format_operand(expr)),
UnaryOp::Negate => format!("-{}", format_operand(expr)),
},
ExprKind::Binary { left, op, right } => {
format!(
"{} {} {}",
format_operand(left),
binary_op(*op),
format_operand(right)
)
}
ExprKind::Call { callee, args } => {
let args = args.iter().map(format_expr).collect::<Vec<_>>().join(", ");
format!("{callee}({args})")
}
ExprKind::If {
condition,
then_branch,
else_branch,
} => {
let mut rendered = format!("if {} ", format_expr(condition));
format_block(then_branch, 0, &mut rendered);
rendered.push_str(" else ");
format_block(else_branch, 0, &mut rendered);
rendered
}
ExprKind::Match { subject, arms } => format_match(subject, arms),
}
}
fn format_match(subject: &Expr, arms: &[MatchArm]) -> String {
let mut rendered = format!("match {} {{\n", format_expr(subject));
for (index, arm) in arms.iter().enumerate() {
rendered.push_str(" ");
rendered.push_str(&format_pattern(&arm.pattern));
rendered.push_str(" => ");
if let Some(expr) = tail_expr(&arm.body) {
rendered.push_str(&format_expr(expr));
} else {
format_block(&arm.body, 1, &mut rendered);
}
if index + 1 != arms.len() {
rendered.push(',');
}
rendered.push('\n');
}
rendered.push('}');
rendered
}
fn tail_expr(block: &Block) -> Option<&Expr> {
match block.statements.as_slice() {
[
Statement::Expr {
value,
trailing_semicolon: false,
..
},
] => Some(value),
_ => None,
}
}
fn format_pattern(pattern: &Pattern) -> String {
match pattern {
Pattern::Wildcard { .. } => "_".to_string(),
Pattern::Variable { name, .. } => name.clone(),
Pattern::Constructor { name, bindings, .. } => {
if bindings.is_empty() {
name.clone()
} else {
let bindings = bindings
.iter()
.map(|binding| binding.name.as_deref().unwrap_or("_"))
.collect::<Vec<_>>()
.join(", ");
format!("{name}({bindings})")
}
}
Pattern::Bool { value, .. } => value.to_string(),
Pattern::Int { value, .. } => value.to_string(),
Pattern::ByteArray { hex, .. } => format!("#{hex}"),
Pattern::String { value, .. } => format!("\"{}\"", escape(value)),
Pattern::Unit { .. } => "()".to_string(),
}
}
fn format_operand(expr: &Expr) -> String {
match expr.kind {
ExprKind::Binary { .. } | ExprKind::If { .. } | ExprKind::Match { .. } => {
format!("({})", format_expr(expr))
}
_ => format_expr(expr),
}
}
fn binary_op(op: BinaryOp) -> &'static str {
match op {
BinaryOp::Or => "||",
BinaryOp::And => "&&",
BinaryOp::Equal => "==",
BinaryOp::NotEqual => "!=",
BinaryOp::Less => "<",
BinaryOp::LessEqual => "<=",
BinaryOp::Greater => ">",
BinaryOp::GreaterEqual => ">=",
BinaryOp::Add => "+",
BinaryOp::Subtract => "-",
BinaryOp::Multiply => "*",
BinaryOp::Divide => "/",
BinaryOp::Remainder => "%",
}
}
fn spaces(indent: usize) -> String {
" ".repeat(indent)
}
fn indent_multiline(value: &str, indent: usize) -> String {
value.replace('\n', &format!("\n{}", spaces(indent)))
}
fn escape(value: &str) -> String {
value
.chars()
.flat_map(|ch| match ch {
'\\' => "\\\\".chars().collect::<Vec<_>>(),
'"' => "\\\"".chars().collect(),
'\n' => "\\n".chars().collect(),
'\r' => "\\r".chars().collect(),
'\t' => "\\t".chars().collect(),
_ => vec![ch],
})
.collect()
}