From a0b7eb2c2c4bdc53c3847080755a0b81bcd9d3cf Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Wed, 29 May 2024 09:46:07 +0200
Subject: [PATCH] Add utility function to extract children::Vector in a type
 stable way.

---
 src/Runic.jl     |  6 ++----
 src/chisels.jl   | 19 ++++++++++++++-----
 src/runestone.jl |  4 ++--
 test/runtests.jl | 12 ++++++++++--
 4 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/src/Runic.jl b/src/Runic.jl
index 93a7acf..1356fe6 100644
--- a/src/Runic.jl
+++ b/src/Runic.jl
@@ -115,11 +115,9 @@ function format_node_with_children!(ctx::Context, node::JuliaSyntax.GreenNode)
     span_sum = 0
 
     # The new node parts. `children′` aliases `children` and only copied below if any of the
-    # nodes change ("copy-on-write"). Since we return directly if the node don't have
-    # children we can assert that `JuliaSyntax.children` returns a vector (and not the
-    # empty tuple) to help the compiler.
+    # nodes change ("copy-on-write").
     head′ = JuliaSyntax.head(node)
-    children = JuliaSyntax.children(node)::AbstractVector
+    children = verified_children(node)
     children′ = children
     any_child_changed = false
 
diff --git a/src/chisels.jl b/src/chisels.jl
index c6ca07f..9da101b 100644
--- a/src/chisels.jl
+++ b/src/chisels.jl
@@ -45,15 +45,23 @@ function first_leaf(node::JuliaSyntax.GreenNode)
     if is_leaf(node)
         return node
     else
-        return first_leaf(first(JuliaSyntax.children(node)::AbstractVector))
+        return first_leaf(first(verified_children(node)))
     end
 end
 
+# This function exist so that we can type-assert the return value to narrow it down from
+# `Union{Tuple{}, Vector{JuliaSyntax.GreenNode}}` to `Vector{JuliaSyntax.GreenNode}`. Must
+# only be called after verifying that the node has children.
+function verified_children(node::JuliaSyntax.GreenNode)
+    @assert JuliaSyntax.haschildren(node)
+    return JuliaSyntax.children(node)::AbstractVector
+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′ = copy(verified_children(node))
         children′[1] = replace_first_leaf(children′[1], child′)
         @assert length(children′) > 0
         span′ = mapreduce(JuliaSyntax.span, +, children′; init = 0)
@@ -65,12 +73,13 @@ function last_leaf(node::JuliaSyntax.GreenNode)
     if is_leaf(node)
         return node
     else
-        return last_leaf(last(JuliaSyntax.children(node)::AbstractVector))
+        return last_leaf(last(verified_children(node)))
     end
 end
 
 function is_assignment(node::JuliaSyntax.GreenNode)
     return JuliaSyntax.is_prec_assignment(node)
+    return !is_leaf(node) && JuliaSyntax.is_prec_assignment(node)
 end
 
 # Just like `JuliaSyntax.is_infix_op_call`, but also check that the node is K"call"
@@ -81,7 +90,7 @@ end
 
 function infix_op_call_op(node::JuliaSyntax.GreenNode)
     @assert is_infix_op_call(node)
-    children = JuliaSyntax.children(node)::AbstractVector
+    children = verified_children(node)
     first_operand_index = findfirst(!JuliaSyntax.is_whitespace, children)
     op_index = findnext(JuliaSyntax.is_operator, children, first_operand_index + 1)
     return children[op_index]
@@ -97,7 +106,7 @@ end
 
 function first_non_whitespace_child(node::JuliaSyntax.GreenNode)
     @assert !is_leaf(node)
-    children = JuliaSyntax.children(node)::AbstractVector
+    children = verified_children(node)
     idx = findfirst(!JuliaSyntax.is_whitespace, children)::Int
     return children[idx]
 end
diff --git a/src/runestone.jl b/src/runestone.jl
index caafcd3..1bf99da 100644
--- a/src/runestone.jl
+++ b/src/runestone.jl
@@ -156,7 +156,7 @@ function spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, is_x::F) whe
         return nothing
     end
 
-    children = JuliaSyntax.children(node)::AbstractVector
+    children = verified_children(node)
     children′ = children
     any_changes = false
     original_bytes = node_bytes(ctx, node)
@@ -299,7 +299,7 @@ function no_spaces_around_x(ctx::Context, node::JuliaSyntax.GreenNode, is_x::F)
         return nothing
     end
 
-    children = JuliaSyntax.children(node)::AbstractVector
+    children = verified_children(node)
     children′ = children
     any_changes = false
     original_bytes = node_bytes(ctx, node)
diff --git a/test/runtests.jl b/test/runtests.jl
index fe4ae7d..1b78f9b 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,9 +1,17 @@
 # SPDX-License-Identifier: MIT
 
 using Runic:
-    format_string
+    Runic, format_string
 using Test:
-    @test, @testset, @test_broken
+    @test, @testset, @test_broken, @inferred
+using JuliaSyntax:
+    JuliaSyntax
+
+@testset "Chisels" begin
+    # Type stability of verified_children
+    node = JuliaSyntax.parseall(JuliaSyntax.GreenNode, "a = 1 + b\n")
+    @test typeof(@inferred Runic.verified_children(node)) <: Vector{<:JuliaSyntax.GreenNode}
+end
 
 @testset "Trailing whitespace" begin
     io = IOBuffer()