1 use crate::formatters::trivia::{FormatTriviaType, UpdateLeadingTrivia, UpdateTrailingTrivia};
2 #[cfg(feature = "luau")]
3 use full_moon::ast::span::ContainedSpan;
4 #[cfg(feature = "luau")]
5 use full_moon::ast::types::{IndexedTypeInfo, TypeDeclaration, TypeInfo};
6 use full_moon::ast::{Block, FunctionBody};
7 use full_moon::{
8     ast::{
9         BinOp, Call, Expression, Field, FunctionArgs, Index, LastStmt, Prefix, Stmt, Suffix,
10         TableConstructor, UnOp, Value, Var,
11     },
12     node::Node,
13     tokenizer::{Token, TokenKind, TokenReference, TokenType},
14 };
15 
trivia_is_whitespace(trivia: &Token) -> bool16 pub fn trivia_is_whitespace(trivia: &Token) -> bool {
17     matches!(trivia.token_kind(), TokenKind::Whitespace)
18 }
19 
trivia_is_comment(trivia: &Token) -> bool20 pub fn trivia_is_comment(trivia: &Token) -> bool {
21     matches!(
22         trivia.token_kind(),
23         TokenKind::SingleLineComment | TokenKind::MultiLineComment
24     )
25 }
26 
trivia_is_newline(trivia: &Token) -> bool27 pub fn trivia_is_newline(trivia: &Token) -> bool {
28     if let TokenType::Whitespace { characters } = trivia.token_type() {
29         if characters.find('\n').is_some() {
30             return true;
31         }
32     }
33     false
34 }
35 
trivia_contains_newline<'a>(trivia_vec: impl Iterator<Item = &'a Token>) -> bool36 pub fn trivia_contains_newline<'a>(trivia_vec: impl Iterator<Item = &'a Token>) -> bool {
37     for trivia in trivia_vec {
38         if trivia_is_newline(trivia) {
39             return true;
40         }
41     }
42     false
43 }
44 
can_hang_expression(expression: &Expression) -> bool45 pub fn can_hang_expression(expression: &Expression) -> bool {
46     match expression {
47         Expression::Parentheses { expression, .. } => can_hang_expression(expression),
48         Expression::UnaryOperator { expression, .. } => can_hang_expression(expression),
49         Expression::BinaryOperator { .. } => true, // If a binop is present, then we can hang the expression
50         Expression::Value { value, .. } => match &**value {
51             Value::ParenthesesExpression(expression) => can_hang_expression(expression),
52             Value::FunctionCall(function_call) => match function_call.prefix() {
53                 Prefix::Expression(expression) => can_hang_expression(expression),
54                 _ => false,
55             },
56             Value::Var(Var::Expression(expression)) => match expression.prefix() {
57                 Prefix::Expression(expression) => can_hang_expression(expression),
58                 _ => false,
59             },
60             _ => false,
61         },
62         other => panic!("unknown node {:?}", other),
63     }
64 }
65 
is_block_empty(block: &Block) -> bool66 pub fn is_block_empty(block: &Block) -> bool {
67     block.stmts().next().is_none() && block.last_stmt().is_none()
68 }
69 
is_function_empty(function_body: &FunctionBody) -> bool70 pub fn is_function_empty(function_body: &FunctionBody) -> bool {
71     is_block_empty(function_body.block())
72         && !function_body
73             .parameters_parentheses()
74             .tokens()
75             .1
76             .trailing_trivia()
77             .any(trivia_is_comment)
78         && !function_body
79             .end_token()
80             .leading_trivia()
81             .any(trivia_is_comment)
82 }
83 
84 // TODO: Can we clean this up? A lot of this code is repeated in trivia_formatter
function_args_trailing_trivia(function_args: &FunctionArgs) -> Vec<Token>85 fn function_args_trailing_trivia(function_args: &FunctionArgs) -> Vec<Token> {
86     match function_args {
87         FunctionArgs::Parentheses { parentheses, .. } => {
88             let (_, end_brace) = parentheses.tokens();
89             end_brace.trailing_trivia().map(|x| x.to_owned()).collect()
90         }
91         FunctionArgs::String(token_reference) => token_reference
92             .trailing_trivia()
93             .map(|x| x.to_owned())
94             .collect(),
95         FunctionArgs::TableConstructor(table_constructor) => {
96             let (_, end_brace) = table_constructor.braces().tokens();
97             end_brace.trailing_trivia().map(|x| x.to_owned()).collect()
98         }
99         other => panic!("unknown node {:?}", other),
100     }
101 }
102 
suffix_trailing_trivia(suffix: &Suffix) -> Vec<Token>103 fn suffix_trailing_trivia(suffix: &Suffix) -> Vec<Token> {
104     match suffix {
105         Suffix::Index(index) => match index {
106             Index::Brackets { brackets, .. } => {
107                 let (_, end_brace) = brackets.tokens();
108                 end_brace.trailing_trivia().map(|x| x.to_owned()).collect()
109             }
110             Index::Dot { name, .. } => name.trailing_trivia().map(|x| x.to_owned()).collect(),
111             other => panic!("unknown node {:?}", other),
112         },
113         Suffix::Call(call) => match call {
114             Call::AnonymousCall(function_args) => function_args_trailing_trivia(function_args),
115             Call::MethodCall(method_call) => function_args_trailing_trivia(method_call.args()),
116             other => panic!("unknown node {:?}", other),
117         },
118         other => panic!("unknown node {:?}", other),
119     }
120 }
121 
122 #[cfg(feature = "luau")]
indexed_type_info_trailing_trivia(indexed_type_info: &IndexedTypeInfo) -> Vec<Token>123 fn indexed_type_info_trailing_trivia(indexed_type_info: &IndexedTypeInfo) -> Vec<Token> {
124     match indexed_type_info {
125         IndexedTypeInfo::Basic(token_reference) => token_reference
126             .trailing_trivia()
127             .map(|x| x.to_owned())
128             .collect(),
129         IndexedTypeInfo::Generic { arrows, .. } => {
130             let (_, end_brace) = arrows.tokens();
131             end_brace.trailing_trivia().map(|x| x.to_owned()).collect()
132         }
133         other => panic!("unknown node {:?}", other),
134     }
135 }
136 
137 #[cfg(feature = "luau")]
type_info_trailing_trivia(type_info: &TypeInfo) -> Vec<Token>138 pub fn type_info_trailing_trivia(type_info: &TypeInfo) -> Vec<Token> {
139     match type_info {
140         TypeInfo::Array { braces, .. } => {
141             let (_, end_brace) = braces.tokens();
142             end_brace.trailing_trivia().map(|x| x.to_owned()).collect()
143         }
144         TypeInfo::Basic(token_reference) => token_reference
145             .trailing_trivia()
146             .map(|x| x.to_owned())
147             .collect(),
148         TypeInfo::Callback { return_type, .. } => type_info_trailing_trivia(return_type),
149         TypeInfo::Generic { arrows, .. } => {
150             let (_, end_brace) = arrows.tokens();
151             end_brace.trailing_trivia().map(|x| x.to_owned()).collect()
152         }
153 
154         TypeInfo::Intersection { right, .. } => type_info_trailing_trivia(right),
155 
156         TypeInfo::Module { type_info, .. } => indexed_type_info_trailing_trivia(type_info),
157 
158         TypeInfo::Optional { question_mark, .. } => question_mark
159             .trailing_trivia()
160             .map(|x| x.to_owned())
161             .collect(),
162 
163         TypeInfo::Table { braces, .. } => {
164             let (_, end_brace) = braces.tokens();
165             end_brace.trailing_trivia().map(|x| x.to_owned()).collect()
166         }
167 
168         TypeInfo::Typeof { parentheses, .. } => {
169             let (_, end_brace) = parentheses.tokens();
170             end_brace.trailing_trivia().map(|x| x.to_owned()).collect()
171         }
172 
173         TypeInfo::Tuple { parentheses, .. } => {
174             let (_, end_brace) = parentheses.tokens();
175             end_brace.trailing_trivia().map(|x| x.to_owned()).collect()
176         }
177 
178         TypeInfo::Union { right, .. } => type_info_trailing_trivia(right),
179         TypeInfo::Variadic { type_info, .. } => type_info_trailing_trivia(type_info),
180 
181         other => panic!("unknown node {:?}", other),
182     }
183 }
184 
var_trailing_trivia(var: &Var) -> Vec<Token>185 fn var_trailing_trivia(var: &Var) -> Vec<Token> {
186     match var {
187         Var::Name(token_reference) => token_reference
188             .trailing_trivia()
189             .map(|x| x.to_owned())
190             .collect(),
191         Var::Expression(var_expr) => {
192             if let Some(last_suffix) = var_expr.suffixes().last() {
193                 suffix_trailing_trivia(last_suffix)
194             } else {
195                 // TODO: is it possible for this to happen?
196                 vec![]
197             }
198         }
199         other => panic!("unknown node {:?}", other),
200     }
201 }
202 
get_value_trailing_trivia(value: &Value) -> Vec<Token>203 pub fn get_value_trailing_trivia(value: &Value) -> Vec<Token> {
204     match value {
205         Value::Function((_, function_body)) => function_body
206             .end_token()
207             .trailing_trivia()
208             .map(|x| x.to_owned())
209             .collect(),
210         Value::FunctionCall(function_call) => {
211             if let Some(last_suffix) = function_call.suffixes().last() {
212                 suffix_trailing_trivia(last_suffix)
213             } else {
214                 // TODO: is it possible for this to happen?
215                 vec![]
216             }
217         }
218         Value::String(token_reference) => token_reference
219             .trailing_trivia()
220             .map(|x| x.to_owned())
221             .collect(),
222         Value::TableConstructor(table_constructor) => {
223             let (_, end_brace) = table_constructor.braces().tokens();
224             end_brace.trailing_trivia().map(|x| x.to_owned()).collect()
225         }
226         Value::Number(token_reference) => token_reference
227             .trailing_trivia()
228             .map(|x| x.to_owned())
229             .collect(),
230         Value::ParenthesesExpression(expr) => get_expression_trailing_trivia(expr),
231         Value::Symbol(token_reference) => token_reference
232             .trailing_trivia()
233             .map(|x| x.to_owned())
234             .collect(),
235         Value::Var(var) => var_trailing_trivia(var),
236         other => panic!("unknown node {:?}", other),
237     }
238 }
239 
get_expression_trailing_trivia(expression: &Expression) -> Vec<Token>240 pub fn get_expression_trailing_trivia(expression: &Expression) -> Vec<Token> {
241     match expression {
242         Expression::Parentheses { contained, .. } => {
243             let (_, end_parentheses) = contained.tokens();
244             end_parentheses
245                 .trailing_trivia()
246                 .map(|x| x.to_owned())
247                 .collect()
248         }
249         Expression::UnaryOperator { expression, .. } => get_expression_trailing_trivia(expression),
250         Expression::BinaryOperator { rhs, .. } => get_expression_trailing_trivia(rhs),
251         Expression::Value {
252             value,
253             #[cfg(feature = "luau")]
254             type_assertion,
255         } => {
256             #[cfg(feature = "luau")]
257             if let Some(type_assertion) = type_assertion {
258                 return type_info_trailing_trivia(type_assertion.cast_to());
259             }
260 
261             get_value_trailing_trivia(value)
262         }
263         other => panic!("unknown node {:?}", other),
264     }
265 }
266 
get_expression_leading_trivia(expression: &Expression) -> Vec<Token>267 pub fn get_expression_leading_trivia(expression: &Expression) -> Vec<Token> {
268     match expression {
269         Expression::Parentheses { contained, .. } => contained
270             .tokens()
271             .0
272             .leading_trivia()
273             .map(|x| x.to_owned())
274             .collect(),
275         Expression::UnaryOperator { unop, .. } => match unop {
276             UnOp::Minus(token_ref) | UnOp::Not(token_ref) | UnOp::Hash(token_ref) => {
277                 token_ref.leading_trivia().map(|x| x.to_owned()).collect()
278             }
279             other => panic!("unknown node {:?}", other),
280         },
281         Expression::BinaryOperator { lhs, .. } => get_expression_leading_trivia(lhs),
282         Expression::Value { value, .. } => match &**value {
283             Value::Function((token_ref, _)) => {
284                 token_ref.leading_trivia().map(|x| x.to_owned()).collect()
285             }
286             Value::FunctionCall(function_call) => match function_call.prefix() {
287                 Prefix::Name(token_ref) => {
288                     token_ref.leading_trivia().map(|x| x.to_owned()).collect()
289                 }
290                 Prefix::Expression(expr) => get_expression_leading_trivia(expr),
291                 other => panic!("unknown node {:?}", other),
292             },
293             Value::TableConstructor(table) => table
294                 .braces()
295                 .tokens()
296                 .0
297                 .leading_trivia()
298                 .map(|x| x.to_owned())
299                 .collect(),
300             Value::Number(token_ref) => token_ref.leading_trivia().map(|x| x.to_owned()).collect(),
301             Value::ParenthesesExpression(expr) => get_expression_leading_trivia(expr),
302             Value::String(token_ref) => token_ref.leading_trivia().map(|x| x.to_owned()).collect(),
303             Value::Symbol(token_ref) => token_ref.leading_trivia().map(|x| x.to_owned()).collect(),
304             Value::Var(var) => match var {
305                 Var::Name(token_ref) => token_ref.leading_trivia().map(|x| x.to_owned()).collect(),
306                 Var::Expression(var_expr) => match var_expr.prefix() {
307                     Prefix::Name(token_ref) => {
308                         token_ref.leading_trivia().map(|x| x.to_owned()).collect()
309                     }
310                     Prefix::Expression(expr) => get_expression_leading_trivia(expr),
311                     other => panic!("unknown node {:?}", other),
312                 },
313                 other => panic!("unknown node {:?}", other),
314             },
315             other => panic!("unknown node {:?}", other),
316         },
317         other => panic!("unknown node {:?}", other),
318     }
319 }
320 
binop_leading_comments(binop: &BinOp) -> Vec<Token>321 pub fn binop_leading_comments(binop: &BinOp) -> Vec<Token> {
322     match binop {
323         BinOp::And(token)
324         | BinOp::Caret(token)
325         | BinOp::GreaterThan(token)
326         | BinOp::GreaterThanEqual(token)
327         | BinOp::LessThan(token)
328         | BinOp::LessThanEqual(token)
329         | BinOp::Minus(token)
330         | BinOp::Or(token)
331         | BinOp::Percent(token)
332         | BinOp::Plus(token)
333         | BinOp::Slash(token)
334         | BinOp::Star(token)
335         | BinOp::TildeEqual(token)
336         | BinOp::TwoDots(token)
337         | BinOp::TwoEqual(token) => token
338             .leading_trivia()
339             .filter(|token| trivia_is_comment(token))
340             .map(|x| x.to_owned())
341             .collect(),
342         other => panic!("unknown node {:?}", other),
343     }
344 }
345 
binop_trailing_comments(binop: &BinOp) -> Vec<Token>346 pub fn binop_trailing_comments(binop: &BinOp) -> Vec<Token> {
347     match binop {
348         BinOp::And(token)
349         | BinOp::Caret(token)
350         | BinOp::GreaterThan(token)
351         | BinOp::GreaterThanEqual(token)
352         | BinOp::LessThan(token)
353         | BinOp::LessThanEqual(token)
354         | BinOp::Minus(token)
355         | BinOp::Or(token)
356         | BinOp::Percent(token)
357         | BinOp::Plus(token)
358         | BinOp::Slash(token)
359         | BinOp::Star(token)
360         | BinOp::TildeEqual(token)
361         | BinOp::TwoDots(token)
362         | BinOp::TwoEqual(token) => {
363             token
364                 .trailing_trivia()
365                 .filter(|token| trivia_is_comment(token))
366                 .flat_map(|x| {
367                     // Prepend a single space beforehand
368                     vec![Token::new(TokenType::spaces(1)), x.to_owned()]
369                 })
370                 .collect()
371         }
372         other => panic!("unknown node {:?}", other),
373     }
374 }
375 
expression_leading_comments(expression: &Expression) -> Vec<Token>376 pub fn expression_leading_comments(expression: &Expression) -> Vec<Token> {
377     get_expression_leading_trivia(expression)
378         .iter()
379         .filter(|token| trivia_is_comment(token))
380         .map(|x| x.to_owned())
381         .collect()
382 }
383 
take_expression_leading_comments(expression: &Expression) -> (Expression, Vec<Token>)384 pub fn take_expression_leading_comments(expression: &Expression) -> (Expression, Vec<Token>) {
385     let trailing_comments = get_expression_leading_trivia(expression)
386         .iter()
387         .filter(|token| trivia_is_comment(token))
388         .map(|x| x.to_owned())
389         .collect();
390 
391     (
392         expression.update_leading_trivia(
393             FormatTriviaType::Replace(vec![]), // TODO: Do we need to keep some trivia?
394         ),
395         trailing_comments,
396     )
397 }
398 
take_expression_trailing_comments(expression: &Expression) -> (Expression, Vec<Token>)399 pub fn take_expression_trailing_comments(expression: &Expression) -> (Expression, Vec<Token>) {
400     let trailing_comments = get_expression_trailing_trivia(expression)
401         .iter()
402         .filter(|token| trivia_is_comment(token))
403         .map(|x| {
404             // Prepend a single space beforehand
405             vec![Token::new(TokenType::spaces(1)), x.to_owned()]
406         })
407         .flatten()
408         .collect();
409 
410     (
411         expression.update_trailing_trivia(
412             FormatTriviaType::Replace(vec![]), // TODO: Do we need to keep some trivia?
413         ),
414         trailing_comments,
415     )
416 }
417 
418 /// Macro for retrieving trailing trivia out of a stmt which ends with an `end` token
419 macro_rules! end_stmt_trailing_trivia {
420     ($enum:ident, $value:ident) => {{
421         let end_token = $value.end_token();
422         let trailing_trivia = end_token.trailing_trivia().map(|x| x.to_owned()).collect();
423         let new_end_token = end_token.update_trailing_trivia(FormatTriviaType::Replace(vec![]));
424 
425         (
426             Stmt::$enum($value.with_end_token(new_end_token)),
427             trailing_trivia,
428         )
429     }};
430 }
431 
432 #[cfg(feature = "luau")]
get_indexed_type_info_trailing_trivia( type_info: IndexedTypeInfo, ) -> (IndexedTypeInfo, Vec<Token>)433 fn get_indexed_type_info_trailing_trivia(
434     type_info: IndexedTypeInfo,
435 ) -> (IndexedTypeInfo, Vec<Token>) {
436     match type_info {
437         IndexedTypeInfo::Basic(token) => {
438             let trailing_trivia = token.trailing_trivia().map(|x| x.to_owned()).collect();
439             let token = token.update_trailing_trivia(FormatTriviaType::Replace(vec![]));
440             (IndexedTypeInfo::Basic(token), trailing_trivia)
441         }
442         IndexedTypeInfo::Generic {
443             base,
444             arrows,
445             generics,
446         } => {
447             let (start_brace, end_brace) = arrows.tokens();
448             let trailing_trivia = end_brace.trailing_trivia().map(|x| x.to_owned()).collect();
449             let braces = ContainedSpan::new(
450                 start_brace.to_owned(),
451                 end_brace.update_trailing_trivia(FormatTriviaType::Replace(vec![])),
452             );
453 
454             (
455                 IndexedTypeInfo::Generic {
456                     base,
457                     arrows: braces,
458                     generics,
459                 },
460                 trailing_trivia,
461             )
462         }
463         other => panic!("unknown node {:?}", other),
464     }
465 }
466 
467 #[cfg(feature = "luau")]
get_type_info_trailing_trivia(type_info: TypeInfo) -> (TypeInfo, Vec<Token>)468 fn get_type_info_trailing_trivia(type_info: TypeInfo) -> (TypeInfo, Vec<Token>) {
469     match type_info {
470         TypeInfo::Array { braces, type_info } => {
471             let (start_brace, end_brace) = braces.tokens();
472             let trailing_trivia = end_brace.trailing_trivia().map(|x| x.to_owned()).collect();
473             let braces = ContainedSpan::new(
474                 start_brace.to_owned(),
475                 end_brace.update_trailing_trivia(FormatTriviaType::Replace(vec![])),
476             );
477 
478             (TypeInfo::Array { braces, type_info }, trailing_trivia)
479         }
480         TypeInfo::Basic(token) => {
481             let trailing_trivia = token.trailing_trivia().map(|x| x.to_owned()).collect();
482             let token = token.update_trailing_trivia(FormatTriviaType::Replace(vec![]));
483             (TypeInfo::Basic(token), trailing_trivia)
484         }
485         TypeInfo::Callback {
486             parentheses,
487             arguments,
488             arrow,
489             return_type,
490         } => {
491             let (return_type, trailing_trivia) = get_type_info_trailing_trivia(*return_type);
492             (
493                 TypeInfo::Callback {
494                     parentheses,
495                     arguments,
496                     arrow,
497                     return_type: Box::new(return_type),
498                 },
499                 trailing_trivia,
500             )
501         }
502         TypeInfo::Generic {
503             base,
504             arrows,
505             generics,
506         } => {
507             let (start_brace, end_brace) = arrows.tokens();
508             let trailing_trivia = end_brace.trailing_trivia().map(|x| x.to_owned()).collect();
509             let braces = ContainedSpan::new(
510                 start_brace.to_owned(),
511                 end_brace.update_trailing_trivia(FormatTriviaType::Replace(vec![])),
512             );
513 
514             (
515                 TypeInfo::Generic {
516                     base,
517                     arrows: braces,
518                     generics,
519                 },
520                 trailing_trivia,
521             )
522         }
523         TypeInfo::Intersection {
524             left,
525             ampersand,
526             right,
527         } => {
528             let (right, trailing_trivia) = get_type_info_trailing_trivia(*right);
529             (
530                 TypeInfo::Intersection {
531                     left,
532                     ampersand,
533                     right: Box::new(right),
534                 },
535                 trailing_trivia,
536             )
537         }
538         TypeInfo::Module {
539             module,
540             punctuation,
541             type_info,
542         } => {
543             let (type_info, trailing_trivia) = get_indexed_type_info_trailing_trivia(*type_info);
544             (
545                 TypeInfo::Module {
546                     module,
547                     punctuation,
548                     type_info: Box::new(type_info),
549                 },
550                 trailing_trivia,
551             )
552         }
553         TypeInfo::Optional {
554             base,
555             question_mark,
556         } => {
557             let trailing_trivia = question_mark
558                 .trailing_trivia()
559                 .map(|x| x.to_owned())
560                 .collect();
561             let question_mark =
562                 question_mark.update_trailing_trivia(FormatTriviaType::Replace(vec![]));
563             (
564                 TypeInfo::Optional {
565                     base,
566                     question_mark,
567                 },
568                 trailing_trivia,
569             )
570         }
571         TypeInfo::Table { braces, fields } => {
572             let (start_brace, end_brace) = braces.tokens();
573             let trailing_trivia = end_brace.trailing_trivia().map(|x| x.to_owned()).collect();
574             let braces = ContainedSpan::new(
575                 start_brace.to_owned(),
576                 end_brace.update_trailing_trivia(FormatTriviaType::Replace(vec![])),
577             );
578 
579             (TypeInfo::Table { braces, fields }, trailing_trivia)
580         }
581         TypeInfo::Typeof {
582             typeof_token,
583             parentheses,
584             inner,
585         } => {
586             let (start_brace, end_brace) = parentheses.tokens();
587             let trailing_trivia = end_brace.trailing_trivia().map(|x| x.to_owned()).collect();
588             let braces = ContainedSpan::new(
589                 start_brace.to_owned(),
590                 end_brace.update_trailing_trivia(FormatTriviaType::Replace(vec![])),
591             );
592 
593             (
594                 TypeInfo::Typeof {
595                     typeof_token,
596                     parentheses: braces,
597                     inner,
598                 },
599                 trailing_trivia,
600             )
601         }
602         TypeInfo::Tuple { parentheses, types } => {
603             let (start_brace, end_brace) = parentheses.tokens();
604             let trailing_trivia = end_brace.trailing_trivia().map(|x| x.to_owned()).collect();
605             let braces = ContainedSpan::new(
606                 start_brace.to_owned(),
607                 end_brace.update_trailing_trivia(FormatTriviaType::Replace(vec![])),
608             );
609 
610             (
611                 TypeInfo::Tuple {
612                     types,
613                     parentheses: braces,
614                 },
615                 trailing_trivia,
616             )
617         }
618         TypeInfo::Union { left, pipe, right } => {
619             let (right, trailing_trivia) = get_type_info_trailing_trivia(*right);
620             (
621                 TypeInfo::Union {
622                     left,
623                     pipe,
624                     right: Box::new(right),
625                 },
626                 trailing_trivia,
627             )
628         }
629         TypeInfo::Variadic { ellipse, type_info } => {
630             let (type_info, trailing_trivia) = get_type_info_trailing_trivia(*type_info);
631             (
632                 TypeInfo::Variadic {
633                     ellipse,
634                     type_info: Box::new(type_info),
635                 },
636                 trailing_trivia,
637             )
638         }
639         other => panic!("unknown node {:?}", other),
640     }
641 }
642 
643 #[cfg(feature = "luau")]
get_type_declaration_trailing_trivia( type_declaration: TypeDeclaration, ) -> (TypeDeclaration, Vec<Token>)644 fn get_type_declaration_trailing_trivia(
645     type_declaration: TypeDeclaration,
646 ) -> (TypeDeclaration, Vec<Token>) {
647     let (type_definition, trailing_trivia) =
648         get_type_info_trailing_trivia(type_declaration.type_definition().to_owned());
649     (
650         type_declaration.with_type_definition(type_definition),
651         trailing_trivia,
652     )
653 }
654 
get_stmt_trailing_trivia(stmt: Stmt) -> (Stmt, Vec<Token>)655 pub fn get_stmt_trailing_trivia(stmt: Stmt) -> (Stmt, Vec<Token>) {
656     let (updated_stmt, trailing_trivia) = match stmt {
657         Stmt::Assignment(assignment) => {
658             let mut formatted_expression_list = assignment.expressions().to_owned();
659             let mut trailing_trivia = Vec::new();
660             if let Some(last_pair) = formatted_expression_list.pop() {
661                 let pair = last_pair.map(|value| {
662                     trailing_trivia = get_expression_trailing_trivia(&value);
663                     value.update_trailing_trivia(FormatTriviaType::Replace(vec![]))
664                 });
665                 formatted_expression_list.push(pair);
666             }
667 
668             (
669                 Stmt::Assignment(assignment.with_expressions(formatted_expression_list)),
670                 trailing_trivia,
671             )
672         }
673 
674         Stmt::LocalAssignment(local_assignment) => {
675             let mut trailing_trivia = Vec::new();
676             let new_assignment = if local_assignment.expressions().is_empty() {
677                 // Unassigned local variable
678                 let mut formatted_name_list = local_assignment.names().to_owned();
679                 // Retrieve last item and take its trailing comments
680                 if let Some(last_pair) = formatted_name_list.pop() {
681                     let pair = last_pair.map(|value| {
682                         trailing_trivia = value.trailing_trivia().map(|x| x.to_owned()).collect();
683                         value.update_trailing_trivia(FormatTriviaType::Replace(vec![]))
684                     });
685                     formatted_name_list.push(pair);
686                 }
687                 local_assignment.with_names(formatted_name_list)
688             } else {
689                 // Add newline at the end of LocalAssignment expression list
690                 // Expression list should already be formatted
691                 let mut formatted_expression_list = local_assignment.expressions().to_owned();
692 
693                 // Retrieve last item and remove trailing trivia
694                 if let Some(last_pair) = formatted_expression_list.pop() {
695                     let pair = last_pair.map(|value| {
696                         trailing_trivia = get_expression_trailing_trivia(&value);
697                         value.update_trailing_trivia(FormatTriviaType::Replace(vec![]))
698                     });
699                     formatted_expression_list.push(pair);
700                 }
701 
702                 local_assignment.with_expressions(formatted_expression_list)
703             };
704 
705             (Stmt::LocalAssignment(new_assignment), trailing_trivia)
706         }
707 
708         Stmt::FunctionCall(function_call) => {
709             let last_suffix = function_call.suffixes().last();
710             let trailing_trivia = match last_suffix {
711                 Some(suffix) => suffix_trailing_trivia(suffix),
712                 None => Vec::new(),
713             };
714 
715             (
716                 Stmt::FunctionCall(
717                     function_call.update_trailing_trivia(FormatTriviaType::Replace(vec![])),
718                 ),
719                 trailing_trivia,
720             )
721         }
722         Stmt::Repeat(repeat_block) => {
723             let trailing_trivia = get_expression_trailing_trivia(repeat_block.until());
724             let until_expr = repeat_block
725                 .until()
726                 .update_trailing_trivia(FormatTriviaType::Replace(vec![]));
727 
728             (
729                 Stmt::Repeat(repeat_block.with_until(until_expr)),
730                 trailing_trivia,
731             )
732         }
733 
734         Stmt::Do(stmt) => {
735             end_stmt_trailing_trivia!(Do, stmt)
736         }
737         Stmt::GenericFor(stmt) => {
738             end_stmt_trailing_trivia!(GenericFor, stmt)
739         }
740         Stmt::If(stmt) => {
741             end_stmt_trailing_trivia!(If, stmt)
742         }
743         Stmt::FunctionDeclaration(stmt) => {
744             let end_token = stmt.body().end_token();
745             let trailing_trivia = end_token.trailing_trivia().map(|x| x.to_owned()).collect();
746             let new_end_token = end_token.update_trailing_trivia(FormatTriviaType::Replace(vec![]));
747 
748             let body = stmt.body().to_owned().with_end_token(new_end_token);
749             (
750                 Stmt::FunctionDeclaration(stmt.with_body(body)),
751                 trailing_trivia,
752             )
753         }
754         Stmt::LocalFunction(stmt) => {
755             let end_token = stmt.body().end_token();
756             let trailing_trivia = end_token.trailing_trivia().map(|x| x.to_owned()).collect();
757             let new_end_token = end_token.update_trailing_trivia(FormatTriviaType::Replace(vec![]));
758 
759             let body = stmt.body().to_owned().with_end_token(new_end_token);
760             (Stmt::LocalFunction(stmt.with_body(body)), trailing_trivia)
761         }
762         Stmt::NumericFor(stmt) => {
763             end_stmt_trailing_trivia!(NumericFor, stmt)
764         }
765         Stmt::While(stmt) => {
766             end_stmt_trailing_trivia!(While, stmt)
767         }
768 
769         #[cfg(feature = "luau")]
770         Stmt::CompoundAssignment(stmt) => {
771             let trailing_trivia = get_expression_trailing_trivia(stmt.rhs());
772             let expr = stmt
773                 .rhs()
774                 .update_trailing_trivia(FormatTriviaType::Replace(vec![]));
775             (
776                 Stmt::CompoundAssignment(stmt.with_rhs(expr)),
777                 trailing_trivia,
778             )
779         }
780         #[cfg(feature = "luau")]
781         Stmt::ExportedTypeDeclaration(stmt) => {
782             let (type_declaration, trailing_trivia) =
783                 get_type_declaration_trailing_trivia(stmt.type_declaration().to_owned());
784             (
785                 Stmt::ExportedTypeDeclaration(stmt.with_type_declaration(type_declaration)),
786                 trailing_trivia,
787             )
788         }
789         #[cfg(feature = "luau")]
790         Stmt::TypeDeclaration(stmt) => {
791             let (type_declaration, trailing_trivia) = get_type_declaration_trailing_trivia(stmt);
792             (Stmt::TypeDeclaration(type_declaration), trailing_trivia)
793         }
794         #[cfg(feature = "lua52")]
795         Stmt::Goto(stmt) => {
796             let trailing_trivia = stmt
797                 .label_name()
798                 .trailing_trivia()
799                 .map(|x| x.to_owned())
800                 .collect();
801             let label_name = stmt
802                 .label_name()
803                 .update_trailing_trivia(FormatTriviaType::Replace(vec![]));
804             (
805                 Stmt::Goto(stmt.with_label_name(label_name)),
806                 trailing_trivia,
807             )
808         }
809         #[cfg(feature = "lua52")]
810         Stmt::Label(stmt) => {
811             let trailing_trivia = stmt
812                 .right_colons()
813                 .trailing_trivia()
814                 .map(|x| x.to_owned())
815                 .collect();
816             let right_colons = stmt
817                 .right_colons()
818                 .update_trailing_trivia(FormatTriviaType::Replace(vec![]));
819             (
820                 Stmt::Label(stmt.with_right_colons(right_colons)),
821                 trailing_trivia,
822             )
823         }
824 
825         other => panic!("unknown node {:?}", other),
826     };
827 
828     (updated_stmt, trailing_trivia)
829 }
830 
get_last_stmt_trailing_trivia(last_stmt: LastStmt) -> (LastStmt, Vec<Token>)831 pub fn get_last_stmt_trailing_trivia(last_stmt: LastStmt) -> (LastStmt, Vec<Token>) {
832     match last_stmt {
833         LastStmt::Return(ret) => {
834             let mut return_token = ret.token().to_owned();
835             let mut formatted_expression_list = ret.returns().to_owned();
836             let mut trailing_trivia = Vec::new();
837 
838             // Retrieve last item and remove trailing trivia
839             if let Some(last_pair) = formatted_expression_list.pop() {
840                 let pair = last_pair.map(|value| {
841                     trailing_trivia = get_expression_trailing_trivia(&value);
842                     value.update_trailing_trivia(FormatTriviaType::Replace(vec![]))
843                 });
844                 formatted_expression_list.push(pair);
845             } else {
846                 trailing_trivia = return_token
847                     .trailing_trivia()
848                     .map(|x| x.to_owned())
849                     .collect();
850                 return_token =
851                     return_token.update_trailing_trivia(FormatTriviaType::Replace(vec![]));
852             }
853 
854             (
855                 LastStmt::Return(
856                     ret.with_token(return_token)
857                         .with_returns(formatted_expression_list),
858                 ),
859                 trailing_trivia,
860             )
861         }
862         LastStmt::Break(token) => {
863             let trailing_trivia = token.trailing_trivia().map(|x| x.to_owned()).collect();
864             let token = token.update_trailing_trivia(FormatTriviaType::Replace(vec![]));
865 
866             (LastStmt::Break(token), trailing_trivia)
867         }
868         #[cfg(feature = "luau")]
869         LastStmt::Continue(token) => {
870             let trailing_trivia = token.trailing_trivia().map(|x| x.to_owned()).collect();
871             let token = token.update_trailing_trivia(FormatTriviaType::Replace(vec![]));
872 
873             (LastStmt::Continue(token), trailing_trivia)
874         }
875         other => panic!("unknown node {:?}", other),
876     }
877 }
878 
token_trivia_contains_comments<'a>(trivia: impl Iterator<Item = &'a Token>) -> bool879 pub fn token_trivia_contains_comments<'a>(trivia: impl Iterator<Item = &'a Token>) -> bool {
880     for trivia in trivia {
881         if trivia_is_comment(trivia) {
882             return true;
883         }
884     }
885     false
886 }
887 
token_contains_leading_comments(token_ref: &TokenReference) -> bool888 pub fn token_contains_leading_comments(token_ref: &TokenReference) -> bool {
889     token_trivia_contains_comments(token_ref.leading_trivia())
890 }
891 
token_contains_trailing_comments(token_ref: &TokenReference) -> bool892 pub fn token_contains_trailing_comments(token_ref: &TokenReference) -> bool {
893     token_trivia_contains_comments(token_ref.trailing_trivia())
894 }
895 
token_contains_comments(token_ref: &TokenReference) -> bool896 pub fn token_contains_comments(token_ref: &TokenReference) -> bool {
897     token_trivia_contains_comments(token_ref.leading_trivia())
898         || token_trivia_contains_comments(token_ref.trailing_trivia())
899 }
900 
contains_comments(node: impl Node) -> bool901 pub fn contains_comments(node: impl Node) -> bool {
902     node.tokens().into_iter().any(token_contains_comments)
903 }
904 
905 /// Checks whether any [`Field`] within a [`TableConstructor`] contains comments, without checking the braces
table_fields_contains_comments(table_constructor: &TableConstructor) -> bool906 pub fn table_fields_contains_comments(table_constructor: &TableConstructor) -> bool {
907     table_constructor.fields().pairs().any(|field| {
908         let comments = match field.value() {
909             Field::ExpressionKey {
910                 brackets,
911                 key,
912                 equal,
913                 value,
914             } => {
915                 contains_comments(brackets)
916                     || contains_comments(key)
917                     || contains_comments(equal)
918                     || contains_comments(value)
919             }
920             Field::NameKey { key, equal, value } => {
921                 contains_comments(key) || contains_comments(equal) || contains_comments(value)
922             }
923             Field::NoKey(expression) => contains_comments(expression),
924             other => panic!("unknown node {:?}", other),
925         };
926 
927         comments || field.punctuation().map_or(false, contains_comments)
928     })
929 }
930 
931 // Checks to see whether an expression contains comments inline inside of it
932 // This can only happen if the expression is a BinOp
933 // We should ignore any comments which are trailing for the whole expression, as they are not inline
expression_contains_inline_comments(expression: &Expression) -> bool934 pub fn expression_contains_inline_comments(expression: &Expression) -> bool {
935     match expression {
936         Expression::BinaryOperator { lhs, binop, rhs } => {
937             contains_comments(binop) || contains_comments(lhs)
938             // Check if the binop chain still continues
939             // If so, we should keep checking the expresion
940             // Otherwise, stop checking
941             || match &**rhs {
942                 Expression::BinaryOperator { .. } => expression_contains_inline_comments(rhs),
943                 Expression::UnaryOperator { unop, expression } => {
944                     let op_contains_comments = match unop {
945                         UnOp::Minus(token) | UnOp::Not(token) | UnOp::Hash(token) => contains_comments(token),
946                         other => panic!("unknown node {:?}", other)
947                     };
948                     op_contains_comments || expression_contains_inline_comments(expression)
949                 }
950                 Expression::Value{ .. } => false,
951                 Expression::Parentheses { .. } => contains_comments(rhs),
952                 other => panic!("unknown node {:?}", other),
953             }
954         }
955         _ => false,
956     }
957 }
958