[NewOptimizer] Track inlining info

base/compiler/optimize.jl

@@ -23,10 +23,11 @@ mutable struct OptimizationState
             s_edges = []
             frame.stmt_edges[1] = s_edges
         end
-        next_label = max(label_counter(frame.src.code), length(frame.src.code)) + 10
+        src = frame.src
+        next_label = max(label_counter(src.code), length(src.code)) + 10
         return new(frame.linfo, frame.vararg_type_container,
                    s_edges::Vector{Any},
-                   frame.src, frame.mod, frame.nargs,
+                   src, frame.mod, frame.nargs,
                    next_label, frame.min_valid, frame.max_valid,
                    frame.params)
     end
@@ -53,7 +54,7 @@ mutable struct OptimizationState
             inmodule = linfo.def::Module
             nargs = 0
         end
-        next_label = max(label_counter(frame.src.code), length(frame.src.code)) + 10
+        next_label = max(label_counter(src.code), length(src.code)) + 10
         vararg_type_container = nothing # if you want something more accurate, set it yourself :P
         return new(linfo, vararg_type_container,
                    s_edges::Vector{Any},
@@ -262,7 +263,7 @@ function isinlineable(m::Method, src::CodeInfo, mod::Module, params::Params, bon
     return inlineable
 end
 
-const enable_new_optimizer = RefValue{Bool}(false)
+const enable_new_optimizer = RefValue(false)
 
 # converge the optimization work
 function optimize(me::InferenceState)
@@ -294,12 +295,13 @@ function optimize(me::InferenceState)
             reindex_labels!(opt)
             nargs = Int(opt.nargs) - 1
             if def isa Method
-                topline = LineNumberNode(Int(def.line), def.file)
+                topline = LineInfoNode(opt.mod, def.name, def.file, def.line, 0)
             else
-                topline = LineNumberNode(0)
+                topline = LineInfoNode(opt.mod, NullLineInfo.name, NullLineInfo.file, 0, 0)
             end
-            ir = run_passes(opt.src, opt.mod, nargs, topline)
-            replace_code!(opt.src, ir, nargs, topline)
+            linetable = [topline]
+            ir = run_passes(opt.src, nargs, linetable)
+            replace_code!(opt.src, ir, nargs, linetable)
             push!(opt.src.code, LabelNode(length(opt.src.code) + 1))
             any_phi = true
         elseif !any_phi
@@ -1474,6 +1476,7 @@ function inlineable(@nospecialize(f), @nospecialize(ft), e::Expr, atypes::Vector
         end
     end
     local end_label # if it ends in a goto, we might need to add a come-from label
+    npops = 0 # we don't require them to balance, so find out how many need to be added
     for i = 1:body_len
         a = body.args[i]
         if isa(a, GotoNode)
@@ -1496,6 +1499,27 @@ function inlineable(@nospecialize(f), @nospecialize(ft), e::Expr, atypes::Vector
                 end
                 edges = Any[newlabels[edge::Int + 1] - 1 for edge in edges]
                 a.args[2] = PhiNode(edges, a.args[2].values)
+            elseif a.head === :meta && length(a.args) > 0
+                a1 = a.args[1]
+                if a1 === :push_loc
+                    npops += 1
+                elseif a1 === :pop_loc
+                    if length(a.args) > 1
+                        npops_loc = a.args[2]::Int
+                        if npops_loc > npops # corrupt IR - try to normalize it to limit the impact
+                            a.args[2] = npops
+                            npops = 0
+                        else
+                            npops -= npops_loc
+                        end
+                    else
+                        if npops == 0 # corrupt IR - try to normalize it to limit the impact
+                            body.args[i] = nothing
+                        else
+                            npops -= 1
+                        end
+                    end
+                end
             end
         end
     end
@@ -1575,27 +1599,23 @@ function inlineable(@nospecialize(f), @nospecialize(ft), e::Expr, atypes::Vector
             isa(linenode.file, Symbol) && (file = linenode.file)
         end
     end
-    if do_coverage
+    npops += 1
+    if do_coverage || !isempty(stmts)
+        pop_loc = (npops == 1) ? Expr(:meta, :pop_loc) : Expr(:meta, :pop_loc, npops)
         # Check if we are switching module, which is necessary to catch user
         # code inlined into `Base` with `--code-coverage=user`.
-        # Assume we are inlining directly into `enclosing` instead of another
-        # function inlined in it
         mod = method.module
-        if mod === sv.mod
+        if !do_coverage || mod === sv.mod
             pushfirst!(stmts, Expr(:meta, :push_loc, file,
                                  method.name, line))
         else
             pushfirst!(stmts, Expr(:meta, :push_loc, file,
                                  method.name, line, mod))
         end
-        push!(stmts, Expr(:meta, :pop_loc))
-    elseif !isempty(stmts)
-        pushfirst!(stmts, Expr(:meta, :push_loc, file,
-                             method.name, line))
-        if isa(stmts[end], LineNumberNode)
-            stmts[end] = Expr(:meta, :pop_loc)
+        if !do_coverage && !isempty(stmts) && isa(stmts[end], LineNumberNode)
+            stmts[end] = pop_loc
         else
-            push!(stmts, Expr(:meta, :pop_loc))
+            push!(stmts, pop_loc)
         end
     end
 

base/compiler/ssair/driver.jl

@@ -1,3 +1,12 @@
+struct LineInfoNode
+    mod::Module
+    method::Symbol
+    file::Symbol
+    line::Int
+    inlined_at::Int
+end
+const NullLineInfo = LineInfoNode(@__MODULE__, Symbol(""), Symbol(""), 0, 0)
+
 include("compiler/ssair/ir.jl")
 include("compiler/ssair/domtree.jl")
 include("compiler/ssair/slot2ssa.jl")
@@ -10,19 +19,41 @@ macro show(s)
     # return :(println($(QuoteNode(s)), " = ", $(esc(s))))
 end
 
-function normalize(@nospecialize(stmt), meta::Vector{Any}, inline::Vector{Any}, loc::RefValue{LineNumberNode})
+function normalize(@nospecialize(stmt), meta::Vector{Any}, table::Vector{LineInfoNode}, loc::RefValue{Int})
     if isa(stmt, Expr)
         if stmt.head == :meta
             args = stmt.args
             if length(args) > 0
                 a1 = args[1]
                 if a1 === :push_loc
-                    push!(inline, stmt)
+                    let
+                        current = loc[]
+                        filename = args[2]::Symbol
+                        methodname = NullLineInfo.method
+                        mod = table[current].mod
+                        line = 0
+                        for i = 3:length(args)
+                            ai = args[i]
+                            if ai isa Symbol
+                                methodname = ai
+                            elseif ai isa Int32
+                                line = Int(ai)
+                            elseif ai isa Int64
+                                line = Int(ai)
+                            elseif ai isa Module
+                                mod = ai
+                            end
+                        end
+                        push!(table, LineInfoNode(mod, methodname, filename, line, current))
+                        loc[] = length(table)
+                    end
                 elseif a1 === :pop_loc
                     n = (length(args) > 1) ? args[2]::Int : 1
                     for i in 1:n
-                        isempty(inline) && break
-                        pop!(inline)
+                        current = loc[]
+                        current = table[current].inlined_at
+                        current == 0 && break
+                        loc[] = current
                     end
                 else
                     push!(meta, stmt)
@@ -39,22 +70,35 @@ function normalize(@nospecialize(stmt), meta::Vector{Any}, inline::Vector{Any},
     elseif isa(stmt, LabelNode)
         return nothing
     elseif isa(stmt, LineNumberNode)
-        loc[] = stmt
+        let # need to expand this node so that it is source-location independent
+            current = loc[]
+            info = table[current]
+            methodname = info.method
+            mod = info.mod
+            file = stmt.file
+            file isa Symbol || (file = info.file)
+            line = stmt.line
+            push!(table, LineInfoNode(mod, methodname, file, line, info.inlined_at))
+            loc[] = length(table)
+        end
         return nothing
     end
     return stmt
 end
 
-function run_passes(ci::CodeInfo, mod::Module, nargs::Int, toploc::LineNumberNode)
+function run_passes(ci::CodeInfo, nargs::Int, linetable::Vector{LineInfoNode})
+    mod = linetable[1].mod
     ci.code = copy(ci.code)
     meta = Any[]
-    lines = fill(LineNumberNode(0), length(ci.code))
-    let inline = Any[], loc = RefValue(toploc)
+    lines = fill(0, length(ci.code))
+    let loc = RefValue(1)
         for i = 1:length(ci.code)
             stmt = ci.code[i]
-            stmt = normalize(stmt, meta, inline, loc)
+            stmt = normalize(stmt, meta, linetable, loc)
             ci.code[i] = stmt
-            stmt === nothing || (lines[i] = loc[])
+            if !(stmt === nothing)
+                lines[i] = loc[]
+            end
         end
     end
     ci.code = strip_trailing_junk!(ci.code, lines)

base/compiler/ssair/ir.jl

@@ -120,22 +120,22 @@ function first_insert_for_bb(code, cfg, block)
 end
 
 
-const NewNode = Tuple{Int, Any, Any, LineNumberNode}
+const NewNode = Tuple{Int, Any, Any, #=LineNumber=#Int}
 
 struct IRCode
     stmts::Vector{Any}
     types::Vector{Any}
-    lines::Vector{LineNumberNode}
+    lines::Vector{Int}
     argtypes::Vector{Any}
     cfg::CFG
     new_nodes::Vector{NewNode}
     mod::Module
     meta::Vector{Any}
 
-    function IRCode(stmts::Vector{Any}, lines::Vector{LineNumberNode}, cfg::CFG, argtypes::Vector{Any}, mod::Module, meta::Vector{Any})
+    function IRCode(stmts::Vector{Any}, lines::Vector{Int}, cfg::CFG, argtypes::Vector{Any}, mod::Module, meta::Vector{Any})
         return new(stmts, Any[], lines, argtypes, cfg, NewNode[], mod, meta)
     end
-    function IRCode(ir::IRCode, stmts::Vector{Any}, types::Vector{Any}, lines::Vector{LineNumberNode}, cfg::CFG, new_nodes::Vector{NewNode})
+    function IRCode(ir::IRCode, stmts::Vector{Any}, types::Vector{Any}, lines::Vector{Int}, cfg::CFG, new_nodes::Vector{NewNode})
         return new(stmts, types, lines, ir.argtypes, cfg, new_nodes, ir.mod, ir.meta)
     end
 end
@@ -317,7 +317,7 @@ mutable struct IncrementalCompact
     ir::IRCode
     result::Vector{Any}
     result_types::Vector{Any}
-    result_lines::Vector{LineNumberNode}
+    result_lines::Vector{Int}
     ssa_rename::Vector{Any}
     used_ssas::Vector{Int}
     late_fixup::Vector{Int}
@@ -331,7 +331,7 @@ mutable struct IncrementalCompact
         new_len = length(code.stmts) + length(code.new_nodes)
         result = Array{Any}(uninitialized, new_len)
         result_types = Array{Any}(uninitialized, new_len)
-        result_lines = Array{LineNumberNode}(uninitialized, new_len)
+        result_lines = Array{Int}(uninitialized, new_len)
         ssa_rename = Any[SSAValue(i) for i = 1:new_len]
         used_ssas = fill(0, new_len)
         late_fixup = Vector{Int}()

base/compiler/ssair/legacy.jl

@@ -10,7 +10,74 @@ function ssaargmap(f, @nospecialize(stmt))
     urs[]
 end
 
-function replace_code!(ci::CodeInfo, code::IRCode, nargs::Int, topline::LineNumberNode)
+function line_to_vector(line::Int, linetable::Vector{LineInfoNode})
+    lines = Int[]
+    while line != 0
+        push!(lines, line)
+        line = linetable[line].inlined_at
+    end
+    return lines
+end
+
+function push_new_lineinfo!(new_code::Vector{Any}, topline::Int, line::Int, linetable::Vector{LineInfoNode})
+    # separate the info into three sets: pops, line-change, pushes
+    do_coverage = coverage_enabled()
+    topmod = linetable[line].mod
+    toplines = line_to_vector(topline, linetable)
+    lines = line_to_vector(line, linetable)
+    while !isempty(lines) && !isempty(toplines) && lines[end] == toplines[end]
+        # remove common frames, recording changes to topmod
+        topmod = linetable[pop!(lines)].mod
+        pop!(toplines)
+    end
+    # check whether the outermost frame changed, or just the line number
+    newframe = true
+    topfile = NullLineInfo.file
+    if !isempty(lines) && !isempty(toplines)
+        let topline = linetable[toplines[end]]
+            line = linetable[lines[end]]
+            if topline.inlined_at == 0 || (topline.mod === line.mod && topline.method === line.method)
+                # we could track frame_id precisely, but llvm / dwarf has no support for that,
+                # and it wouldn't really be that meaningful after statements moved around,
+                # so we just do fuzzy matching here in the legacy-format writer
+                newframe = false
+                topfile = topline.file
+            end
+        end
+    end
+    # first pop the old frame(s)
+    npops = length(toplines) + newframe - 1
+    if npops > 0
+        push!(new_code, (npops == 1) ? Expr(:meta, :pop_loc) : Expr(:meta, :pop_loc, npops))
+    end
+    # then change the line number
+    if !newframe
+        let line = linetable[pop!(lines)]
+            if line.file === topfile
+                loc = LineNumberNode(line.line)
+            else
+                loc = LineNumberNode(line.line, line.file)
+            end
+            push!(new_code, loc)
+            topmod = line.mod
+        end
+    end
+    # then push the new frames
+    while !isempty(lines)
+        let line = linetable[pop!(lines)]
+            if !do_coverage || line.mod == topmod
+                loc = Expr(:meta, :push_loc, line.file, line.method, line.line)
+            else
+                loc = Expr(:meta, :push_loc, line.file, line.method, line.line, line.mod)
+            end
+            push!(new_code, loc)
+            topmod = line.mod
+        end
+    end
+    nothing
+end
+
+function replace_code!(ci::CodeInfo, code::IRCode, nargs::Int, linetable::Vector{LineInfoNode})
     if !isempty(code.new_nodes)
         code = compact!(code)
     end
@@ -68,6 +135,7 @@ function replace_code!(ci::CodeInfo, code::IRCode, nargs::Int, topline::LineNumb
     label_mapping = IdDict{Int, Int}()
     terminator_mapping = IdDict{Int, Int}()
     fixup = Int[]
+    topline = 1
     for (idx, stmt) in pairs(code.stmts)
         line = code.lines[idx]
         # push labels first
@@ -78,8 +146,8 @@ function replace_code!(ci::CodeInfo, code::IRCode, nargs::Int, topline::LineNumb
             push!(new_code, LabelNode(length(new_code) + 1))
         end
         # then metadata
-        if !(line.file === nothing && line.line === 0) && !(line === topline)
-            push!(new_code, line)
+        if line != 0 && line != topline
+            push_new_lineinfo!(new_code, topline, line, linetable)
             topline = line
         end
         # record if this'll need a fixup after stmt number

base/compiler/ssair/slot2ssa.jl

@@ -183,7 +183,7 @@ function rename_uses!(ir::IRCode, ci::CodeInfo, idx::Int, @nospecialize(stmt), r
     return fixemup!(stmt->true, stmt->renames[slot_id(stmt)], ir, ci, idx, stmt)
 end
 
-function strip_trailing_junk!(code::Vector{Any}, lines::Vector{LineNumberNode})
+function strip_trailing_junk!(code::Vector{Any}, lines::Vector{Int})
     # Remove `nothing`s at the end, we don't handle them well
     # (we expect the last instruction to be a terminator)
     for i = length(code):-1:1
@@ -198,7 +198,7 @@ function strip_trailing_junk!(code::Vector{Any}, lines::Vector{LineNumberNode})
     term = code[end]
     if !isa(term, GotoIfNot) && !isa(term, GotoNode) && !isa(term, ReturnNode)
         push!(code, ReturnNode{Any}())
-        push!(lines, LineNumberNode(0))
+        push!(lines, 0)
     end
     return code
 end