From 3ee9531e4d81fa79237b0e75d26463bd9e3b4f2b Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Sat, 3 Aug 2024 14:43:41 +0200
Subject: [PATCH] Fix function indent with unusual infix definitions

---
 src/chisels.jl   | 6 ++++--
 src/runestone.jl | 4 ++--
 test/runtests.jl | 3 +++
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/chisels.jl b/src/chisels.jl
index 5d8aace..9337ccb 100644
--- a/src/chisels.jl
+++ b/src/chisels.jl
@@ -316,7 +316,7 @@ end
 function unwrap_to_call_or_tuple(x)
     is_leaf(x) && return nothing
     @assert !is_leaf(x)
-    if kind(x) in KSet"call tuple"
+    if kind(x) in KSet"call tuple parens"
         return x
     end
     xkids = verified_kids(x)
@@ -324,6 +324,8 @@ function unwrap_to_call_or_tuple(x)
     return unwrap_to_call_or_tuple(xkids[xi])
 end
 
+# TODO: This should be reworked to be more specific, in particular K"parens" is maybe not
+# correct (found in e.g. `function(x * b)\n\nend`).
 function is_longform_anon_function(node::Node)
     is_leaf(node) && return false
     kind(node) === K"function" || return false
@@ -336,7 +338,7 @@ function is_longform_anon_function(node::Node)
     if maybe_tuple === nothing
         return false
     else
-        return kind(maybe_tuple) === K"tuple"
+        return kind(maybe_tuple) in KSet"tuple parens"
     end
 end
 
diff --git a/src/runestone.jl b/src/runestone.jl
index 8cb0e09..3c79620 100644
--- a/src/runestone.jl
+++ b/src/runestone.jl
@@ -1760,7 +1760,7 @@ function indent_function_or_macro(ctx::Context, node::Node)
     sig_node = kids[sig_idx]
     # Identifier for regular names but "not function call" for empty functions with Unicode
     # symbols??
-    if kind(sig_node) === K"Identifier" || !(kind(sig_node) in KSet"call where :: tuple")
+    if kind(sig_node) === K"Identifier" || !(kind(sig_node) in KSet"call where :: tuple parens")
         # Empty function definition like `function f end`.
         # TODO: Make sure the spaces around are correct
         end_idx = findnext(x -> kind(x) === K"end", kids, sig_idx + 1)::Int
@@ -1773,7 +1773,7 @@ function indent_function_or_macro(ctx::Context, node::Node)
         return any_kid_changed ? node : nothing
     end
     # K"tuple" when this is an anonymous function
-    @assert !is_leaf(sig_node) && kind(sig_node) in KSet"call where :: tuple"
+    @assert !is_leaf(sig_node) && kind(sig_node) in KSet"call where :: tuple parens"
     # Fourth node is the function/macro body block.
     block_idx = sig_idx + 1
     block_node′ = indent_block(ctx, kids[block_idx])
diff --git a/test/runtests.jl b/test/runtests.jl
index 9c6a615..fdd0882 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -717,6 +717,9 @@ end
         # Functors
         @test format_string("function$(sp)(a::A)(b)\nx\nend") ==
             "function (a::A)(b)\n    x\nend"
+        # TODO: Spaces after function keyword isn't removed.
+        @test format_string("function$(sp)(a * b)\nreturn\nend") ==
+            "function$(sp)(a * b)\n    return\nend"
         # Multiline strings inside lists
         for trip in ("\"\"\"", "```")
             @test format_string("println(io, $(trip)\n$(sp)a\n$(sp)\n$(sp)b\n$(sp)$(trip))") ==