Consistent spacing around keywords

src/Runic.jl

@@ -298,6 +298,7 @@ function format_node!(ctx::Context, node::Node)::Union{Node, Nothing, NullNode}
     @return_something spaces_around_operators(ctx, node)
     @return_something spaces_around_assignments(ctx, node)
     @return_something spaces_around_anonymous_function(ctx, node)
+    @return_something spaces_around_keywords(ctx, node)
     @return_something no_spaces_around_colon_etc(ctx, node)
     @return_something for_loop_use_in(ctx, node)
     @return_something four_space_indent(ctx, node)

src/runestone.jl

@@ -751,6 +751,144 @@ function no_spaces_around_colon_etc(ctx::Context, node::Node)
     return no_spaces_around_x(ctx, node, is_x)
 end
 
+# Single space around keywords:
+# Both sides of: `where`, `do` (if followed by arguments)
+# Right hand side of: `mutable`, `struct`, `abstract`, `primitive`, `type`, `function`,
+# `if`, `elseif`, `catch` (if followed by variable)
+function spaces_around_keywords(ctx::Context, node::Node)
+    is_leaf(node) && return nothing
+    keyword_set = KSet"where do mutable struct abstract primitive type function if elseif catch"
+    if !(kind(node) in keyword_set)
+        return nothing
+    end
+    kids = verified_kids(node)
+    kids′ = kids
+    any_changes = false
+    pos = position(ctx.fmt_io)
+    ws = Node(JuliaSyntax.SyntaxHead(K"Whitespace", JuliaSyntax.TRIVIA_FLAG), 1)
+
+    peek_kinds = KSet"where do"
+    state = kind(node) in peek_kinds ? (:peeking_for_keyword) : (:looking_for_keyword)
+    keep_looking_for_keywords = false
+    space_after = true
+
+    for i in eachindex(kids)
+        kid = kids[i]
+        if state === :peeking_for_keyword
+            nkid = kids[i + 1]
+            if kind(nkid) in peek_kinds
+                state = :looking_for_space
+                keep_looking_for_keywords = true
+                space_after = false
+            else
+                accept_node!(ctx, kid)
+                any_changes && push!(kids′, kid)
+                continue
+            end
+        end
+        if state === :looking_for_keyword
+            if kind(kid) in keyword_set
+                accept_node!(ctx, kid)
+                any_changes && push!(kids′, kid)
+                if kind(kid) in KSet"mutable abstract primitive"
+                    # These keywords are always followed by another keyword
+                    keep_looking_for_keywords = true
+                end
+                state = :looking_for_space
+                # `do` should only be followed by space if the argument-tuple is non-empty
+                if kind(node) === K"do"
+                    nkid = kids[i + 1]
+                    @assert kind(nkid) === K"tuple"
+                    if !any(!JuliaSyntax.is_whitespace, verified_kids(nkid))
+                        state = :closing
+                    end
+                end
+                # `catch` should only be followed by space if the error is caught in a var
+                if kind(node) === K"catch"
+                    nkid = kids[i + 1]
+                    if kind(nkid) === K"false" && span(nkid) == 0
+                        state = :closing
+                    end
+                end
+            else
+                accept_node!(ctx, kid)
+                any_changes && push!(kids′, kid)
+            end
+        elseif state === :looking_for_space
+            if kind(kid) === K"Whitespace" && span(kid) == 1
+                # TODO: Include NewlineWs here?
+                accept_node!(ctx, kid)
+                any_changes && push!(kids′, kid)
+            elseif kind(kid) === K"Whitespace"
+                # Replace with single space.
+                any_changes = true
+                if kids′ === kids
+                    kids′ = kids[1:i - 1]
+                end
+                replace_bytes!(ctx, " ", span(kid))
+                push!(kids′, ws)
+                accept_node!(ctx, ws)
+            elseif space_after && kind(first_leaf(kid)) === K"Whitespace"
+                kid_ws = first_leaf(kid)
+                if span(kid_ws) == 1
+                    accept_node!(ctx, kid)
+                    any_changes && push!(kids′, kid)
+                else
+                    kid′ = replace_first_leaf(kid, ws)
+                    @assert span(kid′) == span(kid) - span(kid_ws) + 1
+                    replace_bytes!(ctx, " ", span(kid_ws))
+                    accept_node!(ctx, kid′)
+                    any_changes = true
+                    if kids′ === kids
+                        kids′ = kids[1:i - 1]
+                    end
+                    push!(kids′, kid′)
+                end
+            elseif !space_after && kind(last_leaf(kid)) === K"Whitespace"
+                @assert false # Unreachable?
+            else
+                # Reachable in e.g. `T where{T}`, insert space
+                @assert kind(node) === K"where"
+                any_changes = true
+                if kids′ === kids
+                    kids′ = kids[1:i - 1]
+                end
+                # Insert the space before/after the kid depending on whether we are looking
+                # for a space before or after a keyword
+                if !space_after
+                    push!(kids′, kid)
+                    accept_node!(ctx, kid)
+                end
+                replace_bytes!(ctx, " ", 0)
+                push!(kids′, ws)
+                accept_node!(ctx, ws)
+                if space_after
+                    push!(kids′, kid)
+                    accept_node!(ctx, kid)
+                end
+            end
+            state = keep_looking_for_keywords ? (:looking_for_keyword) : (:closing)
+            keep_looking_for_keywords = false
+            space_after = true
+        else
+            @assert state === :closing
+            accept_node!(ctx, kid)
+            any_changes && push!(kids′, kid)
+        end
+    end
+
+    # Reset stream
+    seek(ctx.fmt_io, pos)
+    # Return
+    if any_changes
+        # Construct the new node
+        node′ = make_node(node, kids′)
+        return node′
+    else
+        return nothing
+    end
+end
+
 # Replace the K"=" operator with `in`
 function replace_with_in(ctx::Context, node::Node)
     @assert kind(node) === K"=" && !is_leaf(node) && meta_nargs(node) == 3

test/runtests.jl

@@ -339,6 +339,26 @@ end
     @test format_string("a >:  T   >:    S") == "a >: T >: S"
 end
 
+@testset "spaces around keywords" begin
+    for sp in (" ", "  ")
+        @test format_string("struct$(sp)A end") == "struct A end"
+        @test format_string("mutable$(sp)struct$(sp)A end") == "mutable struct A end"
+        @test format_string("abstract$(sp)type$(sp)A end") == "abstract type A end"
+        @test format_string("primitive$(sp)type$(sp)A 64 end") == "primitive type A 64 end"
+        @test format_string("function$(sp)A() end") == "function A() end"
+        @test format_string("if$(sp)a\nelseif$(sp)b\nend") == "if a\nelseif b\nend"
+        @test format_string("if$(sp)a && b\nelseif$(sp)c || d\nend") == "if a && b\nelseif c || d\nend"
+        @test format_string("try\nerror()\ncatch$(sp)e\nend") == "try\n    error()\ncatch e\nend"
+        @test format_string("A$(sp)where$(sp){T}") == "A where {T}"
+        @test format_string("A$(sp)where$(sp){T}$(sp)where$(sp){S}") == "A where {T} where {S}"
+        @test format_string("f()$(sp)do$(sp)x\ny\nend") == "f() do x\n    y\nend"
+        @test format_string("f()$(sp)do\ny\nend") == "f() do\n    y\nend"
+    end
+    @test format_string("try\nerror()\ncatch\nend") == "try\n    error()\ncatch\nend"
+    @test format_string("A where{T}") == "A where {T}"
+    @test format_string("A{T}where{T}") == "A{T} where {T}"
+end
+
 @testset "replace ∈ and = with in in for loops and generators" begin
     for sp in ("", " ", "  "), op in ("∈", "=", "in")
         op == "in" && sp == "" && continue