Keep multiline comments in place before commas (#822)

* Keep multiline comments in place when formatting punctuated sequences * Update changelog
JohnnyMorganz · Nov 11, 2023 · edf11fe · edf11fe
1 parent 095e2b0
commit edf11fe
Show file tree

Hide file tree

Showing 7 changed files with 107 additions and 23 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -21,9 +21,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - Fix Luau `\z` escape parsing
 
 - Simplified access and modification patterns for StyLua configuration. You can now access the properties directly
+
   - **Deprecated:** the old access patterns of `.property()` and `.with_property()` are now deprecated
   - **Breaking Change (WASM):** due to JS/TS lack of differentiation between `.property` / `.property()` implementation, the `.property()` functions were removed from WASM output.
 
+- Multiline comments before commas will now remain in place and not move to after the comma. This is to support type-assertions-via-comments that is commonly used by some language servers. ([#778](https://github.com/JohnnyMorganz/StyLua/issues/778))
+
 ### Fixed
 
 - Wasm build now correctly supports configuring sort requires ([#818](https://github.com/JohnnyMorganz/StyLua/issues/818))

diff --git a/src/formatters/general.rs b/src/formatters/general.rs
@@ -5,8 +5,8 @@ use crate::{
     formatters::{
         trivia::{FormatTriviaType, UpdateLeadingTrivia, UpdateTrailingTrivia, UpdateTrivia},
         trivia_util::{
-            self, punctuated_inline_comments, take_trailing_comments, GetLeadingTrivia,
-            GetTrailingTrivia, HasInlineComments,
+            self, punctuated_inline_comments, CommentSearch, GetLeadingTrivia, GetTrailingTrivia,
+            HasInlineComments,
         },
     },
     shape::Shape,
@@ -522,9 +522,15 @@ where
                 ctx, shape,
             )]));
 
-        // Take any trailing trivia (i.e. comments) from the argument, and append it to the end of the punctuation
-        let (formatted_argument, mut trailing_comments) =
-            take_trailing_comments(&formatted_argument);
+        // Any singleline comments must be moved to after the punctuation
+        // We should keep multiline comments in the same location
+        let multiline_comments =
+            formatted_argument.trailing_comments_search(CommentSearch::Multiline);
+        let singleline_comments =
+            formatted_argument.trailing_comments_search(CommentSearch::Single);
+
+        let formatted_argument = formatted_argument
+            .update_trailing_trivia(FormatTriviaType::Replace(multiline_comments));
 
         let punctuation = match argument.punctuation() {
             Some(punctuation) => {
@@ -544,8 +550,8 @@ where
                             x,
                         ]
                     })
+                    .chain(singleline_comments)
                     .collect();
-                trailing_trivia.append(&mut trailing_comments);
                 trailing_trivia.push(create_newline_trivia(ctx));
 
                 let symbol = symbol.update_trivia(
@@ -559,7 +565,7 @@ where
             // We need to do this because in function declarations, we format parameters but if they have a type
             // specifier we don't have access to put it after the type specifier
             None => Some(TokenReference::new(
-                trailing_comments,
+                singleline_comments,
                 create_newline_trivia(ctx),
                 vec![],
             )),

diff --git a/src/formatters/luau.rs b/src/formatters/luau.rs
@@ -839,15 +839,18 @@ pub fn format_type_field(
     let shape = shape + (strip_leading_trivia(&key).to_string().len() + 2);
     let mut value = format_type_info(ctx, type_field.value(), shape);
 
-    let trailing_trivia = value.trailing_trivia();
+    // Trailing trivia consists only of single line comments - multiline comments are kept in place
+    let trailing_trivia = value.trailing_comments_search(CommentSearch::Single);
 
     if let TableType::MultiLine = table_type {
         // If still over budget, hang the type
         if can_hang_type(type_field.value()) && shape.test_over_budget(&value) {
             value = hang_type_info(ctx, type_field.value(), TypeInfoContext::new(), shape, 1)
         };
 
-        value = value.update_trailing_trivia(FormatTriviaType::Replace(vec![]))
+        // Keep multiline comments in place
+        let multiline_comments = value.trailing_comments_search(CommentSearch::Multiline);
+        value = value.update_trailing_trivia(FormatTriviaType::Replace(multiline_comments))
     }
 
     (

diff --git a/src/formatters/table.rs b/src/formatters/table.rs
@@ -9,7 +9,7 @@ use crate::{
             trivia_to_vec, EndTokenType, FormatTokenType,
         },
         trivia::{strip_trivia, FormatTriviaType, UpdateLeadingTrivia, UpdateTrailingTrivia},
-        trivia_util::{self, GetTrailingTrivia, HasInlineComments},
+        trivia_util::{self, CommentSearch, GetTrailingTrivia, HasInlineComments},
     },
     shape::Shape,
 };
@@ -40,17 +40,19 @@ fn format_field_expression_value(
     expression: &Expression,
     shape: Shape,
 ) -> Expression {
-    // Remove all trivia from the output expression as it will be moved after the comma
+    // Remove singleline comments from the output expression as it will be moved after the comma
+    // Retain multiline comments in place
+    let multiline_comments = expression.trailing_comments_search(CommentSearch::Multiline);
+    let trailing_trivia = FormatTriviaType::Replace(multiline_comments);
 
     if trivia_util::can_hang_expression(expression) {
         if expression.has_inline_comments() {
-            hang_expression(ctx, expression, shape, Some(1))
-                .update_trailing_trivia(FormatTriviaType::Replace(vec![]))
+            hang_expression(ctx, expression, shape, Some(1)).update_trailing_trivia(trailing_trivia)
         } else {
             let singleline_value = format_expression(ctx, expression, shape)
-                .update_trailing_trivia(FormatTriviaType::Replace(vec![]));
+                .update_trailing_trivia(trailing_trivia.clone());
             let hanging_value = hang_expression(ctx, expression, shape, Some(1))
-                .update_trailing_trivia(FormatTriviaType::Replace(vec![]));
+                .update_trailing_trivia(trailing_trivia);
 
             if shape.test_over_budget(&singleline_value)
                 || format!("{hanging_value}").lines().count()
@@ -62,8 +64,7 @@ fn format_field_expression_value(
             }
         }
     } else {
-        format_expression(ctx, expression, shape)
-            .update_trailing_trivia(FormatTriviaType::Replace(vec![]))
+        format_expression(ctx, expression, shape).update_trailing_trivia(trailing_trivia)
     }
 }
 
@@ -139,6 +140,8 @@ fn format_field(
         _ => FormatTriviaType::NoChange,
     };
 
+    // Trailing trivia is taken out and moved to after the comma
+    // We only move singleline comments, multiline comments remain in place
     let trailing_trivia;
     let field = match field {
         Field::ExpressionKey {
@@ -147,7 +150,7 @@ fn format_field(
             equal,
             value,
         } => {
-            trailing_trivia = value.trailing_trivia();
+            trailing_trivia = value.trailing_comments_search(CommentSearch::Single);
             let brackets = format_contained_span(ctx, brackets, shape);
 
             let space_brackets = is_brackets_string(key);
@@ -185,7 +188,7 @@ fn format_field(
             }
         }
         Field::NameKey { key, equal, value } => {
-            trailing_trivia = value.trailing_trivia();
+            trailing_trivia = value.trailing_comments_search(CommentSearch::Single);
             let key = format_token_reference(ctx, key, shape);
 
             // Get the new leading comments to add before the key, and the equal token
@@ -205,7 +208,7 @@ fn format_field(
             Field::NameKey { key, equal, value }
         }
         Field::NoKey(expression) => {
-            trailing_trivia = expression.trailing_trivia();
+            trailing_trivia = expression.trailing_comments_search(CommentSearch::Single);
 
             if let TableType::MultiLine = table_type {
                 let formatted_expression = format_field_expression_value(ctx, expression, shape);
@@ -305,7 +308,8 @@ where
 
         // Format the field. We will ignore the taken trailing trivia, as we do not need it.
         // (If there were any comments present, this function should never have been called)
-        let formatted_field = formatter(ctx, field, table_type, shape).0;
+        let (formatted_field, trailing_trivia) = formatter(ctx, field, table_type, shape);
+        assert!(trailing_trivia.is_empty());
 
         let formatted_punctuation = match current_fields.peek() {
             Some(_) => {
@@ -370,6 +374,7 @@ where
             trailing_trivia = Vec::new();
         } else {
             // Filter trailing trivia for any newlines
+            // NOTE: in practice, this should only consist of singleline comments
             trailing_trivia = trailing_trivia
                 .iter()
                 .filter(|x| !trivia_util::trivia_is_whitespace(x))

diff --git a/src/formatters/trivia_util.rs b/src/formatters/trivia_util.rs
@@ -44,16 +44,20 @@ pub trait GetTrailingTrivia {
 
     // Retrieves all the trailing comments from the token
     // Prepends a space before each comment
-    fn trailing_comments(&self) -> Vec<Token> {
+    fn trailing_comments_search(&self, search: CommentSearch) -> Vec<Token> {
         self.trailing_trivia()
             .iter()
-            .filter(|token| trivia_is_comment(token))
+            .filter(|token| trivia_is_comment_search(token, search))
             .flat_map(|x| {
                 // Prepend a single space beforehand
                 vec![Token::new(TokenType::spaces(1)), x.to_owned()]
             })
             .collect()
     }
+
+    fn trailing_comments(&self) -> Vec<Token> {
+        self.trailing_comments_search(CommentSearch::All)
+    }
 }
 
 pub fn trivia_is_whitespace(trivia: &Token) -> bool {
@@ -64,13 +68,25 @@ pub fn trivia_is_singleline_comment(trivia: &Token) -> bool {
     matches!(trivia.token_kind(), TokenKind::SingleLineComment)
 }
 
+fn trivia_is_multiline_comment(trivia: &Token) -> bool {
+    matches!(trivia.token_kind(), TokenKind::MultiLineComment)
+}
+
 pub fn trivia_is_comment(trivia: &Token) -> bool {
     matches!(
         trivia.token_kind(),
         TokenKind::SingleLineComment | TokenKind::MultiLineComment
     )
 }
 
+fn trivia_is_comment_search(trivia: &Token, search: CommentSearch) -> bool {
+    match search {
+        CommentSearch::Single => trivia_is_singleline_comment(trivia),
+        CommentSearch::Multiline => trivia_is_multiline_comment(trivia),
+        CommentSearch::All => trivia_is_comment(trivia),
+    }
+}
+
 pub fn trivia_is_newline(trivia: &Token) -> bool {
     if let TokenType::Whitespace { characters } = trivia.token_type() {
         if characters.find('\n').is_some() {
@@ -825,6 +841,8 @@ impl GetTrailingTrivia for LastStmt {
 pub enum CommentSearch {
     // Only care about singleline comments
     Single,
+    // Only care about multiline comments
+    Multiline,
     // Looking for all comments
     All,
 }
@@ -835,6 +853,7 @@ fn trivia_contains_comments<'a>(
 ) -> bool {
     let tester = match search {
         CommentSearch::Single => trivia_is_singleline_comment,
+        CommentSearch::Multiline => trivia_is_multiline_comment,
         CommentSearch::All => trivia_is_comment,
     };
 

diff --git a/tests/inputs/comments-before-punctuation.lua b/tests/inputs/comments-before-punctuation.lua
@@ -0,0 +1,21 @@
+-- https://github.com/JohnnyMorganz/StyLua/issues/778
+-- comments should stay before punctuation to ensure type assertions work in sumneko-lua
+
+function fun(
+	a --[[ a commnet]],
+	b
+)
+end
+
+local tab = {
+	a = 1 --[[@as integer ]],
+	b = 1,
+}
+
+call(
+	long_argument_name --[[@as integer ]],
+	long_argument_name,
+	long_argument_name,
+	long_argument_name,
+	long_argument_name
+)
diff --git a/tests/snapshots/[email protected] b/tests/snapshots/[email protected]
@@ -0,0 +1,27 @@
+---
+source: tests/tests.rs
+expression: format(&contents)
+input_file: tests/inputs/comments-before-punctuation.lua
+---
+-- https://github.com/JohnnyMorganz/StyLua/issues/778
+-- comments should stay before punctuation to ensure type assertions work in sumneko-lua
+
+function fun(
+	a --[[ a commnet]],
+	b
+)
+end
+
+local tab = {
+	a = 1 --[[@as integer ]],
+	b = 1,
+}
+
+call(
+	long_argument_name --[[@as integer ]],
+	long_argument_name,
+	long_argument_name,
+	long_argument_name,
+	long_argument_name
+)
+