SyntaxGraph: Usability and performance tweaks

src/ast.jl

@@ -428,27 +428,47 @@ end
 
 
 """
-Copy AST `ex` into `ctx`
+Recursively copy AST `ex` into `ctx`.
+
+Special provenance handling: If `copy_source` is true, treat the `.source`
+attribute as a reference and recurse on its contents.  Otherwise, treat it like
+any other attribute.
 """
-function copy_ast(ctx, ex)
-    # TODO: Do we need to keep a mapping of node IDs to ensure we don't
-    # double-copy here in the case when some tree nodes are pointed to by
-    # multiple parents? (How much does this actually happen in practice?)
-    s = ex.source
-    # TODO: Figure out how to use provenance() here?
-    srcref = s isa NodeId ? copy_ast(ctx, SyntaxTree(ex._graph, s))            :
-             s isa Tuple  ? map(i->copy_ast(ctx, SyntaxTree(ex._graph, i)), s) :
-             s
-    if !is_leaf(ex)
-        cs = SyntaxList(ctx)
-        for e in children(ex)
-            push!(cs, copy_ast(ctx, e))
-        end
-        ex2 = makenode(ctx, srcref, ex, cs)
+function copy_ast(ctx, ex::SyntaxTree; copy_source=true)
+    graph1 = syntax_graph(ex)
+    graph2 = syntax_graph(ctx)
+    !copy_source && check_same_graph(graph1, graph2)
+    id2 = _copy_ast(graph2, graph1, ex._id, Dict{NodeId, NodeId}(), copy_source)
+    return SyntaxTree(graph2, id2)
+end
+
+function _copy_ast(graph2::SyntaxGraph, graph1::SyntaxGraph,
+                   id1::NodeId, seen, copy_source)
+    let copied = get(seen, id1, nothing)
+        isnothing(copied) || return copied
+    end
+    id2 = newnode!(graph2)
+    seen[id1] = id2
+    src1 = get(SyntaxTree(graph1, id1), :source, nothing)
+    src2 = if !copy_source
+        src1
+    elseif src1 isa NodeId
+        _copy_ast(graph2, graph1, src1, seen, copy_source)
+    elseif src1 isa Tuple
+        map(i->_copy_ast(graph2, graph1, i, seen, copy_source), src1)
     else
-        ex2 = makeleaf(ctx, srcref, ex)
+        src1
     end
-    return ex2
+    copy_attrs!(SyntaxTree(graph2, id2), SyntaxTree(graph1, id1), true)
+    setattr!(graph2, id2; source=src2)
+    if !is_leaf(graph1, id1)
+        cs = NodeId[]
+        for cid in children(graph1, id1)
+            push!(cs, _copy_ast(graph2, graph1, cid, seen, copy_source))
+        end
+        setchildren!(graph2, id2, cs)
+    end
+    return id2
 end
 
 #-------------------------------------------------------------------------------

src/syntax_graph.jl

@@ -22,6 +22,12 @@ function freeze_attrs(graph::SyntaxGraph)
     SyntaxGraph(graph.edge_ranges, graph.edges, frozen_attrs)
 end
 
+# Create a copy of `graph` where the attribute list is mutable
+function unfreeze_attrs(graph::SyntaxGraph)
+    unfrozen_attrs = Dict{Symbol,Any}(pairs(graph.attributes)...)
+    SyntaxGraph(graph.edge_ranges, graph.edges, unfrozen_attrs)
+end
+
 function _show_attrs(io, attributes::Dict)
     show(io, MIME("text/plain"), attributes)
 end
@@ -33,6 +39,10 @@ function attrnames(graph::SyntaxGraph)
     keys(graph.attributes)
 end
 
+function attrdefs(graph::SyntaxGraph)
+    [(k=>typeof(v).parameters[2]) for (k, v) in pairs(graph.attributes)]
+end
+
 function Base.show(io::IO, ::MIME"text/plain", graph::SyntaxGraph)
     print(io, typeof(graph),
           " with $(length(graph.edge_ranges)) vertices, $(length(graph.edges)) edges, and attributes:\n")
@@ -46,25 +56,42 @@ function ensure_attributes!(graph::SyntaxGraph; kws...)
         if haskey(graph.attributes, k)
             v0 = valtype(graph.attributes[k])
             v == v0 || throw(ErrorException("Attribute type mismatch $v != $v0"))
+        elseif graph.attributes isa NamedTuple
+            throw(ErrorException("""
+                ensure_attributes!: $k is not an existing attribute, and the graph's attributes are frozen. \
+                Consider calling non-mutating `ensure_attributes` instead."""))
         else
             graph.attributes[k] = Dict{NodeId,v}()
         end
     end
     graph
 end
 
-function ensure_attributes(graph::SyntaxGraph; kws...)
-    g = SyntaxGraph(graph.edge_ranges, graph.edges, Dict(pairs(graph.attributes)...))
+function ensure_attributes(graph::SyntaxGraph{<:Dict}; kws...)
+    g = unfreeze_attrs(graph)
+    ensure_attributes!(g; kws...)
+end
+
+function ensure_attributes(graph::SyntaxGraph{<:NamedTuple}; kws...)
+    g = unfreeze_attrs(graph)
     ensure_attributes!(g; kws...)
     freeze_attrs(g)
 end
 
-function delete_attributes(graph::SyntaxGraph, attr_names...)
-    attributes = Dict(pairs(graph.attributes)...)
+function delete_attributes!(graph::SyntaxGraph{<:Dict}, attr_names::Symbol...)
     for name in attr_names
-        delete!(attributes, name)
+        delete!(graph.attributes, name)
     end
-    SyntaxGraph(graph.edge_ranges, graph.edges, (; pairs(attributes)...))
+    graph
+end
+
+function delete_attributes(graph::SyntaxGraph{<:Dict}, attr_names::Symbol...)
+    delete_attributes!(unfreeze_attrs(graph), attr_names...)
+end
+
+function delete_attributes(graph::SyntaxGraph{<:NamedTuple}, attr_names::Symbol...)
+    g = delete_attributes!(unfreeze_attrs(graph), attr_names...)
+    freeze_attrs(g)
 end
 
 function newnode!(graph::SyntaxGraph)
@@ -205,7 +232,9 @@ function Base.getproperty(ex::SyntaxTree, name::Symbol)
     name === :_id  && return getfield(ex, :_id)
     _id = getfield(ex, :_id)
     return get(getproperty(getfield(ex, :_graph), name), _id) do
-        error("Property `$name[$_id]` not found")
+        attrstr = join(["\n    $n = $(getproperty(ex, n))"
+                        for n in attrnames(ex)], ",")
+        error("Property `$name[$_id]` not found. Available attributes:$attrstr")
     end
 end
 
@@ -415,7 +444,7 @@ const SourceAttrType = Union{SourceRef,LineNumberNode,NodeId,Tuple}
 function SyntaxTree(graph::SyntaxGraph, node::SyntaxNode)
     ensure_attributes!(graph, kind=Kind, syntax_flags=UInt16, source=SourceAttrType,
                        value=Any, name_val=String)
-    id = _convert_nodes(freeze_attrs(graph), node)
+    id = _convert_nodes(graph, node)
     return SyntaxTree(graph, id)
 end
 
@@ -428,7 +457,7 @@ attrsummary(name, value::Number) = "$name=$value"
 
 function _value_string(ex)
     k = kind(ex)
-    str = k == K"Identifier" || k == K"MacroName" || is_operator(k) ? ex.name_val :
+    str = k in KSet"Identifier MacroName StringMacroName CmdMacroName" || is_operator(k) ? ex.name_val :
           k == K"Placeholder" ? ex.name_val           :
           k == K"SSAValue"    ? "%"                   :
           k == K"BindingId"   ? "#"                   :

test/syntax_graph.jl

@@ -1,3 +1,45 @@
+@testset "SyntaxGraph attrs" begin
+    st = parsestmt(SyntaxTree, "function foo end")
+    g_init = JuliaLowering.unfreeze_attrs(st._graph)
+    gf1 = JuliaLowering.freeze_attrs(g_init)
+    gu1 = JuliaLowering.unfreeze_attrs(gf1)
+
+    # Check that freeze/unfreeze do their jobs
+    @test gf1.attributes isa NamedTuple
+    @test gu1.attributes isa Dict
+    @test Set(keys(gf1.attributes)) == Set(keys(gu1.attributes))
+
+    # ensure_attributes
+    gf2 = JuliaLowering.ensure_attributes(gf1, test_attr=Symbol, foo=Type)
+    gu2 = JuliaLowering.ensure_attributes(gu1, test_attr=Symbol, foo=Type)
+    # returns a graph with the same attribute storage
+    @test gf2.attributes isa NamedTuple
+    @test gu2.attributes isa Dict
+    # does its job
+    @test (:test_attr=>Symbol) in JuliaLowering.attrdefs(gf2)
+    @test (:foo=>Type) in JuliaLowering.attrdefs(gf2)
+    @test Set(keys(gf2.attributes)) == Set(keys(gu2.attributes))
+    # no mutation
+    @test !((:test_attr=>Symbol) in JuliaLowering.attrdefs(gf1))
+    @test !((:foo=>Type) in JuliaLowering.attrdefs(gf1))
+    @test Set(keys(gf1.attributes)) == Set(keys(gu1.attributes))
+
+    # delete_attributes
+    gf3 = JuliaLowering.delete_attributes(gf2, :test_attr, :foo)
+    gu3 = JuliaLowering.delete_attributes(gu2, :test_attr, :foo)
+    # returns a graph with the same attribute storage
+    @test gf3.attributes isa NamedTuple
+    @test gu3.attributes isa Dict
+    # does its job
+    @test !((:test_attr=>Symbol) in JuliaLowering.attrdefs(gf3))
+    @test !((:foo=>Type) in JuliaLowering.attrdefs(gf3))
+    @test Set(keys(gf3.attributes)) == Set(keys(gu3.attributes))
+    # no mutation
+    @test (:test_attr=>Symbol) in JuliaLowering.attrdefs(gf2)
+    @test (:foo=>Type) in JuliaLowering.attrdefs(gf2)
+    @test Set(keys(gf2.attributes)) == Set(keys(gu2.attributes))
+end
+
 @testset "SyntaxTree" begin
     # Expr conversion
     @test Expr(parsestmt(SyntaxTree, "begin a + b ; c end", filename="none")) ==
@@ -16,4 +58,53 @@
     @test kind(tree2) == K"block"
     @test kind(tree2[1]) == K"Identifier" && tree2[1].name_val == "x"
     @test kind(tree2[2]) == K"Identifier" && tree2[2].name_val == "some_unique_identifier"
+
+    "For filling required attrs in graphs created by hand"
+    function testgraph(edge_ranges, edges, more_attrs...)
+        kinds = Dict(map(i->(i=>K"block"), eachindex(edge_ranges)))
+        sources = Dict(map(i->(i=>LineNumberNode(i)), eachindex(edge_ranges)))
+        SyntaxGraph(
+            edge_ranges,
+            edges,
+            Dict(:kind => kinds, :source => sources, more_attrs...))
+    end
+
+    @testset "copy_ast" begin
+        # 1 --> 2 --> 3     src(7-9) = line 7-9
+        # 4 --> 5 --> 6     src(i) = i + 3
+        # 7 --> 8 --> 9
+        g = testgraph([1:1, 2:2, 0:-1, 3:3, 4:4, 0:-1, 5:5, 6:6, 0:-1],
+                      [2, 3, 5, 6, 8, 9],
+                      :source => Dict(enumerate([
+                          map(i->i+3, 1:6)...
+                          map(LineNumberNode, 7:9)...])))
+        st = SyntaxTree(g, 1)
+        stcopy = JuliaLowering.copy_ast(g, st)
+        # Each node should be copied once
+        @test length(g.edge_ranges) === 18
+        @test st._id != stcopy._id
+        @test st ≈ stcopy
+        @test st.source !== stcopy.source
+        @test st.source[1] !== stcopy.source[1]
+        @test st.source[1][1] !== stcopy.source[1][1]
+
+        stcopy2 = JuliaLowering.copy_ast(g, st; copy_source=false)
+        # Only nodes 1-3 should be copied
+        @test length(g.edge_ranges) === 21
+        @test st._id != stcopy2._id
+        @test st ≈ stcopy2
+        @test st.source === stcopy2.source
+        @test st.source[1] === stcopy2.source[1]
+        @test st.source[1][1] === stcopy2.source[1][1]
+
+        # Copy into a new graph
+        new_g = ensure_attributes!(SyntaxGraph(); JuliaLowering.attrdefs(g)...)
+        stcopy3 = JuliaLowering.copy_ast(new_g, st)
+        @test length(new_g.edge_ranges) === 9
+        @test st ≈ stcopy3
+
+        new_g = ensure_attributes!(SyntaxGraph(); JuliaLowering.attrdefs(g)...)
+        # Disallow for now, since we can't prevent dangling sourcerefs
+        @test_throws ErrorException JuliaLowering.copy_ast(new_g, st; copy_source=false)
+    end
 end