From 66c699685c85d21419e682e9714d713b753ad62d Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Thu, 20 Jun 2024 12:06:55 +0200
Subject: [PATCH] Always wrap RHS of `where` in braces

---
 src/Runic.jl     |  6 ++++++
 src/runestone.jl | 46 ++++++++++++++++++++++++++++++++++++++++++++--
 test/runtests.jl |  7 +++++++
 3 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/src/Runic.jl b/src/Runic.jl
index 630f202..35f9aec 100644
--- a/src/Runic.jl
+++ b/src/Runic.jl
@@ -43,6 +43,11 @@ function Node(head::JuliaSyntax.SyntaxHead, span::Integer)
     return Node(head, span % UInt32, (), 0 % TagType)
 end
 
+function Node(head::JuliaSyntax.SyntaxHead, kids::Vector{Node})
+    spn = mapreduce(span, +, kids; init = 0)
+    return Node(head, spn % UInt32, kids, 0 % TagType)
+end
+
 # Re-package a GreenNode as a Node
 function Node(node::JuliaSyntax.GreenNode)
     tags = 0 % TagType
@@ -301,6 +306,7 @@ function format_node!(ctx::Context, node::Node)::Union{Node, Nothing, NullNode}
     @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 braces_around_where_rhs(ctx, node)
     @return_something four_space_indent(ctx, node)
     @return_something spaces_in_listlike(ctx, node)
     ctx.call_depth -= 1
diff --git a/src/runestone.jl b/src/runestone.jl
index f17c670..e060a2d 100644
--- a/src/runestone.jl
+++ b/src/runestone.jl
@@ -152,7 +152,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::Node, is_x::F, n_leaves_per_x::Int = 1) where F
+function spaces_around_x(ctx::Context, node::Node, is_x::F, n_leaves_per_x::Int = 1) where {F}
     # TODO: So much boilerplate here...
     @assert !is_leaf(node)
 
@@ -684,7 +684,7 @@ function spaces_around_anonymous_function(ctx::Context, node::Node)
 end
 
 # Opposite of `spaces_around_x`: remove spaces around `x`
-function no_spaces_around_x(ctx::Context, node::Node, is_x::F) where F
+function no_spaces_around_x(ctx::Context, node::Node, is_x::F) where {F}
     @assert !is_leaf(node)
     # TODO: Can't handle NewlineWs here right now
     if any(kind(c) === K"NewlineWs" for c in verified_kids(node))
@@ -1001,6 +1001,48 @@ function for_loop_use_in(ctx::Context, node::Node)
     return node′
 end
 
+function braces_around_where_rhs(ctx::Context, node::Node)
+    if !(kind(node) === K"where" && !is_leaf(node))
+        return nothing
+    end
+    kids = verified_kids(node)
+    kids′ = kids
+    # any_changes = false
+    pos = position(ctx.fmt_io)
+    where_idx = findfirst(x -> is_leaf(x) && kind(x) === K"where", kids)::Int
+    rhs_idx = findnext(!JuliaSyntax.is_whitespace, kids, where_idx + 1)::Int
+    rhs = kids[rhs_idx]
+    if kind(rhs) === K"braces"
+        return nothing
+    end
+    # Wrap the rhs in a braces node
+    kids′ = kids[1:rhs_idx - 1]
+    for i in 1:rhs_idx - 1
+        accept_node!(ctx, kids[i])
+    end
+    opening_brace = Node(JuliaSyntax.SyntaxHead(K"{", 0), 1)
+    closing_brace = Node(JuliaSyntax.SyntaxHead(K"}", 0), 1)
+    rhs′ = Node(
+        JuliaSyntax.SyntaxHead(K"braces", 0),
+        [opening_brace, rhs, closing_brace],
+    )
+    push!(kids′, rhs′)
+    # Write the new node
+    replace_bytes!(ctx, "{", 0)
+    accept_node!(ctx, opening_brace)
+    accept_node!(ctx, rhs)
+    replace_bytes!(ctx, "}", 0)
+    accept_node!(ctx, closing_brace)
+    # Accept any remaining kids
+    for i in rhs_idx + 1:length(kids)
+        accept_node!(ctx, kids[i])
+        push!(kids′, kids[i])
+    end
+    # Reset stream and return
+    seek(ctx.fmt_io, pos)
+    return make_node(node, kids′)
+end
+
 # This function materialized all indentations marked by `insert_delete_mark_newlines`.
 function four_space_indent(ctx::Context, node::Node)
     kind(node) === K"NewlineWs" || return nothing
diff --git a/test/runtests.jl b/test/runtests.jl
index c6f73c8..0b6090e 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -379,6 +379,13 @@ end
     end
 end
 
+@testset "braces around where rhs" begin
+    @test format_string("A where B") == "A where {B}"
+    @test format_string("A where B <: C") == "A where {B <: C}"
+    @test format_string("A where B >: C") == "A where {B >: C}"
+    @test format_string("A where B where C") == "A where {B} where {C}"
+end
+
 @testset "block/hard indentation" begin
     for sp in ("", "  ", "    ", "      ")
         # function-end