Multiline formatting of generic for syntax

CHANGELOG.md

@@ -6,8 +6,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 ### Added
-- Added option `call_parentheses`
-Specify whether to apply parentheses on function calls with single string or table arg. Possible options:`Always`, `NoSingleString`, `NoSingleTable`, `None`.
+- Added option `call_parentheses`:
+Specify whether to apply parentheses on function calls with single string or table arg. Possible options: `Always` (default), `NoSingleString`, `NoSingleTable`, `None`. ([#329](https://github.com/JohnnyMorganz/StyLua/issues/329))
+- Added proper multiline hanging of generic for syntax. ([#322](https://github.com/JohnnyMorganz/StyLua/issues/322))
 
 ### Fixed
 - Fixed generic variadics not being handled under the `luau` feature flag. ([#333](https://github.com/JohnnyMorganz/StyLua/issues/333))

src/formatters/general.rs

@@ -327,66 +327,6 @@ pub fn format_token_reference(
 
     TokenReference::new(formatted_leading_trivia, token, formatted_trailing_trivia)
 }
-// Formats a punctuation for a Punctuated sequence
-// Removes any trailing comments to be stored in a comments buffer
-pub fn format_punctuation(punctuation: &TokenReference) -> (TokenReference, Vec<Token>) {
-    let trailing_comments = punctuation
-        .trailing_trivia()
-        .filter(|x| trivia_util::trivia_is_comment(x))
-        .map(|x| {
-            // Prepend a single space beforehand
-            vec![Token::new(TokenType::spaces(1)), x.to_owned()]
-        })
-        .flatten()
-        .collect();
-
-    (
-        TokenReference::new(
-            Vec::new(),
-            punctuation.token().to_owned(),
-            vec![Token::new(TokenType::spaces(1))], // Single space whitespace
-        ),
-        trailing_comments,
-    )
-}
-
-// Formats a Punctuated sequence with correct punctuated values
-// If there are any comments in between tied to the punctuation, they will be removed and stored in a returned comments buffer
-pub fn format_punctuated_buffer<T, F>(
-    ctx: &Context,
-    old: &Punctuated<T>,
-    shape: Shape,
-    value_formatter: F,
-) -> (Punctuated<T>, Vec<Token>)
-where
-    T: std::fmt::Display,
-    F: Fn(&Context, &T, Shape) -> T,
-{
-    let mut formatted: Punctuated<T> = Punctuated::new();
-    let mut comments_buffer = Vec::new();
-    let mut shape = shape;
-
-    for pair in old.pairs() {
-        match pair {
-            Pair::Punctuated(value, punctuation) => {
-                // Format punctuation and store any comments into buffer
-                let (formatted_punctuation, mut comments) = format_punctuation(punctuation);
-                comments_buffer.append(&mut comments);
-
-                let formatted_value = value_formatter(ctx, value, shape);
-                shape = shape + (formatted_value.to_string().len() + 2); // 2 = ", "
-
-                formatted.push(Pair::new(formatted_value, Some(formatted_punctuation)));
-            }
-            Pair::End(value) => {
-                let formatted_value = value_formatter(ctx, value, shape);
-                formatted.push(Pair::new(formatted_value, None));
-            }
-        }
-    }
-
-    (formatted, comments_buffer)
-}
 
 /// Formats a Punctuated sequence with correct punctuated values.
 /// This function assumes that there are no comments present which would lead to a syntax error if the list was collapsed.

src/formatters/stmt.rs

@@ -14,7 +14,8 @@ use crate::{
         expression::{format_expression, hang_expression_trailing_newline},
         functions::{format_function_call, format_function_declaration, format_local_function},
         general::{
-            format_end_token, format_punctuated_buffer, format_token_reference, EndTokenType,
+            format_end_token, format_punctuated, format_punctuated_multiline,
+            format_token_reference, EndTokenType,
         },
         trivia::{
             strip_trivia, FormatTriviaType, UpdateLeadingTrivia, UpdateTrailingTrivia, UpdateTrivia,
@@ -77,31 +78,109 @@ pub fn format_do_block(ctx: &Context, do_block: &Do, shape: Shape) -> Do {
 pub fn format_generic_for(ctx: &Context, generic_for: &GenericFor, shape: Shape) -> GenericFor {
     // Create trivia
     let leading_trivia = vec![create_indent_trivia(ctx, shape)];
-    let mut trailing_trivia = vec![create_newline_trivia(ctx)];
+    let trailing_trivia = vec![create_newline_trivia(ctx)];
 
-    // TODO: Should we actually update the shape here?
     let for_token = fmt_symbol!(ctx, generic_for.for_token(), "for ", shape)
         .update_leading_trivia(FormatTriviaType::Append(leading_trivia.to_owned()));
-    let (formatted_names, mut names_comments_buf) =
-        format_punctuated_buffer(ctx, generic_for.names(), shape, format_token_reference);
+    let shape = shape + 4; // 4 = "for "
+
+    // Format the names on a single line
+    // If it goes over the column width, or contains comments, then format it multiline
+    let singleline_names =
+        format_punctuated(ctx, generic_for.names(), shape, format_token_reference);
+    let singleline_shape = shape.take_first_line(&singleline_names);
+
+    let require_names_multiline = trivia_util::contains_comments(generic_for.names())
+        || trivia_util::spans_multiple_lines(&singleline_names)
+        || singleline_shape.over_budget();
+
+    let for_token = match require_names_multiline {
+        true => fmt_symbol!(ctx, generic_for.for_token(), "for", shape)
+            .update_leading_trivia(FormatTriviaType::Append(leading_trivia.to_owned())),
+        false => for_token,
+    };
+
+    let names = match require_names_multiline {
+        true => {
+            let shape = shape.reset().increment_additional_indent();
+            format_punctuated_multiline(
+                ctx,
+                generic_for.names(),
+                shape,
+                format_token_reference,
+                None,
+            )
+            .update_leading_trivia(FormatTriviaType::Append(vec![
+                create_newline_trivia(ctx),
+                create_indent_trivia(ctx, shape),
+            ]))
+        }
+        false => singleline_names,
+    };
+
+    let shape = match require_names_multiline {
+        true => shape.reset() + 3,     // 3 = "in "
+        false => singleline_shape + 4, // 4 = " in "
+    };
 
     #[cfg(feature = "luau")]
     let type_specifiers = generic_for
         .type_specifiers()
         .map(|x| x.map(|type_specifier| format_type_specifier(ctx, type_specifier, shape)))
         .collect();
 
-    let in_token = fmt_symbol!(ctx, generic_for.in_token(), " in ", shape);
-    let (formatted_expr_list, mut expr_comments_buf) =
-        format_punctuated_buffer(ctx, generic_for.expressions(), shape, format_expression);
+    // Format the expression list on a single line, and see if it needs expanding
+    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();
+
+    let in_token = match (require_names_multiline, requires_expr_multiline) {
+        (true, true) => fmt_symbol!(ctx, generic_for.in_token(), "in", shape)
+            .update_leading_trivia(FormatTriviaType::Append(vec![
+                create_newline_trivia(ctx),
+                create_indent_trivia(ctx, shape),
+            ])),
+        (true, false) => fmt_symbol!(ctx, generic_for.in_token(), "in ", shape)
+            .update_leading_trivia(FormatTriviaType::Append(vec![
+                create_newline_trivia(ctx),
+                create_indent_trivia(ctx, shape),
+            ])),
+        (false, true) => fmt_symbol!(ctx, generic_for.in_token(), " in", shape),
+        (false, false) => fmt_symbol!(ctx, generic_for.in_token(), " in ", shape),
+    };
 
-    // Create comments buffer and append to end of do token
-    names_comments_buf.append(&mut expr_comments_buf);
-    // Append trailing trivia to the end
-    names_comments_buf.append(&mut trailing_trivia);
+    // Expand the expression list if necessary
+    let expr_list = match requires_expr_multiline {
+        true => {
+            let shape = shape.reset().increment_additional_indent();
+            format_punctuated_multiline(
+                ctx,
+                generic_for.expressions(),
+                shape,
+                format_expression,
+                None,
+            )
+            .update_leading_trivia(FormatTriviaType::Append(vec![
+                create_newline_trivia(ctx),
+                create_indent_trivia(ctx, shape),
+            ]))
+        }
+        false => singleline_expr,
+    };
 
-    let do_token = fmt_symbol!(ctx, generic_for.do_token(), " do", shape)
-        .update_trailing_trivia(FormatTriviaType::Append(names_comments_buf));
+    let do_token = match requires_expr_multiline {
+        true => fmt_symbol!(ctx, generic_for.in_token(), "do", shape).update_leading_trivia(
+            FormatTriviaType::Append(vec![
+                create_newline_trivia(ctx),
+                create_indent_trivia(ctx, shape),
+            ]),
+        ),
+        false => fmt_symbol!(ctx, generic_for.do_token(), " do", shape),
+    }
+    .update_trailing_trivia(FormatTriviaType::Append(trailing_trivia));
 
     let block_shape = shape.reset().increment_block_indent();
     let block = format_block(ctx, generic_for.block(), block_shape);
@@ -117,9 +196,9 @@ pub fn format_generic_for(ctx: &Context, generic_for: &GenericFor, shape: Shape)
     let generic_for = generic_for.with_type_specifiers(type_specifiers);
     generic_for
         .with_for_token(for_token)
-        .with_names(formatted_names)
+        .with_names(names)
         .with_in_token(in_token)
-        .with_expressions(formatted_expr_list)
+        .with_expressions(expr_list)
         .with_do_token(do_token)
         .with_block(block)
         .with_end_token(end_token)

src/formatters/trivia_util.rs

@@ -42,6 +42,12 @@ pub fn trivia_contains_newline<'a>(trivia_vec: impl Iterator<Item = &'a Token>)
     false
 }
 
+/// Determines whether a particular node spans over multiple lines
+pub fn spans_multiple_lines<T: std::fmt::Display>(item: &T) -> bool {
+    let string = format!("{}", item);
+    string.lines().count() > 1
+}
+
 pub fn can_hang_expression(expression: &Expression) -> bool {
     match expression {
         Expression::Parentheses { expression, .. } => can_hang_expression(expression),

tests/inputs/generic-for-multiline.lua

@@ -0,0 +1,14 @@
+local function system(world)
+	for id, model, lasering, transform in world:query(Components.Model, Components.Lasering, Components.Transform, Components.Mothership) do
+	end
+end
+
+local function system(world)
+	for id, model, lasering, transform, id, model, lasering, transform, id, model, lasering, transform, id, model, lasering, transform in world:query(Components.Model, Components.Lasering, Components.Transform, Components.Mothership) do
+	end
+end
+
+local function system(world)
+	for id, model, lasering, transform, id, model, lasering, transform, id, model, lasering, transform, id, model, lasering, transform in world:query(Components.Model, Components.Lasering, Components.Transform, Components.Mothership, Components.Mothership) do
+	end
+end

tests/snapshots/tests__standard@comments-buffer.lua.snap

@@ -19,7 +19,13 @@ local expr_result = 1
 print("text") --a comment
 foo({ bar = baz }) -- comment
 
-for foo, bar in next, value do -- test -- comment
+for
+	foo, -- test
+	bar
+in
+	next, -- comment
+	value
+do
 	print(
 		"test", -- comment
 		"foo"

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

@@ -0,0 +1,64 @@
+---
+source: tests/tests.rs
+expression: format(&contents)
+
+---
+local function system(world)
+	for id, model, lasering, transform in
+		world:query(Components.Model, Components.Lasering, Components.Transform, Components.Mothership)
+	do
+	end
+end
+
+local function system(world)
+	for
+		id,
+		model,
+		lasering,
+		transform,
+		id,
+		model,
+		lasering,
+		transform,
+		id,
+		model,
+		lasering,
+		transform,
+		id,
+		model,
+		lasering,
+		transform
+	in world:query(Components.Model, Components.Lasering, Components.Transform, Components.Mothership) do
+	end
+end
+
+local function system(world)
+	for
+		id,
+		model,
+		lasering,
+		transform,
+		id,
+		model,
+		lasering,
+		transform,
+		id,
+		model,
+		lasering,
+		transform,
+		id,
+		model,
+		lasering,
+		transform
+	in
+		world:query(
+			Components.Model,
+			Components.Lasering,
+			Components.Transform,
+			Components.Mothership,
+			Components.Mothership
+		)
+	do
+	end
+end
+