From 0db3c5484a6dfede1974424f5b2f57fd54ec431b Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Sun, 26 May 2024 12:49:28 +0200
Subject: [PATCH] Fix spaces in comparison chains

---
 src/chisel.jl    |  6 ++++++
 src/runestone.jl | 25 ++++++++++++-------------
 test/runtests.jl | 18 ++++++++++++------
 3 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/src/chisel.jl b/src/chisel.jl
index 3950944..c41654b 100644
--- a/src/chisel.jl
+++ b/src/chisel.jl
@@ -23,3 +23,9 @@ end
 function is_assignment(node::JuliaSyntax.GreenNode)
     return JuliaSyntax.is_prec_assignment(node)
 end
+function is_comparison_leaf(node::JuliaSyntax.GreenNode)
+    return is_leaf(node) && JuliaSyntax.is_prec_comparison(node)
+end
+function is_operator_leaf(node::JuliaSyntax.GreenNode)
+    return is_leaf(node) && JuliaSyntax.is_operator(node)
+end
diff --git a/src/runestone.jl b/src/runestone.jl
index ada2d8d..d10952c 100644
--- a/src/runestone.jl
+++ b/src/runestone.jl
@@ -143,7 +143,7 @@ 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)
+function spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, is_x::F) where F
     # TODO: So much boilerplate here...
     @assert JuliaSyntax.haschildren(node)
     # TODO: Can't handle NewlineWs here right now
@@ -239,13 +239,13 @@ function spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, x::JuliaSynt
                 accept_node!(ctx, child)
                 looking_for_whitespace = JuliaSyntax.kind(last_leaf(child)) !== K"Whitespace"
                 if looking_for_x
-                    @assert JuliaSyntax.kind(child) === x
+                    @assert is_x(child)::Bool
                 end
                 looking_for_x = !looking_for_x
             end
         else # !expect_ws
             if looking_for_x
-                @assert JuliaSyntax.kind(child) === x
+                @assert is_x(child)::Bool
             end
             @assert JuliaSyntax.kind(child) !== K"Whitespace" # This would be weird, I think?
             any_changes && push!(children′, child)
@@ -266,23 +266,22 @@ function spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, x::JuliaSynt
     end
 end
 
+# This pass handles spaces around infix operator calls and comparison chains
 function spaces_around_operators(ctx::Context, node::JuliaSyntax.GreenNode)
-    if !(JuliaSyntax.is_infix_op_call(node))
+    if !(
+        JuliaSyntax.is_infix_op_call(node) ||
+        (JuliaSyntax.kind(node) === K"comparison" && !JuliaSyntax.is_trivia(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)
+    @assert JuliaSyntax.kind(node) in KSet"call comparison"
+    is_x = x -> is_operator_leaf(x) || is_comparison_leaf(x)
+    return spaces_around_x(ctx, node, is_x)
 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))
+    return spaces_around_x(ctx, node, is_assignment)
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index a2f8a64..d612e3d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -99,18 +99,24 @@ end
         @test format_string("  a$(op) b  ") == "  a $(op) b  "
         @test format_string("x=a$(op) b  ") == "x = a $(op) b  "
         @test format_string("a$(op)   b") == "a $(op) b"
-        if op in ("==", "!=", "===", "!==", "<", "<=")
-            # TODO: Broken when chaining these operators
-            @test_broken format_string("a$(op)   b $(op) x") == "a $(op) b $(op) x"
-        else
-            @test format_string("a$(op)   b $(op) x") == "a $(op) b $(op) x"
-        end
+        @test format_string("a$(op)   b $(op) x") == "a $(op) b $(op) x"
         @test format_string("a$(op)   b  *  x") == "a $(op) b * x"
         @test format_string("a$(op)( b *  x)") == "a $(op) ( b * x)"
         @test format_string("sin(π)$(op)cos(pi)") == "sin(π) $(op) cos(pi)"
     end
 end
 
+@testset "whitespace in comparison chains" begin
+    for sp in ("", " ", "  ")
+        @test format_string("a$(sp)==$(sp)b") == "a == b"
+        @test format_string("a$(sp)==$(sp)b$(sp)==$(sp)c") == "a == b == c"
+        @test format_string("a$(sp)<=$(sp)b$(sp)==$(sp)c") == "a <= b == c"
+        @test format_string("a$(sp)<=$(sp)b$(sp)>=$(sp)c") == "a <= b >= c"
+        @test format_string("a$(sp)<$(sp)b$(sp)>=$(sp)c") == "a < b >= c"
+        @test format_string("a$(sp)<$(sp)b$(sp)<$(sp)c") == "a < b < c"
+    end
+end
+
 @testset "whitespace around assignments" begin
     # Regular assignments and dot-assignments
     for a in ("=", "+=", "-=", ".=", ".+=", ".-=")