Improve DebugInfo source provenance

Project.toml

@@ -13,6 +13,7 @@ CentralizedCaches = "d1073d05-2d26-4019-b855-dfa0385fef5e"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compiler = "807dbc54-b67e-4c79-8afb-eafe4df6f2e1"
+ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Cthulhu = "f68482b8-f384-11e8-15f7-abe071a5a75f"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
@@ -59,6 +60,7 @@ CentralizedCaches = "1.1.0"
 ChainRules = "1.50"
 ChainRulesCore = "1.20"
 Compiler = "0"
+ConstructionBase = "1.5.9"
 DiffEqBase = "6.149.2"
 DifferentiationInterface = "0.6.52"
 Diffractor = "0.2.7"

src/DAECompiler.jl

@@ -6,16 +6,16 @@ module DAECompiler
     using Diffractor
     using OrderedCollections
     using Compiler
-    using Compiler: IRCode, IncrementalCompact, argextype, singleton_type, isexpr, widenconst
+    using Compiler: IRCode, IncrementalCompact, DebugInfoStream, NewInstruction, argextype, singleton_type, isexpr, widenconst
     using Core.IR
     using SciMLBase
     using AutoHashEquals
     using LinearAlgebra: LinearAlgebra
     using InteractiveUtils: gen_call_with_extracted_types_and_kwargs
 
+    include("settings.jl")
     include("utils.jl")
     include("intrinsics.jl")
-    include("settings.jl")
     include("analysis/utils.jl")
     include("analysis/lattice.jl")
     include("analysis/ADAnalyzer.jl")

src/analysis/flattening.jl

@@ -1,4 +1,4 @@
-function _flatten_parameter!(𝕃, compact, argtypes, ntharg, line)
+function _flatten_parameter!(𝕃, compact, argtypes, ntharg, line, settings)
     list = Any[]
     for (argn, argt) in enumerate(argtypes)
         if isa(argt, Const)
@@ -18,22 +18,20 @@ function _flatten_parameter!(𝕃, compact, argtypes, ntharg, line)
                 continue
             end
             this = ntharg(argn)
-            nthfield(i) = insert_node_here!(compact,
-                NewInstruction(Expr(:call, getfield, this, i), Compiler.getfield_tfunc(𝕃, argextype(this, compact), Const(i)), line))
+            nthfield(i) = @insert_instruction_here(compact, line, settings, getfield(this, i)::Compiler.getfield_tfunc(𝕃, argextype(this, compact), Const(i)))
             if isa(argt, PartialStruct)
-                fields = _flatten_parameter!(𝕃, compact, argt.fields, nthfield, line)
+                fields = _flatten_parameter!(𝕃, compact, argt.fields, nthfield, line, settings)
             else
-                fields = _flatten_parameter!(𝕃, compact, fieldtypes(argt), nthfield, line)
+                fields = _flatten_parameter!(𝕃, compact, fieldtypes(argt), nthfield, line, settings)
             end
             append!(list, fields)
         end
     end
     return list
 end
 
-function flatten_parameter!(𝕃, compact, argtypes, ntharg, line)
-    return insert_node_here!(compact,
-        NewInstruction(Expr(:call, tuple, _flatten_parameter!(𝕃, compact, argtypes, ntharg, line)...), Tuple, line))
+function flatten_parameter!(𝕃, compact, argtypes, ntharg, line, settings)
+    return @insert_instruction_here(compact, line, settings, tuple(_flatten_parameter!(𝕃, compact, argtypes, ntharg, line, settings)...)::Tuple)
 end
 
 # Needs to match flatten_arguments!
@@ -76,7 +74,7 @@ struct TransformedArg
     TransformedArg(@nospecialize(arg), new_offset::Int, new_eqoffset::Int) = new(arg, new_offset, new_eqoffset)
 end
 
-function flatten_argument!(compact::Compiler.IncrementalCompact, @nospecialize(argt), offset::Int, eqoffset::Int, argtypes::Vector{Any})::TransformedArg
+function flatten_argument!(compact::Compiler.IncrementalCompact, settings::Settings, @nospecialize(argt), offset::Int, eqoffset::Int, argtypes::Vector{Any})::TransformedArg
     @assert !isa(argt, Incidence) && !isa(argt, Eq)
     if isa(argt, Const)
         return TransformedArg(argt.val, offset, eqoffset)
@@ -86,28 +84,32 @@ function flatten_argument!(compact::Compiler.IncrementalCompact, @nospecialize(a
         push!(argtypes, argt)
         return TransformedArg(Argument(offset+1), offset+1, eqoffset)
     elseif argt === equation
-        ssa = insert_node_here!(compact, NewInstruction(Expr(:invoke, nothing, InternalIntrinsics.external_equation), Eq(eqoffset+1), compact[Compiler.OldSSAValue(1)][:line]))
+        line = compact[Compiler.OldSSAValue(1)][:line]
+        ssa = @insert_instruction_here(compact, line, settings, (:invoke)(nothing, InternalIntrinsics.external_equation)::Eq(eqoffset+1))
         return TransformedArg(ssa, offset, eqoffset+1)
     elseif isabstracttype(argt) || ismutabletype(argt) || (!isa(argt, DataType) && !isa(argt, PartialStruct))
-        ssa = insert_node_here!(compact, NewInstruction(Expr(:call, error, "Cannot IPO model arg type $argt"), Union{}, compact[Compiler.OldSSAValue(1)][:line]))
+        line = compact[Compiler.OldSSAValue(1)][:line]
+        ssa = @insert_instruction_here(compact, line, settings, error("Cannot IPO model arg type $argt")::Union{})
         return TransformedArg(ssa, -1, eqoffset)
     else
         if !isa(argt, PartialStruct) && Base.datatype_fieldcount(argt) === nothing
-            ssa = insert_node_here!(compact, NewInstruction(Expr(:call, error, "Cannot IPO model arg type $argt"), Union{}, compact[Compiler.OldSSAValue(1)][:line]))
+            line = compact[Compiler.OldSSAValue(1)][:line]
+            ssa = @insert_instruction_here(compact, line, settings, error("Cannot IPO model arg type $argt")::Union{})
             return TransformedArg(ssa, -1, eqoffset)
         end
-        (args, _, offset) = flatten_arguments!(compact, isa(argt, PartialStruct) ? argt.fields : collect(Any, fieldtypes(argt)), offset, eqoffset, argtypes)
+        (args, _, offset) = flatten_arguments!(compact, settings, isa(argt, PartialStruct) ? argt.fields : collect(Any, fieldtypes(argt)), offset, eqoffset, argtypes)
         offset == -1 && return TransformedArg(ssa, -1, eqoffset)
         this = Expr(:new, isa(argt, PartialStruct) ? argt.typ : argt, args...)
-        ssa = insert_node_here!(compact, NewInstruction(this, argt, compact[Compiler.OldSSAValue(1)][:line]))
+        line = compact[Compiler.OldSSAValue(1)][:line]
+        ssa = @insert_instruction_here(compact, line, settings, this::argt)
         return TransformedArg(ssa, offset, eqoffset)
     end
 end
 
-function flatten_arguments!(compact::Compiler.IncrementalCompact, argtypes::Vector{Any}, offset::Int=0, eqoffset::Int=0, new_argtypes::Vector{Any} = Any[])
+function flatten_arguments!(compact::Compiler.IncrementalCompact, settings::Settings, argtypes::Vector{Any}, offset::Int=0, eqoffset::Int=0, new_argtypes::Vector{Any} = Any[])
     args = Any[]
     for argt in argtypes
-        (; ssa, offset, eqoffset) = flatten_argument!(compact, argt, offset, eqoffset, new_argtypes)
+        (; ssa, offset, eqoffset) = flatten_argument!(compact, settings, argt, offset, eqoffset, new_argtypes)
         offset == -1 && break
         push!(args, ssa)
     end

src/analysis/refiner.jl

@@ -7,6 +7,7 @@ of structural incidence information.
 """
 struct StructuralRefiner <: Compiler.AbstractInterpreter
     world::UInt
+    settings::Settings
     var_to_diff::DiffGraph
     varkinds::Vector{Intrinsics.VarKind}
     varclassification::Vector{VarEqClassification}
@@ -51,7 +52,7 @@ Compiler.cache_owner(::StructuralRefiner) = StructureCache()
     end
 
     callee_codeinst = invokee::CodeInstance
-    callee_result = structural_analysis!(callee_codeinst, Compiler.get_inference_world(interp))
+    callee_result = structural_analysis!(callee_codeinst, Compiler.get_inference_world(interp), interp.settings)
 
     if isa(callee_result, UncompilableIPOResult) || isa(callee_result.extended_rt, Const) || isa(callee_result.extended_rt, Type)
         # If this is uncompilable, we will be notfiying the user in the outer loop - here we just ignore it

src/analysis/structural.jl

@@ -16,14 +16,14 @@ function find_matching_ci(predicate, mi::MethodInstance, world::UInt)
     return nothing
 end
 
-function structural_analysis!(ci::CodeInstance, world::UInt)
+function structural_analysis!(ci::CodeInstance, world::UInt, settings::Settings)
     # Check if we have aleady done this work - if so return the cached result
     result_ci = find_matching_ci(ci->ci.owner == StructureCache(), ci.def, world)
     if result_ci !== nothing
         return result_ci.inferred
     end
 
-    result = _structural_analysis!(ci, world)
+    result = _structural_analysis!(ci, world, settings)
     # TODO: The world bounds might have been narrowed
     cache_dae_ci!(ci, result, nothing, nothing, StructureCache())
 
@@ -40,7 +40,7 @@ struct EqVarState
     eq_callee_mapping
 end
 
-function _structural_analysis!(ci::CodeInstance, world::UInt)
+function _structural_analysis!(ci::CodeInstance, world::UInt, settings::Settings)
     # Variables
     var_to_diff = DiffGraph(0)
     varclassification = VarEqClassification[]
@@ -83,7 +83,7 @@ function _structural_analysis!(ci::CodeInstance, world::UInt)
     compact = IncrementalCompact(ir)
     old_argtypes = copy(ir.argtypes)
     empty!(ir.argtypes)
-    (arg_replacements, new_argtypes, nexternalargvars, nexternaleqs) = flatten_arguments!(compact, old_argtypes, 0, 0, ir.argtypes)
+    (arg_replacements, new_argtypes, nexternalargvars, nexternaleqs) = flatten_arguments!(compact, settings, old_argtypes, 0, 0, ir.argtypes)
     if nexternalargvars == -1
         return UncompilableIPOResult(warnings, UnsupportedIRException("Unhandled argument types", Compiler.finish(compact)))
     end
@@ -98,7 +98,7 @@ function _structural_analysis!(ci::CodeInstance, world::UInt)
     argtypes = Any[Incidence(new_argtypes[i], i) for i = 1:nexternalargvars]
 
     # Allocate variable and equation numbers of any incoming arguments
-    refiner = StructuralRefiner(world, var_to_diff, varkinds, varclassification, eqkinds, eqclassification)
+    refiner = StructuralRefiner(world, settings, var_to_diff, varkinds, varclassification, eqkinds, eqclassification)
     nexternalargvars = length(var_to_diff)
 
     # Go through the IR, annotating each intrinsic with an appropriate taint
@@ -337,7 +337,7 @@ function _structural_analysis!(ci::CodeInstance, world::UInt)
         if isa(info, MappingInfo)
             (; result, mapping) = info
         else
-            result = structural_analysis!(callee_codeinst, Compiler.get_inference_world(refiner))
+            result = structural_analysis!(callee_codeinst, Compiler.get_inference_world(refiner), settings)
 
             if isa(result, UncompilableIPOResult)
                 # TODO: Stack trace?
@@ -360,9 +360,9 @@ function _structural_analysis!(ci::CodeInstance, world::UInt)
         # Rewrite to flattened ABI
         compact[SSAValue(i)] = nothing
         compact.result_idx -= 1
-        new_args = _flatten_parameter!(Compiler.optimizer_lattice(refiner), compact, callee_codeinst.inferred.ir.argtypes, arg->stmt.args[arg+1], line)
+        new_args = _flatten_parameter!(Compiler.optimizer_lattice(refiner), compact, callee_codeinst.inferred.ir.argtypes, arg->stmt.args[arg+1], line, settings)
 
-        new_call = insert_node_here!(compact,
+        new_call = insert_instruction_here!(compact, settings, @__SOURCE__,
                 NewInstruction(Expr(:invoke, (StructuralSSARef(compact.result_idx), callee_codeinst), new_args...), stmtype, info, line, stmtflags))
         compact.ssa_rename[compact.idx - 1] = new_call
 
@@ -386,8 +386,8 @@ function _structural_analysis!(ci::CodeInstance, world::UInt)
         line = ret_stmt_inst[:line]
         Compiler.delete_inst_here!(compact)
 
-        (new_ret, ultimate_rt) = rewrite_ipo_return!(Compiler.typeinf_lattice(refiner), compact, line, ret_stmt.val, ultimate_rt, eqvars)
-        insert_node_here!(compact, NewInstruction(ReturnNode(new_ret), ultimate_rt, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), true)
+        (new_ret, ultimate_rt) = rewrite_ipo_return!(Compiler.typeinf_lattice(refiner), compact, line, settings, ret_stmt.val, ultimate_rt, eqvars)
+        insert_instruction_here!(compact, settings, @__SOURCE__, NewInstruction(ReturnNode(new_ret), ultimate_rt, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), reverse_affinity = true)
     elseif isa(ultimate_rt, Type)
         # If we don't have any internal variables (in which case we might have to to do a more aggressive rewrite), strengthen the incidence
         # by demoting to full incidence over the argument variables. Incidence is not allowed to propagate through global mutable state, so
@@ -415,7 +415,7 @@ function _structural_analysis!(ci::CodeInstance, world::UInt)
         warnings)
 end
 
-function rewrite_ipo_return!(𝕃, compact::IncrementalCompact, line, ssa, ultimate_rt::Any, eqvars::EqVarState)
+function rewrite_ipo_return!(𝕃, compact::IncrementalCompact, line, settings, ssa, ultimate_rt::Any, eqvars::EqVarState)
     if isa(ultimate_rt, Eq)
         return Pair{Any, Any}(ssa, ultimate_rt)
     end
@@ -425,22 +425,23 @@ function rewrite_ipo_return!(𝕃, compact::IncrementalCompact, line, ssa, ultim
         new_types = Any[]
         for i = 1:length(ultimate_rt.fields)
             ssa_type = Compiler.getfield_tfunc(𝕃, ultimate_rt, Const(i))
-            ssa_field = insert_node_here!(compact,
-                NewInstruction(Expr(:call, getfield, variable), ssa_type, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), true)
+            ssa_field = insert_instruction_here!(compact, settings, @__SOURCE__,
+                NewInstruction(Expr(:call, getfield, variable), ssa_type, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), reverse_affinity = true)
 
-            (new_field, new_type) = rewrite_ipo_return!(𝕃, compact, line, ssa_field, ssa_type, eqvars)
+            (new_field, new_type) = rewrite_ipo_return!(𝕃, compact, line, settings, ssa_field, ssa_type, eqvars)
             push!(new_fields, new_field)
             push!(new_types, new_type)
         end
         newT = Compiler.PartialStruct(ultimate_rt.typ, new_types)
         if widenconst(ultimate_rt) <: Tuple
-            retssa = insert_node_here!(compact,
-                NewInstruction(Expr(:call, tuple, new_fields...), newT, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), true)
+            retssa = insert_instruction_here!(compact, settings, @__SOURCE__,
+                NewInstruction(Expr(:call, tuple, new_fields...), newT, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), reverse_affinity = true)
         else
-            T = insert_node_here!(compact,
-                NewInstruction(Expr(:call, typeof, ssa), Type, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), true)
-            retssa = insert_node_here!(compact,
-                NewInstruction(Expr(:new, T, new_fields...), newT, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), true)
+            T = insert_instruction_here!(compact, settings, @__SOURCE__,
+                NewInstruction(Expr(:call, typeof, ssa), Type, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), reverse_affinity = true)
+
+            retssa = insert_instruction_here!(compact, settings, @__SOURCE__,
+                NewInstruction(Expr(:new, T, new_fields...), newT, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), reverse_affinity = true)
         end
         return Pair{Any, Any}(retssa, newT)
     end
@@ -453,8 +454,8 @@ function rewrite_ipo_return!(𝕃, compact::IncrementalCompact, line, ssa, ultim
     push!(eqvars.varclassification, External)
     push!(eqvars.varkinds, Intrinsics.Continuous)
 
-    new_var_ssa = insert_node_here!(compact,
-        NewInstruction(Expr(:invoke, nothing, variable), Incidence(nonlinrepl), Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), true)
+    new_var_ssa = insert_instruction_here!(compact, settings, @__SOURCE__,
+        NewInstruction(Expr(:invoke, nothing, variable), Incidence(nonlinrepl), Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED); reverse_affinity = true)
 
     eq_incidence = ultimate_rt - Incidence(nonlinrepl)
     push!(eqvars.total_incidence, eq_incidence)
@@ -463,14 +464,14 @@ function rewrite_ipo_return!(𝕃, compact::IncrementalCompact, line, ssa, ultim
     push!(eqvars.eqkinds, Intrinsics.Always)
     new_eq = length(eqvars.total_incidence)
 
-    new_eq_ssa = insert_node_here!(compact,
-        NewInstruction(Expr(:invoke, nothing, equation), Eq(new_eq), Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), true)
+    new_eq_ssa = insert_instruction_here!(compact, settings, @__SOURCE__,
+        NewInstruction(Expr(:invoke, nothing, equation), Eq(new_eq), Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED); reverse_affinity = true)
 
-    eq_val_ssa = insert_node_here!(compact,
-        NewInstruction(Expr(:call, InternalIntrinsics.assign_var, new_var_ssa, ssa), eq_incidence, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), true)
+    eq_val_ssa = insert_instruction_here!(compact, settings, @__SOURCE__,
+        NewInstruction(Expr(:call, InternalIntrinsics.assign_var, new_var_ssa, ssa), eq_incidence, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED); reverse_affinity = true)
 
-    eq_call_ssa = insert_node_here!(compact,
-        NewInstruction(Expr(:invoke, nothing, new_eq_ssa, eq_val_ssa), Nothing, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED), true)
+    eq_call_ssa = insert_instruction_here!(compact, settings, @__SOURCE__,
+        NewInstruction(Expr(:invoke, nothing, new_eq_ssa, eq_val_ssa), Nothing, Compiler.NoCallInfo(), line, Compiler.IR_FLAG_REFINED); reverse_affinity = true)
 
     T = widenconst(ultimate_rt)
     # TODO: We don't have a way to express that the return value is directly this variable for arbitrary types

src/interface.jl

@@ -11,7 +11,7 @@ function factory_gen(@nospecialize(fT), settings::Settings, world::UInt = Base.g
     ci = ad_typeinf(world, Tuple{fT}; force_inline_all=settings.force_inline_all, edges=Core.svec(factory_mi))
 
     # Perform or lookup DAECompiler specific analysis for this system.
-    result = structural_analysis!(ci, world)
+    result = structural_analysis!(ci, world, settings)
 
     if isa(result, UncompilableIPOResult)
         if isa(result.error, FunctionErrorsException)
@@ -58,17 +58,18 @@ function factory_gen(@nospecialize(fT), settings::Settings, world::UInt = Base.g
     end
 
     # Generate the IR implementation of `factory`, returning the DAEFunction/ODEFunction
+    slotnames = nothing
     if settings.mode in (DAE, DAENoInit)
-        ir_factory = dae_factory_gen(tstate, ci, diff_key, world, settings, settings.mode == DAE ? init_key : nothing)
+        ir_factory, slotnames = dae_factory_gen(tstate, ci, diff_key, world, settings, settings.mode == DAE ? init_key : nothing)
     elseif settings.mode in (ODE, ODENoInit)
-        ir_factory = ode_factory_gen(tstate, ci, diff_key, world, settings, settings.mode == ODE ? init_key : nothing)
+        ir_factory, slotnames = ode_factory_gen(tstate, ci, diff_key, world, settings, settings.mode == ODE ? init_key : nothing)
     elseif settings.mode == InitUncompress
         ir_factory = init_uncompress_gen(result, ci, init_key, diff_key, world, settings)
     else
         return :(error("Unknown generation mode: $(settings.mode)"))
     end
 
-    src = ir_to_src(ir_factory, settings)
+    src = ir_to_src(ir_factory, settings; slotnames)
     src.ssavaluetypes = length(src.code)
     src.min_world = @atomic ci.min_world
     src.max_world = @atomic ci.max_world
@@ -107,5 +108,6 @@ function refresh()
         $(Expr(:meta, :generated_only))
         $(Expr(:meta, :generated, factory_gen))
     end
+    return nothing
 end
 refresh()