Format spaces around assignment

This patch generalizes the `spaces_around_operators` pass to
`spaces_around_x` and uses it for both operators and assignment.
The new `spaces_around_assignments` matches a lot of things:
 - regular assignment (`a=b` -> `a = b`, `a+=b` -> `a += b`,
   `a=+b` -> `a = + b`)
 - dotted assignment (`a.=b` -> `a .= b`, `a.+=b` -> `a .+= b`,
   `a.=+b` -> `a .= + b`)
 - keywords in function definitions and at callsites
   (`f(x=1) = x` -> f(x = 1)`, `g(x=1)` -> `g(x = 1)`).
 - named tuples (`(x=1, y=2)` -> `(x = 1, y = 2)`)
 - etc...

src/Runic.jl

@@ -194,6 +194,7 @@ function format_node!(ctx::Context, node::JuliaSyntax.GreenNode)::Union{JuliaSyn
     @return_something format_oct_literals(ctx, node)
     @return_something format_float_literals(ctx, node)
     @return_something spaces_around_operators(ctx, node)
+    @return_something spaces_around_assignments(ctx, node)
 
     # If the node is unchanged at this point, just keep going.
 

src/chisel.jl

@@ -19,3 +19,7 @@ function last_leaf(node::JuliaSyntax.GreenNode)
         return last_leaf(last(JuliaSyntax.children(node)::AbstractVector))
     end
 end
+
+function is_assignment(node::JuliaSyntax.GreenNode)
+    return JuliaSyntax.is_prec_assignment(node)
+end

src/runestone.jl

@@ -140,12 +140,12 @@ function format_float_literals(ctx::Context, node::JuliaSyntax.GreenNode)
     return node′
 end
 
+# Insert space around `x`, where `x` can be operators, assignments, etc. with the pattern:
+# `<something><space><x><space><something>`, for example the spaces around `+` and `=` in
+# `a = x + y`.
+function spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, x::JuliaSyntax.Kind)
     # TODO: So much boilerplate here...
-function spaces_around_operators(ctx::Context, node::JuliaSyntax.GreenNode)
-    JuliaSyntax.is_infix_op_call(node) || return nothing
-    @assert JuliaSyntax.kind(node) === K"call"
     @assert JuliaSyntax.haschildren(node)
-
     # TODO: Can't handle NewlineWs here right now
     if any(JuliaSyntax.kind(c) === K"NewlineWs" for c in JuliaSyntax.children(node))
         return nothing
@@ -163,10 +163,10 @@ function spaces_around_operators(ctx::Context, node::JuliaSyntax.GreenNode)
 
     # Toggle for whether we are currently looking for whitespace or not
     looking_for_whitespace = false
+    looking_for_x = false
 
     for (i, child) in pairs(children)
         span_sum += JuliaSyntax.span(child)
-        ckind = JuliaSyntax.kind(child)
         if i == 1 && JuliaSyntax.kind(child) === K"Whitespace"
             # If the first child is whitespace it will be accepted as is even if the span is
             # larger than one since we don't look behind. The whitespace pass for the parent
@@ -238,12 +238,20 @@ function spaces_around_operators(ctx::Context, node::JuliaSyntax.GreenNode)
                 write_and_reset(ctx, remaining_bytes_inclusive)
                 accept_node!(ctx, child)
                 looking_for_whitespace = JuliaSyntax.kind(last_leaf(child)) !== K"Whitespace"
+                if looking_for_x
+                    @assert JuliaSyntax.kind(child) === x
+                end
+                looking_for_x = !looking_for_x
             end
         else # !expect_ws
+            if looking_for_x
+                @assert JuliaSyntax.kind(child) === x
+            end
             @assert JuliaSyntax.kind(child) !== K"Whitespace" # This would be weird, I think?
             any_changes && push!(children′, child)
             accept_node!(ctx, child)
             looking_for_whitespace = JuliaSyntax.kind(last_leaf(child)) !== K"Whitespace"
+            looking_for_x = !looking_for_x
         end
     end
     # Reset stream
@@ -257,3 +265,24 @@ function spaces_around_operators(ctx::Context, node::JuliaSyntax.GreenNode)
         return nothing
     end
 end
+
+function spaces_around_operators(ctx::Context, node::JuliaSyntax.GreenNode)
+    if !(JuliaSyntax.is_infix_op_call(node))
+        return nothing
+    end
+    @assert JuliaSyntax.kind(node) === K"call"
+    # Find the specific operator
+    children = JuliaSyntax.children(node)::AbstractVector
+    start_idx = JuliaSyntax.is_whitespace(first_leaf(node)) ? 3 : 2
+    i = findnext(!JuliaSyntax.is_whitespace, children, start_idx)::Int
+    op = JuliaSyntax.kind(children[i])
+    @assert JuliaSyntax.is_operator(op)
+    return spaces_around_x(ctx, node, op)
+end
+
+function spaces_around_assignments(ctx::Context, node::JuliaSyntax.GreenNode)
+    if !(is_assignment(node) && !JuliaSyntax.is_trivia(node))
+        return nothing
+    end
+    return spaces_around_x(ctx, node, JuliaSyntax.kind(node))
+end

test/runtests.jl

@@ -110,3 +110,24 @@ end
         @test format_string("sin(π)$(op)cos(pi)") == "sin(π) $(op) cos(pi)"
     end
 end
+
+@testset "whitespace around assignments" begin
+    # Regular assignments and dot-assignments
+    for a in ("=", "+=", "-=", ".=", ".+=", ".-=")
+        @test format_string("a$(a)b") == "a $(a) b"
+        @test format_string("a $(a)b") == "a $(a) b"
+        @test format_string("a$(a) b") == "a $(a) b"
+        @test format_string("  a$(a) b") == "  a $(a) b"
+        @test format_string("  a$(a) b  ") == "  a $(a) b  "
+        @test format_string("a$(a)   b") == "a $(a) b"
+        @test format_string("a$(a)   b  *  x") == "a $(a) b * x"
+        @test format_string("a$(a)( b *  x)") == "a $(a) ( b * x)"
+    end
+    # Chained assignments
+    @test format_string("x=a= b  ") == "x = a = b  "
+    @test format_string("a=   b = x") == "a = b = x"
+    # Check the common footgun of permuting the operator and =
+    @test format_string("a =+ c") == "a = + c"
+    # Short form function definitions
+    @test format_string("sin(π)=cos(pi)") == "sin(π) = cos(pi)"
+end