SyntaxGraph: Usability and performance tweaks

src/ast.jl

@@ -428,15 +428,21 @@ end
 
 
 """
-Copy AST `ex` into `ctx`
+Recursively copy AST `ex` into `ctx`.  The resulting tree contains new nodes
+only, and does not contain nodes with multiple parents.
+
+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)
+function copy_ast(ctx, ex; copy_source=true)
     # 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
+    s = get(ex, :source, nothing)
     # TODO: Figure out how to use provenance() here?
-    srcref = s isa NodeId ? copy_ast(ctx, SyntaxTree(ex._graph, s))            :
+    srcref = !copy_source ? s                                                  :
+             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)

src/syntax_graph.jl

@@ -22,6 +22,11 @@ function freeze_attrs(graph::SyntaxGraph)
     SyntaxGraph(graph.edge_ranges, graph.edges, frozen_attrs)
 end
 
+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,13 +38,17 @@ function attrnames(graph::SyntaxGraph)
     keys(graph.attributes)
 end
 
+function attrtypes(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")
     _show_attrs(io, graph.attributes)
 end
 
-function ensure_attributes!(graph::SyntaxGraph; kws...)
+function ensure_attributes!(graph::SyntaxGraph{<:Dict}; kws...)
     for (k,v) in pairs(kws)
         @assert k isa Symbol
         @assert v isa Type
@@ -52,19 +61,31 @@ function ensure_attributes!(graph::SyntaxGraph; kws...)
     end
     graph
 end
+function ensure_attributes(graph::SyntaxGraph{<:Dict}; kws...)
+    g = unfreeze_attrs(graph)
+    ensure_attributes!(g; kws...)
+end
 
-function ensure_attributes(graph::SyntaxGraph; kws...)
-    g = SyntaxGraph(graph.edge_ranges, graph.edges, Dict(pairs(graph.attributes)...))
+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 +226,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
 
@@ -412,10 +435,10 @@ end
 
 const SourceAttrType = Union{SourceRef,LineNumberNode,NodeId,Tuple}
 
-function SyntaxTree(graph::SyntaxGraph, node::SyntaxNode)
+function SyntaxTree(graph::SyntaxGraph{<:Dict}, 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 +451,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.attrtypes(gf2)
+    @test (:foo, Type) in JuliaLowering.attrtypes(gf2)
+    @test Set(keys(gf2.attributes)) == Set(keys(gu2.attributes))
+    # no mutation
+    @test !((:test_attr, Symbol) in JuliaLowering.attrtypes(gf1))
+    @test !((:foo, Type) in JuliaLowering.attrtypes(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.attrtypes(gf3))
+    @test !((:foo, Type) in JuliaLowering.attrtypes(gf3))
+    @test Set(keys(gf3.attributes)) == Set(keys(gu3.attributes))
+    # no mutation
+    @test (:test_attr, Symbol) in JuliaLowering.attrtypes(gf2)
+    @test (:foo, Type) in JuliaLowering.attrtypes(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")) ==