Fix comments ordering when semicolon removed after stmt (#424)

* Extract out stmt trivia formatters

* Add test case

* Fix comments ordering of semicolon

* Update snapshot

* Update changelog

CHANGELOG.md

@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed issue through static linking where Windows binary would not execute due to missing `VCRUNTIME140.dll`. ([#413](https://github.com/JohnnyMorganz/StyLua/issues/413))
 - Fixed assignment with comment sometimes not hanging leading to malformed syntax. ([#416](https://github.com/JohnnyMorganz/StyLua/issues/416))
 - Fixed block ignores not applied when multiple leading block ignore comments are present at once. ([#421](https://github.com/JohnnyMorganz/StyLua/issues/421))
+- Fixed ordering of comments when semicolon after statement is removed. ([#423](https://github.com/JohnnyMorganz/StyLua/issues/423))
 
 ## [0.12.5] - 2022-03-08
 ### Fixed

src/formatters/block.rs

@@ -1,10 +1,12 @@
+#[cfg(feature = "luau")]
+use crate::formatters::general::format_symbol;
 use crate::{
     context::{create_indent_trivia, create_newline_trivia, Context, FormatNode},
     fmt_symbol,
     formatters::{
         assignment::hang_punctuated_list,
         expression::{format_expression, hang_expression},
-        general::{format_punctuated, format_punctuated_multiline, format_symbol},
+        general::{format_punctuated, format_punctuated_multiline},
         stmt::format_stmt,
         trivia::{
             strip_trailing_trivia, strip_trivia, FormatTriviaType, UpdateLeadingTrivia,
@@ -476,18 +478,29 @@ pub fn format_block(ctx: &Context, block: &Block, shape: Shape) -> Block {
                 Some(semi) => {
                     // We used to have a semicolon, but now we are removing it
                     // We want to keep any old comments on the semicolon token, otherwise we will lose it
-                    let (updated_stmt, trivia) = trivia_util::get_stmt_trailing_trivia(stmt);
-                    stmt = updated_stmt;
-                    // We will do a hack here, where we insert an empty token, and add all the remaining trivia onto it
-                    Some(
-                        format_symbol(
-                            &ctx,
-                            semi,
-                            &TokenReference::new(vec![], Token::new(TokenType::spaces(0)), vec![]),
-                            shape,
+                    // Move the comments to the end of the stmt, but before the newline token
+                    // TODO: this is a bit of a hack - we should probably move newline appending to this function
+                    let trivia = trivia_util::get_stmt_trailing_trivia(stmt.to_owned())
+                        .1
+                        .iter()
+                        .rev()
+                        .skip(1) // Remove the newline at the end
+                        .rev()
+                        .cloned()
+                        .chain(
+                            semi.trailing_trivia()
+                                .filter(|token| trivia_util::trivia_is_comment(token))
+                                .flat_map(|x| {
+                                    // Prepend a single space beforehand
+                                    vec![Token::new(TokenType::spaces(1)), x.to_owned()]
+                                }),
                         )
-                        .update_trailing_trivia(FormatTriviaType::Append(trivia)),
-                    )
+                        .chain(std::iter::once(create_newline_trivia(&ctx)))
+                        .collect();
+
+                    stmt = stmt.update_trailing_trivia(FormatTriviaType::Replace(trivia));
+
+                    None
                 }
                 None => None,
             },
@@ -510,24 +523,33 @@ pub fn format_block(ctx: &Context, block: &Block, shape: Shape) -> Block {
             {
                 last_stmt = last_stmt_remove_leading_newlines(last_stmt);
             }
+
             // LastStmt will never need a semicolon
             // We need to check if we previously had a semicolon, and keep the comments if so
             let semicolon = match semi {
                 Some(semi) => {
-                    let (updated_last_stmt, trivia) =
-                        trivia_util::get_last_stmt_trailing_trivia(last_stmt);
-                    last_stmt = updated_last_stmt;
+                    // Append semicolon trailing trivia to the end, but before the newline
+                    // TODO: this is a bit of a hack - we should probably move newline appending to this function
+                    let trivia = trivia_util::last_stmt_trailing_trivia(&last_stmt)
+                        .iter()
+                        .rev()
+                        .skip(1) // Remove the newline at the end
+                        .rev()
+                        .cloned()
+                        .chain(
+                            semi.trailing_trivia()
+                                .filter(|token| trivia_util::trivia_is_comment(token))
+                                .flat_map(|x| {
+                                    // Prepend a single space beforehand
+                                    vec![Token::new(TokenType::spaces(1)), x.to_owned()]
+                                }),
+                        )
+                        .chain(std::iter::once(create_newline_trivia(&ctx)))
+                        .collect();
 
-                    // We want to keep any old comments on the semicolon token, otherwise we will lose it
-                    // We will do a hack here, where we replace the semicolon with an empty symbol
-                    let semicolon_token = format_symbol(
-                        &ctx,
-                        semi,
-                        &TokenReference::new(vec![], Token::new(TokenType::spaces(0)), vec![]),
-                        shape,
-                    )
-                    .update_trailing_trivia(FormatTriviaType::Append(trivia));
-                    Some(semicolon_token)
+                    last_stmt = last_stmt.update_trailing_trivia(FormatTriviaType::Replace(trivia));
+
+                    None
                 }
                 None => None,
             };

src/formatters/trivia.rs

@@ -6,8 +6,8 @@ use full_moon::ast::types::{
 };
 use full_moon::ast::{
     punctuated::Punctuated, span::ContainedSpan, BinOp, Call, Expression, FunctionArgs,
-    FunctionBody, FunctionCall, FunctionName, Index, MethodCall, Parameter, Prefix, Suffix,
-    TableConstructor, UnOp, Value, Var, VarExpression,
+    FunctionBody, FunctionCall, FunctionName, Index, LastStmt, MethodCall, Parameter, Prefix, Stmt,
+    Suffix, TableConstructor, UnOp, Value, Var, VarExpression,
 };
 use full_moon::tokenizer::{Token, TokenReference};
 
@@ -362,6 +362,30 @@ define_update_trivia!(Index, |this, leading, trailing| {
     }
 });
 
+define_update_trailing_trivia!(LastStmt, |this, trailing| {
+    match this {
+        LastStmt::Break(token) => LastStmt::Break(token.update_trailing_trivia(trailing)),
+        #[cfg(feature = "luau")]
+        LastStmt::Continue(token) => LastStmt::Continue(token.update_trailing_trivia(trailing)),
+        LastStmt::Return(r#return) => {
+            if r#return.returns().is_empty() {
+                LastStmt::Return(
+                    r#return
+                        .to_owned()
+                        .with_token(r#return.token().update_trailing_trivia(trailing)),
+                )
+            } else {
+                LastStmt::Return(
+                    r#return
+                        .to_owned()
+                        .with_returns(r#return.returns().update_trailing_trivia(trailing)),
+                )
+            }
+        }
+        other => panic!("unknown node {:?}", other),
+    }
+});
+
 define_update_trivia!(MethodCall, |this, leading, trailing| {
     this.to_owned()
         .with_colon_token(this.colon_token().update_leading_trivia(leading))
@@ -431,6 +455,95 @@ where
     }
 }
 
+define_update_trailing_trivia!(Stmt, |this, trailing| {
+    match this {
+        Stmt::Assignment(assignment) => {
+            let expressions = assignment.expressions().update_trailing_trivia(trailing);
+            Stmt::Assignment(assignment.to_owned().with_expressions(expressions))
+        }
+
+        Stmt::LocalAssignment(local_assignment) => {
+            if local_assignment.expressions().is_empty() {
+                let names = local_assignment.names().update_trailing_trivia(trailing);
+                Stmt::LocalAssignment(local_assignment.to_owned().with_names(names))
+            } else {
+                let expressions = local_assignment
+                    .expressions()
+                    .update_trailing_trivia(trailing);
+                Stmt::LocalAssignment(local_assignment.to_owned().with_expressions(expressions))
+            }
+        }
+        Stmt::FunctionCall(function_call) => {
+            Stmt::FunctionCall(function_call.update_trailing_trivia(trailing))
+        }
+        Stmt::Repeat(repeat_block) => {
+            let until = repeat_block.until().update_trailing_trivia(trailing);
+            Stmt::Repeat(repeat_block.to_owned().with_until(until))
+        }
+        Stmt::Do(stmt) => {
+            let end_token = stmt.end_token().update_trailing_trivia(trailing);
+            Stmt::Do(stmt.to_owned().with_end_token(end_token))
+        }
+        Stmt::GenericFor(stmt) => {
+            let end_token = stmt.end_token().update_trailing_trivia(trailing);
+            Stmt::GenericFor(stmt.to_owned().with_end_token(end_token))
+        }
+        Stmt::If(stmt) => {
+            let end_token = stmt.end_token().update_trailing_trivia(trailing);
+            Stmt::If(stmt.to_owned().with_end_token(end_token))
+        }
+        Stmt::FunctionDeclaration(stmt) => {
+            let end_token = stmt.body().end_token().update_trailing_trivia(trailing);
+            let body = stmt.body().to_owned().with_end_token(end_token);
+            Stmt::FunctionDeclaration(stmt.to_owned().with_body(body))
+        }
+        Stmt::LocalFunction(stmt) => {
+            let end_token = stmt.body().end_token().update_trailing_trivia(trailing);
+            let body = stmt.body().to_owned().with_end_token(end_token);
+            Stmt::LocalFunction(stmt.to_owned().with_body(body))
+        }
+        Stmt::NumericFor(stmt) => {
+            let end_token = stmt.end_token().update_trailing_trivia(trailing);
+            Stmt::NumericFor(stmt.to_owned().with_end_token(end_token))
+        }
+        Stmt::While(stmt) => {
+            let end_token = stmt.end_token().update_trailing_trivia(trailing);
+            Stmt::While(stmt.to_owned().with_end_token(end_token))
+        }
+
+        #[cfg(feature = "luau")]
+        Stmt::CompoundAssignment(stmt) => {
+            let rhs = stmt.rhs().update_trailing_trivia(trailing);
+            Stmt::CompoundAssignment(stmt.to_owned().with_rhs(rhs))
+        }
+        #[cfg(feature = "luau")]
+        Stmt::ExportedTypeDeclaration(stmt) => {
+            let type_declaration = stmt.type_declaration().to_owned().with_type_definition(
+                stmt.type_declaration()
+                    .type_definition()
+                    .update_trailing_trivia(trailing),
+            );
+            Stmt::ExportedTypeDeclaration(stmt.to_owned().with_type_declaration(type_declaration))
+        }
+        #[cfg(feature = "luau")]
+        Stmt::TypeDeclaration(stmt) => Stmt::TypeDeclaration(
+            stmt.to_owned()
+                .with_type_definition(stmt.type_definition().update_trailing_trivia(trailing)),
+        ),
+        #[cfg(feature = "lua52")]
+        Stmt::Goto(stmt) => Stmt::Goto(
+            stmt.to_owned()
+                .with_label_name(stmt.label_name().update_trailing_trivia(trailing)),
+        ),
+        #[cfg(feature = "lua52")]
+        Stmt::Label(stmt) => Stmt::Label(
+            stmt.to_owned()
+                .with_right_colons(stmt.right_colons().update_trailing_trivia(trailing)),
+        ),
+        other => panic!("unknown node {:?}", other),
+    }
+});
+
 define_update_trivia!(Suffix, |this, leading, trailing| {
     match this {
         Suffix::Call(call) => Suffix::Call(call.update_trivia(leading, trailing)),

src/formatters/trivia_util.rs

@@ -982,50 +982,19 @@ pub fn get_stmt_trailing_trivia(stmt: Stmt) -> (Stmt, Vec<Token>) {
     (updated_stmt, trailing_trivia)
 }
 
-pub fn get_last_stmt_trailing_trivia(last_stmt: LastStmt) -> (LastStmt, Vec<Token>) {
+pub fn last_stmt_trailing_trivia(last_stmt: &LastStmt) -> Vec<Token> {
     match last_stmt {
         LastStmt::Return(ret) => {
-            let mut return_token = ret.token().to_owned();
-            let mut formatted_expression_list = ret.returns().to_owned();
-            let mut trailing_trivia = Vec::new();
-
-            // Retrieve last item and remove trailing trivia
-            if let Some(last_pair) = formatted_expression_list.pop() {
-                let pair = last_pair.map(|value| {
-                    trailing_trivia = get_expression_trailing_trivia(&value);
-                    value.update_trailing_trivia(FormatTriviaType::Replace(vec![]))
-                });
-                formatted_expression_list.push(pair);
+            if ret.returns().is_empty() {
+                ret.token().trailing_trivia().cloned().collect()
             } else {
-                trailing_trivia = return_token
-                    .trailing_trivia()
-                    .map(|x| x.to_owned())
-                    .collect();
-                return_token =
-                    return_token.update_trailing_trivia(FormatTriviaType::Replace(vec![]));
+                let last_expression = ret.returns().iter().last().unwrap();
+                get_expression_trailing_trivia(last_expression)
             }
-
-            (
-                LastStmt::Return(
-                    ret.with_token(return_token)
-                        .with_returns(formatted_expression_list),
-                ),
-                trailing_trivia,
-            )
-        }
-        LastStmt::Break(token) => {
-            let trailing_trivia = token.trailing_trivia().map(|x| x.to_owned()).collect();
-            let token = token.update_trailing_trivia(FormatTriviaType::Replace(vec![]));
-
-            (LastStmt::Break(token), trailing_trivia)
         }
+        LastStmt::Break(token) => token.trailing_trivia().cloned().collect(),
         #[cfg(feature = "luau")]
-        LastStmt::Continue(token) => {
-            let trailing_trivia = token.trailing_trivia().map(|x| x.to_owned()).collect();
-            let token = token.update_trailing_trivia(FormatTriviaType::Replace(vec![]));
-
-            (LastStmt::Continue(token), trailing_trivia)
-        }
+        LastStmt::Continue(token) => token.trailing_trivia().cloned().collect(),
         other => panic!("unknown node {:?}", other),
     }
 }

tests/inputs/semicolon-comments-ordering.lua

@@ -0,0 +1,2 @@
+local foo = a <= b --[[ some block comment ]]; -- inline comment
+fn() --[[ some block comment 2 ]]; -- inline comment 2

tests/snapshots/tests__standard@semicolon-comments-ordering.lua.snap

@@ -0,0 +1,9 @@
+---
+source: tests/tests.rs
+assertion_line: 12
+expression: format(&contents)
+
+---
+local foo = a <= b --[[ some block comment ]] -- inline comment
+fn() --[[ some block comment 2 ]] -- inline comment 2
+