Improve handling of for loops (#374)

`for` loops have an implicit call to `iterate` and the matching is poor. This extracts the type of the iteration variable.
JuliaDebug · Mar 10, 2023 · f61f09b · f61f09b · timholy · Mar 10, 2023
1 parent 2ace6f2
commit f61f09b
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 28 deletions.
diff --git a/TypedSyntax/README.md b/TypedSyntax/README.md
@@ -84,29 +84,29 @@ then (on Julia 1.9)
 ```julia
 julia> tsn, mappings = TypedSyntax.tsn_and_mappings(summer, (Vector{Float64},));
 
-julia> hcat(tsn.typedsource.code, mappings)
-16×2 Matrix{Any}:
- :(_4 = 0)                     Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
- :(_2)                         Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[list]
- :(_3 = Base.iterate(%2))      Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[(= x list)]
- :(_3 === nothing)             Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
- :(Base.not_int(%4))           Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
- :(goto %16 if not %5)         Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
- :(_3::Tuple{Float64, Int64})  Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
- :(_5 = Core.getfield(%7, 1))  Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
- :(Core.getfield(%7, 2))       Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
- :(_4 = _4 + _5)               Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[(+= s x)]
- :(_3 = Base.iterate(%2, %9))  Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
- :(_3 === nothing)             Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
- :(Base.not_int(%12))          Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
- :(goto %16 if not %13)        Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
- :(goto %7)                    Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
- :(return _4)                  Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[s]
+julia> hcat(1:length(mappings), tsn.typedsource.code, mappings)
+16×3 Matrix{Any}:
+  1  :(_4 = 0)                     Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[0]
+  2  :(_2)                         Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[list]
+  3  :(_3 = Base.iterate(%2))      Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[(= x list)]
+  4  :(_3 === nothing)             Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
+  5  :(Base.not_int(%4))           Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
+  6  :(goto %16 if not %5)         Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
+  7  :(_3::Tuple{Float64, Int64})  Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
+  8  :(_5 = Core.getfield(%7, 1))  Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[x]
+  9  :(Core.getfield(%7, 2))       Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
+ 10  :(_4 = _4 + _5)               Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[(+= s x)]
+ 11  :(_3 = Base.iterate(%2, %9))  Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
+ 12  :(_3 === nothing)             Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
+ 13  :(Base.not_int(%12))          Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
+ 14  :(goto %16 if not %13)        Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
+ 15  :(goto %7)                    Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[]
+ 16  :(return _4)                  Union{TreeNode{SyntaxData}, TreeNode{TypedSyntaxData}}[s]
 ```
 The left column contains the statements of the type-inferred code, the right column the mappings back to the source.
 You can see that the majority of these mappings are empty, indicating either no good match or that there were multiple possible matches. This is because lowering changes the implementation so significantly that there are few calls that relate directly to the source.
 
-Nevertheless, many statements in the source can be annotated:
+Nevertheless, most of the statements in the source can be annotated:
 
 ```julia
 julia> tsn
@@ -117,11 +117,11 @@ line:col│ tree                                   │ type
    1:17 │    list                                │Vector{Float64}
    1:22 │  [block]
    2:5  │    [=]
-   2:5  │      s
-   2:9  │      0
+   2:5  │      s                                 │Int64
+   2:9  │      0                                 │Int64
    3:5  │    [for]
    3:8  │      [=]                               │Union{Nothing, Tuple{Float64, Int64}}
-   3:9  │        x
+   3:9  │        x                               │Float64
    3:14 │        list                            │Vector{Float64}
    3:18 │      [block]
    4:9  │        [+=]                            │Float64

diff --git a/TypedSyntax/src/node.jl b/TypedSyntax/src/node.jl
@@ -25,16 +25,13 @@ function tsn_and_mappings(@nospecialize(f), @nospecialize(t); kwargs...)
     tsn_and_mappings(m, src, rt; kwargs...)
 end
 
-const tam_args = Ref{Any}()
-
 function tsn_and_mappings(m::Method, src::CodeInfo, @nospecialize(rt); warn::Bool=true, strip_macros::Bool=false, kwargs...)
     def = definition(String, m)
     if isnothing(def)
         warn && @warn "couldn't retrieve source of $m"
         return nothing, nothing
     end
     sourcetext, lineno = def
-    tam_args[] = m, src, sourcetext
     rootnode = JuliaSyntax.parse(SyntaxNode, sourcetext; filename=string(m.file), first_line=lineno, kwargs...)
     if strip_macros
         rootnode = get_function_def(rootnode)
@@ -521,8 +518,9 @@ function map_ssas_to_source(src::CodeInfo, rootnode::SyntaxNode, Δline::Int)
                         for (arg, j) in argjs
                             if is_slot(arg)
                                 sym = src.slotnames[arg.id]
-                                (sym == Symbol("") || sym == Symbol("#self#")) && continue
-                                for t in symlocs[symloc_key(sym)]
+                                itr = get(symlocs, symloc_key(sym), nothing)
+                                itr === nothing && continue
+                                for t in itr
                                     haskey(symtyps, t) && continue
                                     if skipped_parent(t) == node
                                         is_prec_assignment(node) && t == child(node, 1) && continue
@@ -559,6 +557,16 @@ function map_ssas_to_source(src::CodeInfo, rootnode::SyntaxNode, Δline::Int)
                         end
                     end
                 end
+            elseif isempty(mapped) && isexpr(stmt, :(=))
+                lhs = stmt.args[1]
+                if is_slot(lhs)
+                    empty!(argmapping)
+                    append_targets_for_arg!(argmapping, i, lhs)
+                    if length(argmapping) == 1
+                        node = only(argmapping)
+                        mappings[i] = [node]
+                    end
+                end
             end
         end
         i ∈ used || empty!(mappings[i])   # if the result of the call is not used, don't attach a type to it

diff --git a/TypedSyntax/test/runtests.jl b/TypedSyntax/test/runtests.jl
@@ -243,12 +243,17 @@ include("test_module.jl")
     sig, body = children(tsn)
     @test has_name_typ(child(sig, 2), :list, Vector{Float64})
     @test has_name_typ(child(body, 1, 1), :s, Int)
-    @test_broken has_name_typ(child(body, 2, 1, 1), :x, Float64)
+    @test has_name_typ(child(body, 2, 1, 1), :x, Float64)
     node = child(body, 2, 2, 1)
     @test kind(node) == K"+="
     @test has_name_typ(child(node, 1), :s, Float64)   # if this line runs, the LHS now has type `Float64`
     @test has_name_typ(child(node, 2), :x, Float64)
     @test has_name_typ(child(body, 3, 1), :s, Union{Float64, Int})
+    tsn = TypedSyntaxNode(TSN.summer_iterate, (Vector{Float64},))
+    @test tsn.typ == Union{Int,Float64}
+    sig, body = children(tsn)
+    @test has_name_typ(child(body, 2, 1), :ret, Union{Nothing, Tuple{Float64, Int64}})
+    @test has_name_typ(child(body, 3, 2, 1, 1, 1), :x, Float64)
 
     # `where`, unnamed arguments, and types-as-arguments
     tsn = TypedSyntaxNode(TSN.zerowhere, (Vector{Int16},))

diff --git a/TypedSyntax/test/test_module.jl b/TypedSyntax/test/test_module.jl
@@ -32,6 +32,18 @@ function summer(list)
     end
     return s
 end
+function summer_iterate(list)
+    # same as above, but with an explicit call it `iterate` to ensure our handling
+    # of implicit `iterate` doesn't mess up explicit `iterate`
+    s = 0
+    ret = iterate(list)
+    while ret !== nothing
+        x, state = ret
+        s += x
+        ret = iterate(list, state)
+    end
+    return s
+end
 
 zerowhere(::AbstractArray{T}) where T<:Real = zero(T)
 cb(a, i) = checkbounds(Bool, a, i)