use crate::formatters::{
trivia_formatter::{self, FormatTriviaType},
CodeFormatter,
};
use full_moon::ast::{Do, ElseIf, GenericFor, If, NumericFor, Repeat, Stmt, While};
use full_moon::tokenizer::TokenReference;
use std::borrow::Cow;
macro_rules! fmt_stmt {
($fmter:expr, $value:ident, { $($(#[$inner:meta])* $operator:ident = $output:ident,)+ }) => {
match $value {
$(
$(#[$inner])*
Stmt::$operator(stmt) => Stmt::$operator($fmter.$output(stmt)),
)+
}
};
}
impl CodeFormatter {
pub fn format_do_block<'ast>(&self, do_block: &Do<'ast>) -> Do<'ast> {
let do_token = crate::fmt_symbol!(self, do_block.do_token(), "do");
let end_token = self.format_end_token(do_block.end_token());
do_block
.to_owned()
.with_do_token(do_token)
.with_end_token(end_token)
}
pub fn format_generic_for<'ast>(&mut self, generic_for: &GenericFor<'ast>) -> GenericFor<'ast> {
let for_token = crate::fmt_symbol!(self, generic_for.for_token(), "for ");
let (formatted_names, mut names_comments_buf) = self.format_punctuated(
generic_for.names(),
&CodeFormatter::format_token_reference_mut,
);
#[cfg(feature = "luau")]
let type_specifiers = generic_for
.type_specifiers()
.map(|x| match x {
Some(type_specifier) => Some(self.format_type_specifier(type_specifier)),
None => None,
})
.collect();
let in_token = crate::fmt_symbol!(self, generic_for.in_token(), " in ");
let (formatted_expr_list, mut expr_comments_buf) =
self.format_punctuated(generic_for.expr_list(), &CodeFormatter::format_expression);
names_comments_buf.append(&mut expr_comments_buf);
let do_token = crate::fmt_symbol!(self, generic_for.do_token(), " do");
let do_token = Cow::Owned(trivia_formatter::token_reference_add_trivia(
do_token.to_owned().into_owned(),
FormatTriviaType::NoChange,
FormatTriviaType::Append(names_comments_buf),
));
let end_token = self.format_end_token(generic_for.end_token());
let generic_for = generic_for
.to_owned()
.with_for_token(for_token)
.with_names(formatted_names)
.with_in_token(in_token)
.with_expr_list(formatted_expr_list)
.with_do_token(do_token)
.with_end_token(end_token);
#[cfg(feature = "luau")]
let generic_for = generic_for.with_type_specifiers(type_specifiers);
generic_for
}
fn format_else_if<'ast>(&mut self, else_if_node: &ElseIf<'ast>) -> ElseIf<'ast> {
let formatted_else_if_token =
crate::fmt_symbol!(self, else_if_node.else_if_token(), "elseif ");
let formatted_condition = self.format_expression(else_if_node.condition());
let formatted_then_token = crate::fmt_symbol!(self, else_if_node.then_token(), " then");
else_if_node
.to_owned()
.with_else_if_token(formatted_else_if_token)
.with_condition(formatted_condition)
.with_then_token(formatted_then_token)
}
pub fn format_if<'ast>(&mut self, if_node: &If<'ast>) -> If<'ast> {
let formatted_if_token = crate::fmt_symbol!(self, if_node.if_token(), "if ");
let formatted_condition = self.format_expression(if_node.condition());
let formatted_then_token = crate::fmt_symbol!(self, if_node.then_token(), " then");
let formatted_end_token = self.format_end_token(if_node.end_token());
let formatted_else_if = match if_node.else_if() {
Some(else_if) => Some(
else_if
.iter()
.map(|else_if| self.format_else_if(else_if))
.collect(),
),
None => None,
};
let formatted_else_token = match if_node.else_token() {
Some(token) => Some(crate::fmt_symbol!(self, token, "else")),
None => None,
};
if_node
.to_owned()
.with_if_token(formatted_if_token)
.with_condition(formatted_condition)
.with_then_token(formatted_then_token)
.with_else_if(formatted_else_if)
.with_else_token(formatted_else_token)
.with_end_token(formatted_end_token)
}
pub fn format_numeric_for<'ast>(&mut self, numeric_for: &NumericFor<'ast>) -> NumericFor<'ast> {
let for_token = crate::fmt_symbol!(self, numeric_for.for_token(), "for ");
let formatted_index_variable =
Cow::Owned(self.format_plain_token_reference(numeric_for.index_variable()));
#[cfg(feature = "luau")]
let type_specifier = match numeric_for.type_specifier() {
Some(type_specifier) => Some(self.format_type_specifier(type_specifier)),
None => None,
};
let equal_token = crate::fmt_symbol!(self, numeric_for.equal_token(), " = ");
let formatted_start_expression = self.format_expression(numeric_for.start());
let start_end_comma = crate::fmt_symbol!(self, numeric_for.start_end_comma(), ", ");
let formatted_end_expression = self.format_expression(numeric_for.end());
let (end_step_comma, formatted_step_expression) = match numeric_for.step() {
Some(step) => (
Some(crate::fmt_symbol!(
self,
numeric_for.end_step_comma().unwrap(),
", "
)),
Some(self.format_expression(step)),
),
None => (None, None),
};
let do_token = crate::fmt_symbol!(self, numeric_for.do_token(), " do");
let end_token = self.format_end_token(numeric_for.end_token());
let numeric_for = numeric_for
.to_owned()
.with_for_token(for_token)
.with_index_variable(formatted_index_variable)
.with_equal_token(equal_token)
.with_start(formatted_start_expression)
.with_start_end_comma(start_end_comma)
.with_end(formatted_end_expression)
.with_end_step_comma(end_step_comma)
.with_step(formatted_step_expression)
.with_do_token(do_token)
.with_end_token(end_token);
#[cfg(feature = "luau")]
let numeric_for = numeric_for.with_type_specifier(type_specifier);
numeric_for
}
pub fn format_repeat_block<'ast>(&mut self, repeat_block: &Repeat<'ast>) -> Repeat<'ast> {
let repeat_token = crate::fmt_symbol!(self, repeat_block.repeat_token(), "repeat");
let until_token = crate::fmt_symbol!(self, repeat_block.until_token(), "until ");
let formatted_until = self.format_expression(repeat_block.until());
repeat_block
.to_owned()
.with_repeat_token(repeat_token)
.with_until_token(until_token)
.with_until(formatted_until)
}
pub fn format_while_block<'ast>(&mut self, while_block: &While<'ast>) -> While<'ast> {
let while_token = crate::fmt_symbol!(self, while_block.while_token(), "while ");
let formatted_condition = self.format_expression(while_block.condition());
let do_token = crate::fmt_symbol!(self, while_block.do_token(), " do");
let end_token = self.format_end_token(while_block.end_token());
while_block
.to_owned()
.with_while_token(while_token)
.with_condition(formatted_condition)
.with_do_token(do_token)
.with_end_token(end_token)
}
pub fn format_stmt<'ast>(&mut self, stmt: &Stmt<'ast>) -> Stmt<'ast> {
fmt_stmt!(self, stmt, {
Assignment = format_assignment,
Do = format_do_block,
FunctionCall = format_function_call,
FunctionDeclaration = format_function_declaration,
GenericFor = format_generic_for,
If = format_if,
LocalAssignment = format_local_assignment,
LocalFunction = format_local_function,
NumericFor = format_numeric_for,
Repeat = format_repeat_block,
While = format_while_block,
#[cfg(feature = "luau")] CompoundAssignment = format_compound_assignment,
#[cfg(feature = "luau")] ExportedTypeDeclaration = format_exported_type_declaration,
#[cfg(feature = "luau")] TypeDeclaration = format_type_declaration,
})
}
pub fn stmt_add_trivia<'ast>(
&self,
stmt: Stmt<'ast>,
additional_indent_level: Option<usize>,
) -> Stmt<'ast> {
let leading_trivia =
FormatTriviaType::Append(vec![self.create_indent_trivia(additional_indent_level)]);
let trailing_trivia = FormatTriviaType::Append(vec![self.create_newline_trivia()]);
match stmt {
Stmt::Assignment(assignment) => {
Stmt::Assignment(self.assignment_add_trivia(assignment, additional_indent_level))
}
Stmt::Do(do_block) => Stmt::Do(trivia_formatter::do_block_add_trivia(
do_block,
leading_trivia,
trailing_trivia,
)),
Stmt::FunctionCall(function_call) => {
Stmt::FunctionCall(trivia_formatter::function_call_add_trivia(
function_call,
leading_trivia,
trailing_trivia,
))
}
Stmt::FunctionDeclaration(function_declaration) => {
Stmt::FunctionDeclaration(trivia_formatter::function_declaration_add_trivia(
function_declaration,
leading_trivia,
trailing_trivia,
))
}
Stmt::GenericFor(generic_for) => {
Stmt::GenericFor(trivia_formatter::generic_for_add_trivia(
generic_for,
leading_trivia,
trailing_trivia,
))
}
Stmt::If(if_block) => {
Stmt::If(self.if_block_add_trivia(if_block, additional_indent_level))
}
Stmt::LocalAssignment(local_assignment) => Stmt::LocalAssignment(
self.local_assignment_add_trivia(local_assignment, additional_indent_level),
),
Stmt::LocalFunction(local_function) => {
Stmt::LocalFunction(trivia_formatter::local_function_add_trivia(
local_function,
leading_trivia,
trailing_trivia,
))
}
Stmt::NumericFor(numeric_for) => {
Stmt::NumericFor(trivia_formatter::numeric_for_add_trivia(
numeric_for,
leading_trivia,
trailing_trivia,
))
}
Stmt::Repeat(repeat_block) => {
Stmt::Repeat(self.repeat_block_add_trivia(repeat_block, additional_indent_level))
}
Stmt::While(while_block) => {
Stmt::While(self.while_block_add_trivia(while_block, additional_indent_level))
}
#[cfg(feature = "luau")]
Stmt::CompoundAssignment(compound_assignment) => {
Stmt::CompoundAssignment(trivia_formatter::compound_assignment_add_trivia(
compound_assignment,
leading_trivia,
trailing_trivia,
))
}
#[cfg(feature = "luau")]
Stmt::ExportedTypeDeclaration(exported_type_declaration) => {
Stmt::ExportedTypeDeclaration(
trivia_formatter::exported_type_declaration_add_trivia(
exported_type_declaration,
leading_trivia,
trailing_trivia,
),
)
}
#[cfg(feature = "luau")]
Stmt::TypeDeclaration(type_declaration) => {
Stmt::TypeDeclaration(trivia_formatter::type_declaration_add_trivia(
type_declaration,
leading_trivia,
trailing_trivia,
))
}
}
}
}