Fix method call chain formatting with comments

CHANGELOG.md

@@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Added configuration option `space_after_function_names` to specify whether to include a space between a function name and parentheses ([#839](https://github.com/JohnnyMorganz/StyLua/issues/839))
 
+### Fixed
+
+- Fixed formatting of method call chain when there is a comment between the colon token `:` and the function name ([#890](https://github.com/JohnnyMorganz/StyLua/issues/890))
+
 ## [0.20.0] - 2024-01-20
 
 ### Added

src/formatters/expression.rs

@@ -375,6 +375,31 @@ pub fn is_brackets_string(expression: &Expression) -> bool {
     }
 }
 
+pub fn process_dot_name(
+    ctx: &Context,
+    dot: &TokenReference,
+    name: &TokenReference,
+    shape: Shape,
+) -> (TokenReference, TokenReference) {
+    // If there are any comments in between the dot and name,
+    // then taken them out and put them before the dot
+    let (mut dot, mut dot_comments) =
+        take_trailing_comments(&format_token_reference(ctx, dot, shape));
+    let (name, name_comments) = take_leading_comments(&format_token_reference(ctx, name, shape));
+
+    dot_comments.extend(name_comments);
+
+    if !dot_comments.is_empty() {
+        dot = prepend_newline_indent(
+            ctx,
+            &dot.update_leading_trivia(FormatTriviaType::Append(dot_comments)),
+            shape,
+        );
+    }
+
+    (dot, name)
+}
+
 /// Formats an Index Node
 pub fn format_index(ctx: &Context, index: &Index, shape: Shape) -> Index {
     match index {
@@ -436,23 +461,7 @@ pub fn format_index(ctx: &Context, index: &Index, shape: Shape) -> Index {
         }
 
         Index::Dot { dot, name } => {
-            // If there are any comments in between the dot and name,
-            // then taken them out and put them before the dot
-            let (mut dot, mut dot_comments) =
-                take_trailing_comments(&format_token_reference(ctx, dot, shape));
-            let (name, name_comments) =
-                take_leading_comments(&format_token_reference(ctx, name, shape));
-
-            dot_comments.extend(name_comments);
-
-            if !dot_comments.is_empty() {
-                dot = prepend_newline_indent(
-                    ctx,
-                    &dot.update_leading_trivia(FormatTriviaType::Append(dot_comments)),
-                    shape,
-                );
-            }
-
+            let (dot, name) = process_dot_name(ctx, dot, name, shape);
             Index::Dot { dot, name }
         }
         other => panic!("unknown node {:?}", other),

src/formatters/functions.rs

@@ -36,6 +36,8 @@ use crate::{
     CallParenType,
 };
 
+use super::expression::process_dot_name;
+
 /// Formats an Anonymous Function
 /// This doesn't have its own struct, but it is part of Value::Function
 pub fn format_anonymous_function(
@@ -938,6 +940,19 @@ fn should_inline_prefix(ctx: &Context, prefix: &Prefix) -> bool {
         || prefix.len() <= ctx.config().indent_width
 }
 
+fn suffix_contains_comments(suffix: &Suffix) -> bool {
+    match suffix {
+        Suffix::Index(Index::Dot { dot, name }) => {
+            dot.has_trailing_comments(CommentSearch::All)
+                || name.has_leading_comments(CommentSearch::All)
+        }
+        Suffix::Call(Call::MethodCall(method_call)) => method_call
+            .colon_token()
+            .has_trailing_comments(CommentSearch::Single),
+        _ => false,
+    }
+}
+
 /// Formats a FunctionCall node
 pub fn format_function_call(
     ctx: &Context,
@@ -958,7 +973,7 @@ pub fn format_function_call(
             while let Some(suffix) = peekable_suffixes.next() {
                 must_hang = suffix.has_leading_comments(CommentSearch::All)
                 // Check for comment placed inside of suffix
-                || matches!(suffix, Suffix::Index(Index::Dot { dot, name }) if dot.has_trailing_comments(CommentSearch::All) || name.has_leading_comments(CommentSearch::All))
+                || suffix_contains_comments(suffix)
                 // Check for a trailing comment (iff there is still a suffix after this)
                 || (peekable_suffixes.peek().is_some()
                     && suffix.has_trailing_comments(CommentSearch::All));
@@ -1212,15 +1227,15 @@ pub fn format_method_call(
     call_next_node: FunctionCallNextNode,
 ) -> MethodCall {
     let function_call_trivia = vec![create_function_call_trivia(ctx)];
-    let formatted_colon_token = format_token_reference(ctx, method_call.colon_token(), shape);
-    let formatted_name = format_token_reference(ctx, method_call.name(), shape);
-    let shape =
-        shape + (formatted_colon_token.to_string().len() + formatted_name.to_string().len());
+
+    let (colon_token, name) =
+        process_dot_name(ctx, method_call.colon_token(), method_call.name(), shape);
+    let shape = shape + (colon_token.to_string().len() + name.to_string().len());
     let formatted_function_args =
         format_function_args(ctx, method_call.args(), shape, call_next_node)
             .update_leading_trivia(FormatTriviaType::Append(function_call_trivia));
 
-    MethodCall::new(formatted_name, formatted_function_args).with_colon_token(formatted_colon_token)
+    MethodCall::new(name, formatted_function_args).with_colon_token(colon_token)
 }
 
 /// Formats a single Parameter node

tests/inputs/hang-call-chain-comments-3.lua

@@ -0,0 +1,4 @@
+-- https://github.com/JohnnyMorganz/StyLua/issues/890
+
+build(): -- comment
+init():start()

tests/snapshots/tests__standard@hang-call-chain-comments-3.lua.snap

@@ -0,0 +1,11 @@
+---
+source: tests/tests.rs
+expression: format(&contents)
+input_file: tests/inputs/hang-call-chain-comments-3.lua
+---
+-- https://github.com/JohnnyMorganz/StyLua/issues/890
+
+build()
+	-- comment
+	:init()
+	:start()