Annotate something along with syntax tree

annotate.jl

include("../test/jsjl_utils.jl")

function traverse(@specialize(callback), st::JL.SyntaxTree)
    stack = JL.SyntaxList(st)
    push!(stack, st)
    _traverse!(callback, stack)
end
traverse(@specialize(callback), sn::JL.SyntaxNode) = _traverse!(callback, JS.SyntaxNode[sn])

function _traverse!(@specialize(callback), stack)
    while !isempty(stack)
        x = pop!(stack)
        callback(x)
        if JS.numchildren(x) === 0
            continue
        end
        for i = JS.numchildren(x):-1:1
            push!(stack, x[i])
        end
    end
end

cyan(s) = "\e[36m::$s\e[39m"

# ps = parsedstream("""
# function foo(x::T)
#     sin(x)
# end
# """)
# st0 = jlparse(ps)[1]
function annotate(x::Union{JS.SyntaxNode,JL.SyntaxTree}, ps::JS.ParseStream)
    annotations = Pair{UInt,Vector{UInt8}}[]
    traverse(x) do x
        if JS.kind(x) == JS.K"call"
            push!(annotations, JS.last_byte(x) => Vector{UInt8}(cyan("Call")))
        elseif JS.kind(x) == JS.K"Identifier"
            push!(annotations, JS.last_byte(x) => Vector{UInt8}(cyan("X")))
        end
    end
    textbuf = copy(ps.textbuf)
    i = 0
    sort!(annotations; by=first)
    for (b, annotation) in annotations
        idx = b + i
        splice!(textbuf, idx+1:idx, annotation)
        i += length(annotation)
    end
    return String(textbuf)
end