From a1ee251d980d53a48ee0a606edd63d290f78292c Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Mon, 27 May 2024 12:45:34 +0200
Subject: [PATCH] Fix spaces around operators when nesting calls

The next leaf node might not be a direct grand child when nesting
operators. This patches introduces a function `replace_first_leaf` that
replaces the correct leaf.
---
 src/chisels.jl   | 12 ++++++++++++
 src/runestone.jl | 11 +++--------
 test/runtests.jl |  5 +++++
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/src/chisels.jl b/src/chisels.jl
index f91ffa4..367b0a7 100644
--- a/src/chisels.jl
+++ b/src/chisels.jl
@@ -44,6 +44,18 @@ function first_leaf(node::JuliaSyntax.GreenNode)
     end
 end
 
+function replace_first_leaf(node::JuliaSyntax.GreenNode, child′::JuliaSyntax.GreenNode)
+    if is_leaf(node)
+        return child′
+    else
+        children′ = copy(JuliaSyntax.children(node)::AbstractVector)
+        children′[1] = replace_first_leaf(children′[1], child′)
+        @assert length(children′) > 0
+        span′ = mapreduce(JuliaSyntax.span, +, children′; init = 0)
+        return JuliaSyntax.GreenNode(JuliaSyntax.head(node), span′, children′)
+    end
+end
+
 function last_leaf(node::JuliaSyntax.GreenNode)
     if is_leaf(node)
         return node
diff --git a/src/runestone.jl b/src/runestone.jl
index 3aee271..33d4674 100644
--- a/src/runestone.jl
+++ b/src/runestone.jl
@@ -209,18 +209,13 @@ function spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, is_x::F) whe
                     any_changes && push!(children′, child)
                 else
                     # Replace the whitespace node of the child
-                    grand_children = JuliaSyntax.children(child)[2:end]
-                    pushfirst!(grand_children, ws)
-                    span′ = mapreduce(JuliaSyntax.span, +, grand_children; init = 0)
-                    @assert span′ == JuliaSyntax.span(child) - JuliaSyntax.span(child_ws) + 1
-                    bytes_to_skip = JuliaSyntax.span(child) - span′
+                    child′ = replace_first_leaf(child, ws)
+                    @assert JuliaSyntax.span(child′) == JuliaSyntax.span(child) - JuliaSyntax.span(child_ws) + 1
+                    bytes_to_skip = JuliaSyntax.span(child) - JuliaSyntax.span(child′)
                     @assert bytes_to_skip > 0
                     remaining_bytes_inclusive =
                         @view original_bytes[(span_sum + 1 + bytes_to_skip - JuliaSyntax.span(child)):end]
                     write_and_reset(ctx, remaining_bytes_inclusive)
-                    child′ = JuliaSyntax.GreenNode(
-                        JuliaSyntax.head(child), span′, grand_children,
-                    )
                     any_changes = true
                     if children′ === children
                         children′ = children[1:i - 1]
diff --git a/test/runtests.jl b/test/runtests.jl
index b840c2c..ce82b5b 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -130,6 +130,11 @@ end
             "$(sp)(1 + 2):(1 + 3):(1 + 4)$(sp)"
         # a^b
         @test format_string("$(sp)a$(sp)^$(sp)b$(sp)") == "$(sp)a^b$(sp)"
+        # Edgecase when formatting whitespace in the next leaf, when the next leaf is a
+        # grand child or even younger. Note that this test depends a bit on where
+        # JuliaSyntax.jl decides to place the K"Whitespace" node.
+        @test format_string("$(sp)a$(sp)+$(sp)b$(sp)*$(sp)c$(sp)/$(sp)d$(sp)") ==
+            "$(sp)a + b * c / d$(sp)"
     end
 end