Fix general mistransformations (#202)

* Keep parentheses around type assertions

* Fix function types with comments incorrectly formatted

* Fix comment in between equals and expr in assignment leading to mistransformation

* Clean clippy warnings

Author: JohnnyMorganz

CHANGELOG.md

@@ -5,6 +5,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+### Fixed
+- Fixed parentheses around type assertions being classed as unnecessary and removed under the `luau` feature flag.
+- Fixed mistransformation of function type where arguments have comments under the `luau` feature flag. ([#201](https://github.com/JohnnyMorganz/StyLua/issues/201))
+- Fixed comments in an assignment in between the equals token and the expression leading to a mistransformation. ([#200](https://github.com/JohnnyMorganz/StyLua/issues/200))
 
 ## [0.9.1] - 2021-06-17
 ### Added

src/formatters/assignment.rs

@@ -1,10 +1,13 @@
 #[cfg(feature = "luau")]
 use full_moon::ast::types::TypeSpecifier;
-use full_moon::ast::{
-    punctuated::{Pair, Punctuated},
-    Assignment, Expression, LocalAssignment,
+use full_moon::tokenizer::{Token, TokenReference};
+use full_moon::{
+    ast::{
+        punctuated::{Pair, Punctuated},
+        Assignment, Expression, LocalAssignment,
+    },
+    tokenizer::TokenType,
 };
-use full_moon::tokenizer::{Token, TokenKind, TokenReference};
 
 #[cfg(feature = "luau")]
 use crate::formatters::luau::format_type_specifier;
@@ -72,20 +75,22 @@ fn hang_equal_token<'ast>(
     ctx: &Context,
     equal_token: TokenReference<'ast>,
     shape: Shape,
+    indent_first_item: bool,
 ) -> TokenReference<'ast> {
-    let equal_token_trailing_trivia = vec![
-        create_newline_trivia(ctx),
-        create_indent_trivia(ctx, shape.increment_additional_indent()),
-    ]
-    .iter()
-    .chain(
-        // Remove the space that was present after the equal token
-        equal_token
-            .trailing_trivia()
-            .skip_while(|x| x.token_kind() == TokenKind::Whitespace),
-    )
-    .map(|x| x.to_owned())
-    .collect();
+    let mut equal_token_trailing_trivia = vec![create_newline_trivia(ctx)];
+    if indent_first_item {
+        equal_token_trailing_trivia.push(create_indent_trivia(
+            ctx,
+            shape.increment_additional_indent(),
+        ))
+    }
+
+    let equal_token_trailing_trivia = equal_token
+        .trailing_trivia()
+        .filter(|x| trivia_util::trivia_is_comment(x))
+        .flat_map(|x| vec![Token::new(TokenType::spaces(1)), x.to_owned()])
+        .chain(equal_token_trailing_trivia.iter().map(|x| x.to_owned()))
+        .collect();
 
     equal_token.update_trailing_trivia(FormatTriviaType::Replace(equal_token_trailing_trivia))
 }
@@ -102,7 +107,7 @@ fn attempt_assignment_tactics<'ast>(
     // If there is, we should put it on multiple lines
     if expressions.len() > 1 {
         // First try hanging at the equal token, using an infinite width, to see if its enough
-        let hanging_equal_token = hang_equal_token(ctx, equal_token.to_owned(), shape);
+        let hanging_equal_token = hang_equal_token(ctx, equal_token.to_owned(), shape, true);
         let hanging_shape = shape.reset().increment_additional_indent();
         let expr_list = format_punctuated(
             ctx,
@@ -131,6 +136,46 @@ fn attempt_assignment_tactics<'ast>(
         }
     } else {
         // There is only a single element in the list
+        let expression = expressions.iter().next().unwrap();
+
+        // Special case: there is a comment in between the equals and the expression
+        if trivia_util::token_contains_comments(&equal_token)
+            || !trivia_util::expression_leading_comments(expression).is_empty()
+        {
+            // We will hang at the equals token, and then format the expression as necessary
+            let equal_token = hang_equal_token(ctx, equal_token, shape, false);
+
+            let shape = shape.reset().increment_additional_indent();
+
+            // As we know that there is only a single element in the list, we can extract it to work with it
+            let expression = format_expression(ctx, expression, shape);
+
+            // We need to take all the leading trivia from the expr_list
+            let (expression, leading_comments) =
+                trivia_util::take_expression_leading_comments(&expression);
+
+            // Indent each comment and trail them with a newline
+            let leading_comments = leading_comments
+                .iter()
+                .flat_map(|x| {
+                    vec![
+                        create_indent_trivia(ctx, shape),
+                        x.to_owned(),
+                        create_newline_trivia(ctx),
+                    ]
+                })
+                .chain(std::iter::once(create_indent_trivia(ctx, shape)))
+                .collect();
+
+            let expression =
+                expression.update_leading_trivia(FormatTriviaType::Replace(leading_comments));
+
+            // Rebuild expression back into a list
+            let expr_list = std::iter::once(Pair::new(expression, None)).collect();
+
+            return (expr_list, equal_token);
+        }
+
         // Create an example hanging the expression - we need to create a new context so that we don't overwrite it
         let hanging_expr_list = hang_punctuated_list(ctx, expressions, shape);
         let hanging_shape = shape.take_first_line(&strip_trivia(&hanging_expr_list));
@@ -147,7 +192,7 @@ fn attempt_assignment_tactics<'ast>(
             // Normal format is better: but check to see if the formatting is still over budget
             if formatting_shape.over_budget() {
                 // Hang at the equal token
-                let equal_token = hang_equal_token(ctx, equal_token, shape);
+                let equal_token = hang_equal_token(ctx, equal_token, shape, true);
                 // Add the expression list into the indent range, as it will be indented by one
                 let shape = shape.increment_additional_indent();
                 let expr_list = format_punctuated(ctx, expressions, shape, format_expression);
@@ -169,13 +214,14 @@ pub fn format_assignment<'ast>(
     let leading_trivia = vec![create_indent_trivia(ctx, shape)];
     let trailing_trivia = vec![create_newline_trivia(ctx)];
 
-    // Check if the assignment expressions contain comments. If they do, we bail out of determining any tactics
+    // 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 = assignment.expressions().pairs().any(|pair| {
-        pair.punctuation()
-            .map_or(false, |x| trivia_util::token_contains_comments(x))
-            || trivia_util::expression_contains_inline_comments(pair.value())
-    });
+    let contains_comments = trivia_util::token_contains_comments(assignment.equal_token())
+        || assignment.expressions().pairs().any(|pair| {
+            pair.punctuation()
+                .map_or(false, |x| trivia_util::token_contains_comments(x))
+                || trivia_util::expression_contains_inline_comments(pair.value())
+        });
 
     // Firstly attempt to format the assignment onto a single line, using an infinite column width shape
     let mut var_list = try_format_punctuated(
@@ -280,13 +326,16 @@ pub fn format_local_assignment<'ast>(
     if assignment.expressions().is_empty() {
         format_local_no_assignment(ctx, assignment, shape, leading_trivia, trailing_trivia)
     } else {
-        // Check if the assignment expressions contain comments. If they do, we bail out of determining any tactics
+        // Check if the assignment expression or equals token contain comments. If they do, we bail out of determining any tactics
         // and format multiline
-        let contains_comments = assignment.expressions().pairs().any(|pair| {
-            pair.punctuation()
-                .map_or(false, |x| trivia_util::token_contains_comments(x))
-                || trivia_util::expression_contains_inline_comments(pair.value())
-        });
+        let contains_comments = assignment
+            .equal_token()
+            .map_or(false, |x| trivia_util::token_contains_comments(x))
+            || assignment.expressions().pairs().any(|pair| {
+                pair.punctuation()
+                    .map_or(false, |x| trivia_util::token_contains_comments(x))
+                    || trivia_util::expression_contains_inline_comments(pair.value())
+            });
 
         // 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)

src/formatters/expression.rs

@@ -78,7 +78,17 @@ fn check_excess_parentheses(internal_expression: &Expression) -> bool {
         Expression::UnaryOperator { expression, .. } => check_excess_parentheses(expression),
         // Don't bother removing them if there is a binop, as they may be needed. TODO: can we be more intelligent here?
         Expression::BinaryOperator { .. } => false,
-        Expression::Value { value, .. } => {
+        Expression::Value {
+            value,
+            #[cfg(feature = "luau")]
+            type_assertion,
+        } => {
+            // If we have a type assertion, we should always keep parentheses
+            #[cfg(feature = "luau")]
+            if type_assertion.is_some() {
+                return false;
+            }
+
             match &**value {
                 // Internal expression is a function call
                 // We could potentially be culling values, so we should not remove parentheses

src/formatters/luau.rs

@@ -4,15 +4,18 @@ use crate::{
     formatters::{
         expression::{format_expression, format_var},
         general::{
-            format_contained_span, format_end_token, format_symbol, format_token_reference,
+            format_contained_span, format_end_token, format_punctuated,
+            format_punctuated_multiline, format_symbol, format_token_reference,
             try_format_punctuated, EndTokenType,
         },
         table::{create_table_braces, TableType},
         trivia::{
             strip_leading_trivia, strip_trailing_trivia, FormatTriviaType, UpdateLeadingTrivia,
             UpdateTrailingTrivia,
         },
-        trivia_util::{contains_comments, take_type_field_trailing_comments},
+        trivia_util::{
+            contains_comments, take_type_field_trailing_comments, token_trivia_contains_comments,
+        },
     },
     shape::Shape,
 };
@@ -93,11 +96,49 @@ pub fn format_type_info<'ast>(
             arrow,
             return_type,
         } => {
-            let parentheses = format_contained_span(ctx, parentheses, shape);
-            let arguments =
-                try_format_punctuated(ctx, arguments, shape, format_type_argument, None);
+            let (start_parens, end_parens) = parentheses.tokens();
+
+            let force_multiline = token_trivia_contains_comments(start_parens.trailing_trivia())
+                || token_trivia_contains_comments(end_parens.leading_trivia())
+                || contains_comments(arguments);
+
+            let (parentheses, arguments, shape) = if force_multiline {
+                let start_parens = fmt_symbol!(ctx, start_parens, "(", shape)
+                    .update_trailing_trivia(FormatTriviaType::Append(vec![
+                        create_newline_trivia(ctx),
+                        create_indent_trivia(ctx, shape.increment_additional_indent()),
+                    ]));
+                let end_parens =
+                    format_end_token(ctx, end_parens, EndTokenType::ClosingParens, shape)
+                        .update_leading_trivia(FormatTriviaType::Append(vec![
+                            create_newline_trivia(ctx),
+                            create_indent_trivia(ctx, shape),
+                        ]));
+
+                let parentheses = ContainedSpan::new(start_parens, end_parens);
+
+                let arguments = format_punctuated_multiline(
+                    ctx,
+                    arguments,
+                    shape.reset(),
+                    format_type_argument,
+                    Some(1),
+                );
+
+                let shape = shape.reset() + 1; // 1 = ")"
+
+                (parentheses, arguments, shape)
+            } else {
+                let parentheses = format_contained_span(ctx, &parentheses, shape);
+                let arguments = format_punctuated(ctx, arguments, shape + 1, format_type_argument);
+                let shape = shape + (2 + arguments.to_string().len()); // 2 = opening and closing parens
+
+                (parentheses, arguments, shape)
+            };
+
             let arrow = fmt_symbol!(ctx, arrow, " -> ", shape);
             let return_type = Box::new(format_type_info(ctx, return_type, shape));
+
             TypeInfo::Callback {
                 parentheses,
                 arguments,

src/formatters/trivia_util.rs

@@ -1,4 +1,4 @@
-use crate::formatters::trivia::{FormatTriviaType, UpdateTrailingTrivia};
+use crate::formatters::trivia::{FormatTriviaType, UpdateLeadingTrivia, UpdateTrailingTrivia};
 #[cfg(feature = "luau")]
 use full_moon::ast::span::ContainedSpan;
 #[cfg(feature = "luau")]
@@ -383,6 +383,23 @@ pub fn expression_leading_comments<'ast>(expression: &Expression<'ast>) -> Vec<T
         .collect()
 }
 
+pub fn take_expression_leading_comments<'ast>(
+    expression: &Expression<'ast>,
+) -> (Expression<'ast>, Vec<Token<'ast>>) {
+    let trailing_comments = get_expression_leading_trivia(expression)
+        .iter()
+        .filter(|token| trivia_is_comment(token))
+        .map(|x| x.to_owned())
+        .collect();
+
+    (
+        expression.update_leading_trivia(
+            FormatTriviaType::Replace(vec![]), // TODO: Do we need to keep some trivia?
+        ),
+        trailing_comments,
+    )
+}
+
 pub fn take_expression_trailing_comments<'ast>(
     expression: &Expression<'ast>,
 ) -> (Expression<'ast>, Vec<Token<'ast>>) {

tests/inputs-luau/excess-parentheses.lua

@@ -0,0 +1 @@
+local foo = (bar :: any) :: number

tests/inputs-luau/types_comments.lua

@@ -5,4 +5,10 @@ export type IntrospectionNamedTypeRef<
   kind: any, -- deviation: add this type spec later: $PropertyType<T, 'kind'>,
   name: string,
   ofType: T -- TODO: this field is missing
-}
+}
+
+export type ReactScopeQuery = (
+	string, -- type
+	{ [any]: any }, -- props
+	any -- instance
+) -> boolean

tests/inputs/assignment-comments.lua

@@ -0,0 +1,14 @@
+local isValid =
+	-- Allow nil for conditional declaration
+	contextType == nil or
+	(contextType["$$typeof"] == REACT_CONTEXT_TYPE and
+		contextType._context == nil) -- Not a <Context.Consumer>
+
+local isValid = -- Allow nil for conditional declaration
+	foo
+
+local isValid = -- test comment
+	-- Allow nil for conditional declaration
+	contextType == nil or
+	(contextType["$$typeof"] == REACT_CONTEXT_TYPE and
+		contextType._context == nil) -- Not a <Context.Consumer>

tests/snapshots/tests__luau@excess-parentheses.lua.snap

@@ -0,0 +1,7 @@
+---
+source: tests/tests.rs
+expression: format(&contents)
+
+---
+local foo = (bar :: any) :: number
+

tests/snapshots/tests__luau@types_comments.lua.snap

@@ -12,3 +12,9 @@ export type IntrospectionNamedTypeRef<
 	ofType: T, -- TODO: this field is missing
 }
 
+export type ReactScopeQuery = (
+	string, -- type
+	{ [any]: any }, -- props
+	any -- instance
+) -> boolean
+