Fix mistransformation of generic for comments (#580)

* Fix incorrect token used in multiline generic for

* Add test case

* Fix check for comments in generic_for expr

* Fix trivia location for multiline generic for

* Add snapshot

* Update changelog

Author: JohnnyMorganz

CHANGELOG.md

@@ -19,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Precommit hook now supports downloading aarch64 binary for M1 macs ([#558](https://github.com/JohnnyMorganz/StyLua/issues/558))
+- Fixed mistransformations of generic for loop with comments in the expression list ([#579](https://github.com/JohnnyMorganz/StyLua/issues/579))
 
 ## [0.14.3] - 2022-08-27
 

src/formatters/stmt.rs

@@ -221,10 +221,12 @@ pub fn format_generic_for(ctx: &Context, generic_for: &GenericFor, shape: Shape)
     let singleline_expr =
         format_punctuated(ctx, generic_for.expressions(), shape, format_expression);
     let singleline_expr_shape = shape.take_first_line(&singleline_expr);
-    let requires_expr_multiline = (trivia_util::contains_comments(generic_for.expressions())
-        || trivia_util::spans_multiple_lines(&singleline_expr)
-        || singleline_expr_shape.over_budget())
-        && !hug_generic_for(generic_for.expressions());
+    let requires_expr_multiline =
+        (trivia_util::token_contains_trailing_comments(generic_for.in_token())
+            || trivia_util::contains_comments(generic_for.expressions())
+            || trivia_util::spans_multiple_lines(&singleline_expr)
+            || singleline_expr_shape.over_budget())
+            && !hug_generic_for(generic_for.expressions());
 
     let in_token = match (require_names_multiline, requires_expr_multiline) {
         (true, true) => fmt_symbol!(ctx, generic_for.in_token(), "in", shape)
@@ -245,23 +247,25 @@ pub fn format_generic_for(ctx: &Context, generic_for: &GenericFor, shape: Shape)
     let expr_list = match requires_expr_multiline {
         true => {
             let shape = shape.reset().increment_additional_indent();
-            format_punctuated_multiline(
+            let expr_list = format_punctuated_multiline(
                 ctx,
                 generic_for.expressions(),
                 shape,
                 format_expression,
                 None,
+            );
+            trivia_util::prepend_newline_indent(
+                ctx,
+                &expr_list,
+                trivia_util::punctuated_leading_trivia(&expr_list).iter(),
+                shape,
             )
-            .update_leading_trivia(FormatTriviaType::Append(vec![
-                create_newline_trivia(ctx),
-                create_indent_trivia(ctx, shape),
-            ]))
         }
         false => singleline_expr,
     };
 
     let do_token = match requires_expr_multiline {
-        true => fmt_symbol!(ctx, generic_for.in_token(), "do", shape).update_leading_trivia(
+        true => fmt_symbol!(ctx, generic_for.do_token(), "do", shape).update_leading_trivia(
             FormatTriviaType::Append(vec![
                 create_newline_trivia(ctx),
                 create_indent_trivia(ctx, shape),

src/formatters/trivia_util.rs

@@ -12,8 +12,8 @@ use full_moon::ast::types::{
 };
 use full_moon::{
     ast::{
-        BinOp, Block, Call, Expression, Field, FunctionArgs, Index, LastStmt, Parameter, Prefix,
-        Stmt, Suffix, TableConstructor, UnOp, Value, Var,
+        punctuated::Punctuated, BinOp, Block, Call, Expression, Field, FunctionArgs, Index,
+        LastStmt, Parameter, Prefix, Stmt, Suffix, TableConstructor, UnOp, Value, Var,
     },
     node::Node,
     tokenizer::{Token, TokenKind, TokenReference, TokenType},
@@ -420,6 +420,13 @@ 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)

tests/inputs/generic-for-multiline-2.lua

@@ -0,0 +1,19 @@
+-- https://github.com/JohnnyMorganz/StyLua/issues/579
+for _, item in
+	-- comment
+	call()
+do
+end
+
+
+for _, item in -- comment
+	call()
+do
+end
+
+
+for _, item in           -- comment
+
+	call()
+do
+end

tests/snapshots/tests__standard@generic-for-multiline-2.lua.snap

@@ -0,0 +1,21 @@
+---
+source: tests/tests.rs
+expression: format(&contents)
+---
+-- https://github.com/JohnnyMorganz/StyLua/issues/579
+for _, item in
+	-- comment
+	call()
+do
+end
+
+for _, item in -- comment
+	call()
+do
+end
+
+for _, item in -- comment
+	call()
+do
+end
+