Construct Expr directly from ParseStream (#268)

c42f · web-flow · commit 5bb321a66ec3 · 2023-05-07T15:40:12.000+10:00
Generalize `build_tree` so that we can more easily construct tree types
other than `GreenNode`.  Use this to construct `Expr` directly from
`ParseStream` rather than constructing both GreenNode and SyntaxNode
along the way.

Fix a bunch of type instabilities in the `Expr` conversion code along the
way.

With these changes, parsing all of Base to `Expr` is sped up by about 35%
overall and allocations reduced by around 50%. (Parsing to `Expr` is
now comparable with parsing to `SyntaxNode`.)
diff --git a/src/expr.jl b/src/expr.jl
diff --git a/src/green_tree.jl b/src/green_tree.jl
@@ -43,25 +43,10 @@ struct GreenNode{Head}
     args::Union{Tuple{},Vector{GreenNode{Head}}}
 end
 
-function GreenNode{Head}(head::Head, span::Integer) where {Head}
-    GreenNode{Head}(head, span, ())
+function GreenNode(head::Head, span::Integer, args) where {Head}
+    GreenNode{Head}(head, span, args)
 end
 
-function GreenNode(head::Head, span::Integer) where {Head}
-    GreenNode{Head}(head, span, ())
-end
-
-function GreenNode(head::Head, args) where {Head}
-    children = collect(GreenNode{Head}, args)
-    span = isempty(children) ? 0 : sum(x.span for x in children)
-    GreenNode{Head}(head, span, children)
-end
-
-function GreenNode(head::Head, args::GreenNode{Head}...) where {Head}
-    GreenNode{Head}(head, GreenNode{Head}[args...])
-end
-
-
 # Accessors / predicates
 haschildren(node::GreenNode) = !(node.args isa Tuple{})
 children(node::GreenNode)    = node.args
@@ -115,3 +100,11 @@ function Base.show(io::IO, ::MIME"text/plain", node::GreenNode, str::AbstractStr
     _show_green_node(io, node, "", 1, str, show_trivia)
 end
 
+function build_tree(::Type{GreenNode}, stream::ParseStream; kws...)
+    build_tree(GreenNode{SyntaxHead}, stream; kws...) do h, srcrange, cs
+        span = length(srcrange)
+        isnothing(cs) ? GreenNode(h, span, ()) :
+                        GreenNode(h, span, collect(GreenNode{SyntaxHead}, cs))
+    end
+end
+
diff --git a/src/hooks.jl b/src/hooks.jl
@@ -195,7 +195,7 @@ function _core_parser_hook(code, filename::String, lineno::Int, offset::Int, opt
             #
             ex = build_tree(Expr, stream; filename=filename,
                             wrap_toplevel_as_kind=K"None", first_line=lineno)
-            if Meta.isexpr(ex, :None)
+            if @isexpr(ex, :None)
                 # The None wrapping is only to give somewhere for trivia to be
                 # attached; unwrap!
                 ex = only(ex.args)
@@ -286,7 +286,7 @@ function _fl_parse_hook(code, filename, lineno, offset, options)
     else
         if options === :all
             ex = Base.parse_input_line(String(code), filename=filename, depwarn=false)
-            if !Meta.isexpr(ex, :toplevel)
+            if !@isexpr(ex, :toplevel)
                 ex = Expr(:toplevel, ex)
             end
             return ex, sizeof(code)
diff --git a/src/literal_parsing.jl b/src/literal_parsing.jl
@@ -12,7 +12,7 @@ function parse_int_literal(str::AbstractString)
     end
     if isnothing(x)
         x = Base.tryparse(Int128, str)
-        if isnothing(x)
+        if x === nothing
             x = Base.parse(BigInt, str)
         end
     end
@@ -358,3 +358,83 @@ function normalize_identifier(str)
     flags = Base.Unicode.UTF8PROC_STABLE | Base.Unicode.UTF8PROC_COMPOSE
     return isascii(str) ? str : utf8proc_map(str, flags)
 end
+
+
+#-------------------------------------------------------------------------------
+function parse_julia_literal(source, head::SyntaxHead, srcrange)
+    # Leaf node
+    k = kind(head)
+    val_str = view(source, srcrange)
+    # Any errors parsing literals are represented as ErrorVal() - this can
+    # happen when the user sets `ignore_errors=true` during parsing.
+    val = if k == K"Integer"
+        parse_int_literal(val_str)
+    elseif k == K"Float"
+        v, code = parse_float_literal(Float64, source.code, first(srcrange),
+                                      last(srcrange)+1)
+        (code === :ok || code === :underflow) ? v : ErrorVal()
+    elseif k == K"Float32"
+        v, code = parse_float_literal(Float32, source.code, first(srcrange),
+                                      last(srcrange)+1)
+        (code === :ok || code === :underflow) ? v : ErrorVal()
+    elseif k in KSet"BinInt OctInt HexInt"
+        parse_uint_literal(val_str, k)
+    elseif k == K"true"
+        true
+    elseif k == K"false"
+        false
+    elseif k == K"Char"
+        io = IOBuffer()
+        had_error = unescape_julia_string(io, source.code, first(srcrange),
+                                          last(srcrange)+1, Diagnostic[])
+        if had_error
+            ErrorVal()
+        else
+            seek(io, 0)
+            c = read(io, Char)
+            eof(io) ? c : ErrorVal()
+        end
+    elseif k == K"Identifier"
+        if has_flags(head, RAW_STRING_FLAG)
+            io = IOBuffer()
+            unescape_raw_string(io, val_str, false)
+            Symbol(normalize_identifier(String(take!(io))))
+        else
+            Symbol(normalize_identifier(val_str))
+        end
+    elseif is_keyword(k)
+        # This should only happen for tokens nested inside errors
+        Symbol(val_str)
+    elseif k in KSet"String CmdString"
+        io = IOBuffer()
+        had_error = false
+        if has_flags(head, RAW_STRING_FLAG)
+            unescape_raw_string(io, val_str, k == K"CmdString")
+        else
+            had_error = unescape_julia_string(io, source.code, first(srcrange),
+                                              last(srcrange)+1, Diagnostic[])
+        end
+        had_error ? ErrorVal() : String(take!(io))
+    elseif is_operator(k)
+        isempty(srcrange)  ?
+            Symbol(untokenize(k)) : # synthetic invisible tokens
+            Symbol(normalize_identifier(val_str))
+    elseif k == K"error"
+        ErrorVal()
+    elseif k == K"MacroName"
+        Symbol("@$(normalize_identifier(val_str))")
+    elseif k == K"StringMacroName"
+        Symbol("@$(normalize_identifier(val_str))_str")
+    elseif k == K"CmdMacroName"
+        Symbol("@$(normalize_identifier(val_str))_cmd")
+    elseif k == K"core_@cmd"
+        Symbol("core_@cmd")
+    elseif is_syntax_kind(head)
+        nothing
+    else
+        # FIXME: this allows us to recover from trivia is_error nodes
+        # that we insert below
+        ErrorVal()
+    end
+end
+
diff --git a/src/parse_stream.jl b/src/parse_stream.jl
@@ -341,10 +341,6 @@ function token_last_byte(stream::ParseStream, i)
     stream.tokens[i].next_byte - 1
 end
 
-function token_span(stream::ParseStream, i)
-    stream.tokens[i].next_byte - stream.tokens[i-1].next_byte
-end
-
 function lookahead_token_first_byte(stream, i)
     i == 1 ? _next_byte(stream) : stream.lookahead[i-1].next_byte
 end
@@ -961,24 +957,25 @@ end
 # API for extracting results from ParseStream
 
 """
-    build_tree(::Type{NodeType}, stream::ParseStream;
+    build_tree(make_node::Function, ::Type{StackEntry}, stream::ParseStream;
                wrap_toplevel_as_kind=nothing, kws...)
 
-Construct a tree with `NodeType` nodes from a ParseStream using depth-first
-traversal. `NodeType` must have the constructors
+Construct a tree from a ParseStream using depth-first traversal. `make_node`
+must have the signature
+
+    make_node(head::SyntaxHead, span::Integer, children)
 
-    NodeType(head::SyntaxHead, span::Integer)
-    NodeType(head::SyntaxHead, span::Integer, children::Vector{NodeType})
+where `children` is either `nothing` for leaf nodes or an iterable of the
+children of type `StackEntry` for internal nodes. `StackEntry` may be a node
+type, but also may include other information required during building the tree.
 
 A single node which covers the input is expected, but if the ParseStream has
 multiple nodes at the top level, `wrap_toplevel_as_kind` may be used to wrap
 them in a single node.
 
-The tree here is constructed depth-first, but it would also be possible to use
-a bottom-up tree builder interface similar to rust-analyzer. (In that case we'd
-traverse the list of ranges backward rather than forward.)
+The tree here is constructed depth-first in postorder.
 """
-function build_tree(::Type{NodeType}, stream::ParseStream;
+function build_tree(make_node::Function, ::Type{NodeType}, stream::ParseStream;
                     wrap_toplevel_as_kind=nothing, kws...) where NodeType
     stack = Vector{NamedTuple{(:first_token,:node),Tuple{Int,NodeType}}}()
 
@@ -996,8 +993,15 @@ function build_tree(::Type{NodeType}, stream::ParseStream;
                 i += 1
                 continue # Ignore removed tokens
             end
-            node = NodeType(head(t), token_span(stream, i))
-            push!(stack, (first_token=i, node=node))
+            srcrange = (stream.tokens[i-1].next_byte:
+                        stream.tokens[i].next_byte - 1)
+            h = head(t)
+            children = (is_syntax_kind(h) || is_keyword(h)) ?
+                (stack[n].node for n=1:0) : nothing
+            node = make_node(h, srcrange, children)
+            if !isnothing(node)
+                push!(stack, (first_token=i, node=node))
+            end
             i += 1
         end
         if j > lastindex(ranges)
@@ -1018,25 +1022,31 @@ function build_tree(::Type{NodeType}, stream::ParseStream;
             while k > 1 && r.first_token <= stack[k-1].first_token
                 k -= 1
             end
+            srcrange = (stream.tokens[r.first_token-1].next_byte:
+                        stream.tokens[r.last_token].next_byte - 1)
             children = (stack[n].node for n = k:length(stack))
-            node = NodeType(head(r), children)
+            node = make_node(head(r), srcrange, children)
             resize!(stack, k-1)
-            push!(stack, (first_token=r.first_token, node=node))
+            if !isnothing(node)
+                push!(stack, (first_token=r.first_token, node=node))
+            end
             j += 1
         end
     end
     if length(stack) == 1
         return only(stack).node
     elseif !isnothing(wrap_toplevel_as_kind)
         # Mostly for debugging
+        srcrange = (stream.tokens[1].next_byte:
+                    stream.tokens[end].next_byte - 1)
         children = (x.node for x in stack)
-        return NodeType(SyntaxHead(wrap_toplevel_as_kind, EMPTY_FLAGS), children)
+        return make_node(SyntaxHead(wrap_toplevel_as_kind, EMPTY_FLAGS),
+                         srcrange, children)
     else
         error("Found multiple nodes at top level")
     end
 end
 
-
 """
     sourcetext(stream::ParseStream; steal_textbuf=true)
 
diff --git a/src/precompile.jl b/src/precompile.jl
@@ -1,5 +1,5 @@
 # Just parse some file as a precompile workload
-let filename = joinpath(@__DIR__, "literal_parsing.jl")
-    text = read(filename, String)
-    parseall(Expr, text)
-end
+# let filename = joinpath(@__DIR__, "literal_parsing.jl")
+#     text = read(filename, String)
+#     parseall(Expr, text)
+# end
diff --git a/src/syntax_tree.jl b/src/syntax_tree.jl
@@ -46,85 +46,9 @@ Base.show(io::IO, ::ErrorVal) = printstyled(io, "✘", color=:light_red)
 
 function SyntaxNode(source::SourceFile, raw::GreenNode{SyntaxHead}, position::Integer=1)
     if !haschildren(raw) && !(is_syntax_kind(raw) || is_keyword(raw))
-        # Leaf node
-        k = kind(raw)
-        val_range = position:position + span(raw) - 1
-        val_str = view(source, val_range)
         # Here we parse the values eagerly rather than representing them as
         # strings. Maybe this is good. Maybe not.
-        #
-        # Any errors parsing literals are represented as ErrorVal() - this can
-        # happen when the user sets `ignore_errors=true` during parsing.
-        val = if k == K"Integer"
-            parse_int_literal(val_str)
-        elseif k == K"Float"
-            v, code = parse_float_literal(Float64, source.code, position,
-                                          position+span(raw))
-            (code === :ok || code === :underflow) ? v : ErrorVal()
-        elseif k == K"Float32"
-            v, code = parse_float_literal(Float32, source.code, position,
-                                          position+span(raw))
-            (code === :ok || code === :underflow) ? v : ErrorVal()
-        elseif k in KSet"BinInt OctInt HexInt"
-            parse_uint_literal(val_str, k)
-        elseif k == K"true"
-            true
-        elseif k == K"false"
-            false
-        elseif k == K"Char"
-            io = IOBuffer()
-            had_error = unescape_julia_string(io, source.code, position,
-                                              position+span(raw), Diagnostic[])
-            if had_error
-                ErrorVal()
-            else
-                seek(io, 0)
-                c = read(io, Char)
-                eof(io) ? c : ErrorVal()
-            end
-        elseif k == K"Identifier"
-            if has_flags(head(raw), RAW_STRING_FLAG)
-                io = IOBuffer()
-                unescape_raw_string(io, val_str, false)
-                Symbol(normalize_identifier(String(take!(io))))
-            else
-                Symbol(normalize_identifier(val_str))
-            end
-        elseif is_keyword(k)
-            # This should only happen for tokens nested inside errors
-            Symbol(val_str)
-        elseif k in KSet"String CmdString"
-            io = IOBuffer()
-            had_error = false
-            if has_flags(head(raw), RAW_STRING_FLAG)
-                unescape_raw_string(io, val_str, k == K"CmdString")
-            else
-                had_error = unescape_julia_string(io, source.code, position,
-                                                  position+span(raw), Diagnostic[])
-            end
-            had_error ? ErrorVal() : String(take!(io))
-        elseif is_operator(k)
-            isempty(val_range)  ?
-                Symbol(untokenize(k)) : # synthetic invisible tokens
-                Symbol(normalize_identifier(val_str))
-        elseif k == K"error"
-            ErrorVal()
-        elseif k == K"MacroName"
-            Symbol("@$(normalize_identifier(val_str))")
-        elseif k == K"StringMacroName"
-            Symbol("@$(normalize_identifier(val_str))_str")
-        elseif k == K"CmdMacroName"
-            Symbol("@$(normalize_identifier(val_str))_cmd")
-        elseif k == K"core_@cmd"
-            Symbol("core_@cmd")
-        elseif is_syntax_kind(raw)
-            nothing
-        else
-            # FIXME: this allows us to recover from trivia is_error nodes
-            # that we insert below
-            @debug "Leaf node of kind $k unknown to SyntaxNode"
-            ErrorVal()
-        end
+        val = parse_julia_literal(source, head(raw), position:position + span(raw) - 1)
         return SyntaxNode(nothing, nothing, SyntaxData(source, raw, position, val))
     else
         cs = SyntaxNode[]
diff --git a/test/benchmark.jl b/test/benchmark.jl
@@ -16,9 +16,9 @@ end
 all_base_code = concat_base()
 
 b_ParseStream = @benchmark JuliaSyntax.parse!(JuliaSyntax.ParseStream(all_base_code), rule=:all)
-b_GreenNode   = @benchmark JuliaSyntax.parseall(JuliaSyntax.GreenNode, all_base_code)
-b_SyntaxNode  = @benchmark JuliaSyntax.parseall(JuliaSyntax.SyntaxNode, all_base_code)
-b_Expr        = @benchmark JuliaSyntax.parseall(Expr, all_base_code)
+b_GreenNode   = @benchmark JuliaSyntax.parseall(JuliaSyntax.GreenNode, all_base_code, ignore_warnings=true)
+b_SyntaxNode  = @benchmark JuliaSyntax.parseall(JuliaSyntax.SyntaxNode, all_base_code, ignore_warnings=true)
+b_Expr        = @benchmark JuliaSyntax.parseall(Expr, all_base_code, ignore_warnings=true)
 
 @info "Benchmarks" ParseStream=b_ParseStream GreenNode=b_GreenNode SyntaxNode=b_SyntaxNode Expr=b_Expr
 
diff --git a/test/expr.jl b/test/expr.jl
diff --git a/test/parser.jl b/test/parser.jl
