use crate::{Config, IndentType, LineEndings};
use full_moon::ast::{
punctuated::{Pair, Punctuated},
span::ContainedSpan,
Block,
};
use full_moon::tokenizer::{
StringLiteralQuoteType, Symbol, Token, TokenKind, TokenReference, TokenType,
};
use full_moon::visitors::VisitorMut;
use std::borrow::Cow;
use std::collections::HashSet;
pub mod assignment_formatter;
pub mod block_formatter;
#[macro_use]
pub mod expression_formatter;
pub mod functions_formatter;
#[cfg(feature = "luau")]
pub mod luau_formatter;
pub mod stmt_formatter;
pub mod table_formatter;
pub mod trivia_formatter;
pub mod trivia_util;
pub type Range = (usize, usize);
#[derive(Default)]
pub struct CodeFormatter {
config: Config,
indent_level: usize,
indent_ranges: HashSet<Range>,
}
#[derive(Debug)]
enum FormatTokenType {
Token,
LeadingTrivia,
TrailingTrivia,
}
fn get_line_ending_character(line_endings: &LineEndings) -> String {
match line_endings {
LineEndings::Unix => String::from("\n"),
LineEndings::Windows => String::from("\r\n"),
}
}
#[macro_export]
macro_rules! fmt_symbol {
($fmter:expr, $token:expr, $x:expr) => {
$fmter.format_symbol($token, &TokenReference::symbol($x).unwrap())
};
}
impl CodeFormatter {
pub fn new(config: Config) -> Self {
CodeFormatter {
indent_level: 0,
config,
indent_ranges: HashSet::new(),
}
}
pub fn increment_indent_level(&mut self) {
self.indent_level += 1;
}
pub fn decrement_indent_level(&mut self) {
self.indent_level -= 1;
}
pub fn get_indent_width(&self) -> usize {
(self.indent_level - 1) * self.config.indent_width
}
pub fn add_indent_range(&mut self, range: Range) {
self.indent_ranges.insert(range);
}
pub fn get_range_indent_increase(&self, range: Range) -> Option<usize> {
let indent_increase = self
.indent_ranges
.iter()
.filter(|x| range.0 >= x.0 && range.1 <= x.1);
let count = indent_increase.count();
if count > 0 {
Some(count)
} else {
None
}
}
pub fn create_indent_trivia<'ast>(
&self,
additional_indent_level: Option<usize>,
) -> Token<'ast> {
let indent_level = match additional_indent_level {
Some(level) => self.indent_level - 1 + level,
None => self.indent_level - 1,
};
match self.config.indent_type {
IndentType::Tabs => Token::new(TokenType::tabs(indent_level)),
IndentType::Spaces => {
Token::new(TokenType::spaces(indent_level * self.config.indent_width))
}
}
}
pub fn create_newline_trivia<'ast>(&self) -> Token<'ast> {
Token::new(TokenType::Whitespace {
characters: Cow::Owned(get_line_ending_character(&self.config.line_endings)),
})
}
fn format_single_line_comment_string(&self, comment: String) -> String {
comment.trim_end().to_string()
}
fn format_token<'ast>(
&self,
token: Token<'ast>,
format_type: &FormatTokenType,
additional_indent_level: Option<usize>,
) -> (
Token<'ast>,
Option<Vec<Token<'ast>>>,
Option<Vec<Token<'ast>>>,
) {
let mut leading_trivia: Option<Vec<Token<'ast>>> = None;
let mut trailing_trivia: Option<Vec<Token<'ast>>> = None;
let token_type = match token.token_type() {
TokenType::Number { text } => TokenType::Number {
text: Cow::Owned(if text.starts_with('.') {
String::from("0")
+ match text {
Cow::Owned(text) => text.as_str(),
Cow::Borrowed(text) => text,
}
} else if text.starts_with("-.") {
String::from("-0") + text.get(1..).expect("unknown number literal")
} else {
text.to_owned().into_owned()
}),
},
TokenType::StringLiteral {
literal,
multi_line,
quote_type,
} => {
if let StringLiteralQuoteType::Brackets = quote_type {
TokenType::StringLiteral {
literal: literal.to_owned(),
multi_line: *multi_line,
quote_type: StringLiteralQuoteType::Brackets,
}
} else {
lazy_static::lazy_static! {
static ref RE: regex::Regex = regex::Regex::new(r#"\\([\S\s])|(["'])"#).unwrap();
static ref UNNECESSARY_ESCAPES: regex::Regex = regex::Regex::new(r#"^[^\n\r"'0-7\\bfnrt-vx\u2028\u2029]$"#).unwrap();
}
let literal = RE
.replace_all(literal, |caps: ®ex::Captures| {
let escaped = caps.get(1);
let quote = caps.get(2);
match quote {
Some(quote) => {
let quote_type = match quote.as_str() {
"'" => StringLiteralQuoteType::Single,
"\"" => StringLiteralQuoteType::Double,
_ => panic!("unknown quote type"),
};
if let StringLiteralQuoteType::Single = quote_type {
String::from("'")
} else {
String::from("\\\"")
}
}
None => {
let text = escaped
.expect(
"have a match which was neither an escape or a quote",
)
.as_str();
if UNNECESSARY_ESCAPES.is_match(text) {
text.to_owned()
} else {
format!("\\{}", text.to_owned())
}
}
}
})
.into_owned();
TokenType::StringLiteral {
literal: Cow::Owned(literal),
multi_line: None,
quote_type: StringLiteralQuoteType::Double,
}
}
}
TokenType::SingleLineComment { comment } => {
let comment =
self.format_single_line_comment_string(comment.to_owned().into_owned());
match format_type {
FormatTokenType::LeadingTrivia => {
let additional_indent_level = additional_indent_level.unwrap_or(0)
+ self
.get_range_indent_increase((
token.start_position().bytes(),
token.end_position().bytes(),
))
.unwrap_or(0);
leading_trivia = Some(vec![
self.create_indent_trivia(Some(additional_indent_level))
]);
trailing_trivia = Some(vec![self.create_newline_trivia()]);
}
FormatTokenType::TrailingTrivia => {
leading_trivia = Some(vec![Token::new(TokenType::spaces(1))]);
}
_ => (),
}
TokenType::SingleLineComment {
comment: Cow::Owned(comment),
}
}
TokenType::MultiLineComment { blocks, comment } => {
if let FormatTokenType::LeadingTrivia = format_type {
let additional_indent_level = additional_indent_level.unwrap_or(0)
+ self
.get_range_indent_increase((
token.start_position().bytes(),
token.end_position().bytes(),
))
.unwrap_or(0);
leading_trivia = Some(vec![
self.create_indent_trivia(Some(additional_indent_level))
]);
trailing_trivia = Some(vec![self.create_newline_trivia()]);
}
TokenType::MultiLineComment {
blocks: *blocks,
comment: comment.to_owned(),
}
}
TokenType::Whitespace { characters } => TokenType::Whitespace {
characters: characters.to_owned(),
},
_ => token.token_type().to_owned(),
};
(Token::new(token_type), leading_trivia, trailing_trivia)
}
fn load_token_trivia<'ast>(
&self,
current_trivia: Vec<&Token<'ast>>,
format_token_type: FormatTokenType,
additional_indent_level: Option<usize>,
) -> Vec<Token<'ast>> {
let mut token_trivia = Vec::new();
let mut newline_count_in_succession = 0;
for trivia in current_trivia {
match trivia.token_type() {
TokenType::Whitespace { characters } => {
if characters.contains('\n') {
newline_count_in_succession += 1;
if newline_count_in_succession == 2 {
token_trivia.push(Token::new(TokenType::Whitespace {
characters: Cow::Owned(get_line_ending_character(
&self.config.line_endings,
)),
}));
}
}
continue;
}
_ => {
newline_count_in_succession = 0;
}
}
let (token, leading_trivia, trailing_trivia) = self.format_token(
trivia.to_owned(),
&format_token_type,
additional_indent_level,
);
if let Some(mut trivia) = leading_trivia {
token_trivia.append(&mut trivia);
}
token_trivia.push(token);
if let Some(mut trivia) = trailing_trivia {
token_trivia.append(&mut trivia)
}
}
token_trivia
}
fn format_plain_token_reference<'a>(
&self,
token_reference: &TokenReference<'a>,
) -> TokenReference<'a> {
let formatted_leading_trivia: Vec<Token<'a>> = self.load_token_trivia(
token_reference.leading_trivia().collect(),
FormatTokenType::LeadingTrivia,
None,
);
let formatted_trailing_trivia: Vec<Token<'a>> = self.load_token_trivia(
token_reference.trailing_trivia().collect(),
FormatTokenType::TrailingTrivia,
None,
);
let (token, _leading_trivia, _trailing_trivia) = self.format_token(
token_reference.token().to_owned(),
&FormatTokenType::Token,
None,
);
TokenReference::new(formatted_leading_trivia, token, formatted_trailing_trivia)
}
pub fn format_token_reference<'a>(
&self,
token_reference: &TokenReference<'a>,
) -> Cow<'a, TokenReference<'a>> {
Cow::Owned(self.format_plain_token_reference(&token_reference))
}
pub fn format_token_reference_mut<'ast>(
&mut self,
token_reference: &Cow<'ast, TokenReference<'ast>>,
) -> Cow<'ast, TokenReference<'ast>> {
Cow::Owned(self.format_plain_token_reference(&token_reference))
}
pub fn format_punctuation<'ast>(
&self,
punctuation: &TokenReference<'ast>,
) -> (Cow<'ast, TokenReference<'ast>>, Vec<Token<'ast>>) {
let trailing_comments = punctuation
.trailing_trivia()
.filter(|token| {
token.token_kind() == TokenKind::SingleLineComment
|| token.token_kind() == TokenKind::MultiLineComment
})
.map(|x| {
vec![Token::new(TokenType::spaces(1)), x.to_owned()]
})
.flatten()
.collect();
(
Cow::Owned(TokenReference::new(
Vec::new(),
punctuation.token().to_owned(),
vec![Token::new(TokenType::spaces(1))],
)),
trailing_comments,
)
}
pub fn format_punctuated<'a, T>(
&mut self,
old: &Punctuated<'a, T>,
value_formatter: &dyn Fn(&mut Self, &T) -> T,
) -> (Punctuated<'a, T>, Vec<Token<'a>>) {
let mut formatted: Punctuated<T> = Punctuated::new();
let mut comments_buffer = Vec::new();
for pair in old.pairs() {
match pair {
Pair::Punctuated(value, punctuation) => {
let (formatted_punctuation, mut comments) =
self.format_punctuation(punctuation);
comments_buffer.append(&mut comments);
let formatted_value = value_formatter(self, value);
formatted.push(Pair::new(formatted_value, Some(formatted_punctuation)));
}
Pair::End(value) => {
let formatted_value = value_formatter(self, value);
formatted.push(Pair::new(formatted_value, None));
}
}
}
(formatted, comments_buffer)
}
pub fn format_contained_span<'ast>(
&self,
contained_span: &ContainedSpan<'ast>,
) -> ContainedSpan<'ast> {
let (start_token, end_token) = contained_span.tokens();
ContainedSpan::new(
Cow::Owned(self.format_plain_token_reference(start_token)),
Cow::Owned(self.format_plain_token_reference(end_token)),
)
}
pub fn format_symbol<'ast>(
&self,
current_symbol: &TokenReference<'ast>,
wanted_symbol: &TokenReference<'ast>,
) -> Cow<'ast, TokenReference<'ast>> {
let mut formatted_leading_trivia: Vec<Token<'ast>> = self.load_token_trivia(
current_symbol.leading_trivia().collect(),
FormatTokenType::LeadingTrivia,
None,
);
let mut formatted_trailing_trivia: Vec<Token<'ast>> = self.load_token_trivia(
current_symbol.trailing_trivia().collect(),
FormatTokenType::TrailingTrivia,
None,
);
let mut wanted_leading_trivia: Vec<Token<'ast>> = wanted_symbol
.leading_trivia()
.map(|x| x.to_owned())
.collect();
let mut wanted_trailing_trivia: Vec<Token<'ast>> = wanted_symbol
.trailing_trivia()
.map(|x| x.to_owned())
.collect();
wanted_trailing_trivia.append(&mut formatted_trailing_trivia);
formatted_leading_trivia.append(&mut wanted_leading_trivia);
Cow::Owned(TokenReference::new(
formatted_leading_trivia,
wanted_symbol.token().to_owned(),
wanted_trailing_trivia,
))
}
pub fn format_end_token<'ast>(
&self,
current_token: &TokenReference<'ast>,
) -> Cow<'ast, TokenReference<'ast>> {
let formatted_leading_trivia: Vec<Token<'ast>> = self.load_token_trivia(
current_token.leading_trivia().collect(),
crate::formatters::FormatTokenType::LeadingTrivia,
Some(1),
);
let formatted_trailing_trivia: Vec<Token<'ast>> = self.load_token_trivia(
current_token.trailing_trivia().collect(),
crate::formatters::FormatTokenType::TrailingTrivia,
None,
);
Cow::Owned(TokenReference::new(
formatted_leading_trivia,
Token::new(TokenType::Symbol {
symbol: Symbol::End,
}),
formatted_trailing_trivia,
))
}
}
fn pop_until_no_whitespace<'ast>(trivia: &mut Vec<Token<'ast>>) {
if let Some(t) = trivia.pop() {
match t.token_kind() {
TokenKind::Whitespace => pop_until_no_whitespace(trivia),
_ => trivia.push(t),
}
}
}
impl<'ast> VisitorMut<'ast> for CodeFormatter {
fn visit_block(&mut self, node: Block<'ast>) -> Block<'ast> {
self.increment_indent_level();
self.format_block(node)
}
fn visit_block_end(&mut self, node: Block<'ast>) -> Block<'ast> {
self.decrement_indent_level();
node
}
fn visit_eof(&mut self, node: TokenReference<'ast>) -> TokenReference<'ast> {
self.indent_level += 1;
let mut formatted_leading_trivia: Vec<Token<'ast>> = self.load_token_trivia(
node.leading_trivia().collect(),
FormatTokenType::LeadingTrivia,
None,
);
let only_whitespace = formatted_leading_trivia
.iter()
.all(|x| x.token_kind() == TokenKind::Whitespace);
if only_whitespace {
TokenReference::new(Vec::new(), Token::new(TokenType::Eof), Vec::new())
} else {
pop_until_no_whitespace(&mut formatted_leading_trivia);
formatted_leading_trivia.push(Token::new(TokenType::Whitespace {
characters: Cow::Owned(get_line_ending_character(&self.config.line_endings)),
}));
TokenReference::new(
formatted_leading_trivia,
Token::new(TokenType::Eof),
Vec::new(),
)
}
}
}