Fix comments in punctuated list for returns and assignments

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Bumped internal parser dependency which should fix parsing problems for comments with Chinese characters, and multiline string escapes
+- Fixed comments in punctuated lists for return statements or assignments being incorrectly formatted leading to syntax errors ([#662](https://github.com/JohnnyMorganz/StyLua/issues/662))
 
 ## [0.17.0] - 2023-03-11
 

src/formatters/assignment.rs

@@ -31,6 +31,8 @@ use crate::{
     shape::Shape,
 };
 
+use super::trivia_util::{prepend_newline_indent, GetTrailingTrivia};
+
 /// Calculates the hanging level to use when hanging an expression.
 /// By default, we indent one further, but we DO NOT want to do this if the expression is just parentheses (or a unary operation on them)
 /// https://github.com/JohnnyMorganz/StyLua/issues/274
@@ -181,7 +183,7 @@ fn attempt_assignment_tactics(
             format_expression,
         );
 
-        if trivia_util::punctuated_expression_inline_comments(expressions)
+        if trivia_util::punctuated_inline_comments(expressions, true)
             || hanging_shape
                 .take_first_line(&strip_trivia(&expr_list))
                 .over_budget()
@@ -326,12 +328,12 @@ fn attempt_assignment_tactics(
 pub fn format_assignment_no_trivia(
     ctx: &Context,
     assignment: &Assignment,
-    shape: Shape,
+    mut shape: Shape,
 ) -> Assignment {
     // Check if the assignment expressions or equal token contain comments. If they do, we bail out of determining any tactics
     // and format multiline
     let contains_comments = trivia_util::token_contains_comments(assignment.equal_token())
-        || trivia_util::punctuated_expression_inline_comments(assignment.expressions());
+        || trivia_util::punctuated_inline_comments(assignment.expressions(), true);
 
     // Firstly attempt to format the assignment onto a single line, using an infinite column width shape
     let mut var_list = try_format_punctuated(
@@ -349,6 +351,16 @@ pub fn format_assignment_no_trivia(
         format_expression,
     );
 
+    // If the var list ended with a comment, we need to hang the equals token
+    if var_list.has_trailing_comments(trivia_util::CommentSearch::Single) {
+        const EQUAL_TOKEN_LEN: usize = "= ".len();
+        shape = shape
+            .reset()
+            .increment_additional_indent()
+            .add_width(EQUAL_TOKEN_LEN);
+        equal_token = prepend_newline_indent(ctx, &equal_token, shape);
+    }
+
     // Test the assignment to see if its over width
     let singleline_shape = shape
         + (strip_leading_trivia(&var_list).to_string().len()
@@ -420,7 +432,7 @@ fn format_local_no_assignment(
 pub fn format_local_assignment_no_trivia(
     ctx: &Context,
     assignment: &LocalAssignment,
-    shape: Shape,
+    mut shape: Shape,
 ) -> LocalAssignment {
     if assignment.expressions().is_empty() {
         format_local_no_assignment(ctx, assignment, shape)
@@ -430,7 +442,7 @@ pub fn format_local_assignment_no_trivia(
         let contains_comments = assignment
             .equal_token()
             .map_or(false, trivia_util::token_contains_comments)
-            || trivia_util::punctuated_expression_inline_comments(assignment.expressions());
+            || trivia_util::punctuated_inline_comments(assignment.expressions(), true);
 
         // Firstly attempt to format the assignment onto a single line, using an infinite column width shape
         let local_token = fmt_symbol!(ctx, assignment.local_token(), "local ", shape);
@@ -477,6 +489,16 @@ pub fn format_local_assignment_no_trivia(
             });
         }
 
+        // If the var list ended with a comment, we need to hang the equals token
+        if name_list.has_trailing_comments(trivia_util::CommentSearch::Single) {
+            const EQUAL_TOKEN_LEN: usize = "= ".len();
+            shape = shape
+                .reset()
+                .increment_additional_indent()
+                .add_width(EQUAL_TOKEN_LEN);
+            equal_token = prepend_newline_indent(ctx, &equal_token, shape);
+        }
+
         // Test the assignment to see if its over width
         let singleline_shape = shape
             + (strip_leading_trivia(&name_list).to_string().len()

src/formatters/block.rs

@@ -55,7 +55,7 @@ pub fn format_return(ctx: &Context, return_node: &Return, shape: Shape) -> Retur
         );
 
         let contains_comments = return_token_trailing_comments
-            || trivia_util::punctuated_expression_inline_comments(returns);
+            || trivia_util::punctuated_inline_comments(returns, true);
         let is_function_or_table = returns.iter().all(is_function_or_table_constructor);
 
         // See if we need to format multiline

src/formatters/general.rs

@@ -2,7 +2,10 @@ use crate::{
     context::{create_indent_trivia, create_newline_trivia, Context, FormatNode},
     formatters::{
         trivia::{FormatTriviaType, UpdateLeadingTrivia, UpdateTrailingTrivia, UpdateTrivia},
-        trivia_util::{self, GetLeadingTrivia},
+        trivia_util::{
+            self, punctuated_inline_comments, GetLeadingTrivia, GetTrailingTrivia,
+            HasInlineComments,
+        },
     },
     shape::Shape,
     QuoteStyle,
@@ -384,7 +387,7 @@ pub fn format_punctuated_multiline<T, F>(
     hang_level: Option<usize>,
 ) -> Punctuated<T>
 where
-    T: Node + GetLeadingTrivia + UpdateLeadingTrivia,
+    T: Node + GetLeadingTrivia + UpdateLeadingTrivia + GetTrailingTrivia + UpdateTrailingTrivia,
     F: Fn(&Context, &T, Shape) -> T,
 {
     let mut formatted: Punctuated<T> = Punctuated::new();
@@ -412,7 +415,16 @@ where
                     trivia_util::prepend_newline_indent(ctx, &value, hanging_shape)
                 };
 
+                // Handle any comments in between the value and the punctuation
+                // If they are present, then move them to after the punctuation
+                let mut trailing_comments = value.trailing_comments();
+                let value = value.update_trailing_trivia(FormatTriviaType::Replace(vec![]));
+
                 let punctuation = fmt_symbol!(ctx, punctuation, ",", shape);
+                trailing_comments.append(&mut punctuation.trailing_trivia().cloned().collect());
+                let punctuation = punctuation
+                    .update_trailing_trivia(FormatTriviaType::Replace(trailing_comments));
+
                 formatted.push(Pair::new(value, Some(punctuation)));
             }
             Pair::End(value) => {
@@ -440,19 +452,18 @@ pub fn try_format_punctuated<T, F>(
     hang_level: Option<usize>,
 ) -> Punctuated<T>
 where
-    T: Node + GetLeadingTrivia + UpdateLeadingTrivia + std::fmt::Display,
+    T: Node
+        + GetLeadingTrivia
+        + UpdateLeadingTrivia
+        + GetTrailingTrivia
+        + UpdateTrailingTrivia
+        + HasInlineComments
+        + std::fmt::Display,
     F: Fn(&Context, &T, Shape) -> T,
 {
-    let mut format_multiline = false;
-
-    for pair in old.pairs() {
-        if let Pair::Punctuated(_, punctuation) = pair {
-            if trivia_util::contains_comments(punctuation) {
-                format_multiline = true;
-                break;
-            }
-        }
-    }
+    // TODO: we do not check the leading comments of the punctuated list for determining multiline
+    // Maybe we should do later?
+    let format_multiline = punctuated_inline_comments(old, false);
 
     if format_multiline {
         format_punctuated_multiline(ctx, old, shape, value_formatter, hang_level)

src/formatters/trivia_util.rs

@@ -450,13 +450,6 @@ pub fn get_expression_leading_trivia(expression: &Expression) -> Vec<Token> {
     }
 }
 
-pub fn punctuated_leading_trivia(punctuated: &Punctuated<Expression>) -> Vec<Token> {
-    punctuated
-        .iter()
-        .next()
-        .map_or_else(Vec::new, get_expression_leading_trivia)
-}
-
 pub fn binop_leading_comments(binop: &BinOp) -> Vec<Token> {
     match binop {
         BinOp::And(token)
@@ -1307,17 +1300,39 @@ pub fn expression_contains_inline_comments(expression: &Expression) -> bool {
     }
 }
 
-pub fn punctuated_expression_inline_comments(punctuated: &Punctuated<Expression>) -> bool {
-    punctuated.pairs().any(|pair| {
-        pair.punctuation().map_or(false, token_contains_comments)
-            || !expression_leading_comments(pair.value()).is_empty()
-            || expression_contains_inline_comments(pair.value())
-    })
+pub fn punctuated_inline_comments<T: GetLeadingTrivia + GetTrailingTrivia + HasInlineComments>(
+    punctuated: &Punctuated<T>,
+    include_leading: bool,
+) -> bool {
+    let mut iter = punctuated.pairs().peekable();
+    while let Some(pair) = iter.next() {
+        // Only check trailing comments on the expression if this is not the last pair
+        if iter.peek().is_some() && !pair.value().trailing_comments().is_empty() {
+            return true;
+        }
+
+        if pair.punctuation().map_or(false, token_contains_comments)
+            || (include_leading && !pair.value().leading_comments().is_empty())
+            || pair.value().has_inline_comments()
+        {
+            return true;
+        }
+    }
+
+    false
 }
 
 // TODO: can we change this from returning a Vec to just a plain iterator?
 pub trait GetLeadingTrivia {
     fn leading_trivia(&self) -> Vec<Token>;
+
+    fn leading_comments(&self) -> Vec<Token> {
+        self.leading_trivia()
+            .iter()
+            .filter(|token| trivia_is_comment(token))
+            .cloned()
+            .collect()
+    }
 }
 
 impl GetLeadingTrivia for TokenReference {
@@ -1338,9 +1353,11 @@ impl GetLeadingTrivia for Expression {
     }
 }
 
-impl GetLeadingTrivia for Punctuated<Expression> {
+impl<T: GetLeadingTrivia> GetLeadingTrivia for Punctuated<T> {
     fn leading_trivia(&self) -> Vec<Token> {
-        punctuated_leading_trivia(self)
+        self.iter()
+            .next()
+            .map_or_else(Vec::new, GetLeadingTrivia::leading_trivia)
     }
 }
 
@@ -1376,6 +1393,74 @@ impl GetLeadingTrivia for Prefix {
     }
 }
 
+pub trait GetTrailingTrivia {
+    fn trailing_trivia(&self) -> Vec<Token>;
+
+    fn has_trailing_comments(&self, search: CommentSearch) -> bool {
+        trivia_contains_comments(self.trailing_trivia().iter(), search)
+    }
+
+    // Retrieves all the trailing comments from the token
+    // Prepends a space before each comment
+    fn trailing_comments(&self) -> Vec<Token> {
+        self.trailing_trivia()
+            .iter()
+            .filter(|token| trivia_is_comment(token))
+            .flat_map(|x| {
+                // Prepend a single space beforehand
+                vec![Token::new(TokenType::spaces(1)), x.to_owned()]
+            })
+            .collect()
+    }
+}
+
+impl GetTrailingTrivia for TokenReference {
+    fn trailing_trivia(&self) -> Vec<Token> {
+        self.trailing_trivia().cloned().collect()
+    }
+}
+
+impl GetTrailingTrivia for Suffix {
+    fn trailing_trivia(&self) -> Vec<Token> {
+        suffix_trailing_trivia(self)
+    }
+}
+
+impl GetTrailingTrivia for Expression {
+    fn trailing_trivia(&self) -> Vec<Token> {
+        get_expression_trailing_trivia(self)
+    }
+}
+
+impl GetTrailingTrivia for Var {
+    fn trailing_trivia(&self) -> Vec<Token> {
+        var_trailing_trivia(self)
+    }
+}
+
+impl<T: GetTrailingTrivia> GetTrailingTrivia for Punctuated<T> {
+    fn trailing_trivia(&self) -> Vec<Token> {
+        self.iter()
+            .last()
+            .map_or_else(Vec::new, GetTrailingTrivia::trailing_trivia)
+    }
+}
+
+pub trait HasInlineComments {
+    fn has_inline_comments(&self) -> bool {
+        false
+    }
+}
+
+impl HasInlineComments for Expression {
+    fn has_inline_comments(&self) -> bool {
+        expression_contains_inline_comments(self)
+    }
+}
+
+impl HasInlineComments for Var {}
+impl HasInlineComments for TokenReference {}
+
 // Commonly, we update trivia to add in a newline and indent trivia to the leading trivia of a token/node.
 // An issue with this is if we do not properly take into account comments. This function also handles any comments present
 // by also interspersing them with the required newline and indentation, so they are aligned correctly.

tests/inputs/assignment-comments-3.lua

@@ -0,0 +1,8 @@
+-- https://github.com/JohnnyMorganz/StyLua/issues/662
+local a, b
+= 1 -- adoc
+, 2 -- bdoc
+
+local a -- adoc
+, b -- bdoc
+= 1, 2

tests/inputs/return-comments-2.lua

@@ -0,0 +1,5 @@
+-- https://github.com/JohnnyMorganz/StyLua/issues/662
+function f()
+	return a -- adoc
+		, b -- bdoc
+end

tests/snapshots/tests__standard@assignment-comments-3.lua.snap

@@ -0,0 +1,13 @@
+---
+source: tests/tests.rs
+expression: format(&contents)
+---
+-- https://github.com/JohnnyMorganz/StyLua/issues/662
+local a, b =
+	1, -- adoc
+	2 -- bdoc
+
+local a, -- adoc
+	b -- bdoc
+	= 1, 2
+

tests/snapshots/tests__standard@return-comments-2.lua.snap

@@ -0,0 +1,10 @@
+---
+source: tests/tests.rs
+expression: format(&contents)
+---
+-- https://github.com/JohnnyMorganz/StyLua/issues/662
+function f()
+	return a, -- adoc
+		b -- bdoc
+end
+