Make external eqs tests work again (#52)

Author: Keno

Manifest.toml

@@ -50,7 +50,7 @@ Cthulhu = {rev = "master", url = "https://github.com/JuliaDebug/Cthulhu.jl.git"}
 DifferentiationInterface = {rev = "main", subdir = "DifferentiationInterface", url = "https://github.com/Keno/DifferentiationInterface.jl"}
 Diffractor = {rev = "main", url = "https://github.com/JuliaDiff/Diffractor.jl.git"}
 SimpleNonlinearSolve = {rev = "master", subdir = "lib/SimpleNonlinearSolve", url = "https://github.com/SciML/NonlinearSolve.jl.git"}
-StateSelection = {rev = "main", url = "https://github.com/JuliaComputing/StateSelection.jl"}
+StateSelection = {rev = "main", url = "https://github.com/JuliaComputing/StateSelection.jl.git"}
 
 [compat]
 Accessors = "0.1.36"

Project.toml

@@ -48,12 +48,12 @@ struct DAEIPOResult
     varkinds::Vector{Union{Intrinsics.VarKind, Nothing}}
     eqkinds::Vector{Union{Intrinsics.EqKind, Nothing}}
     # TODO: Chain these rather than copying them
-    warnings::Vector{UnsupportedIRException}
+    warnings::Vector{BadDAECompilerInputException}
 end
 
 struct UncompilableIPOResult
-    warnings::Vector{UnsupportedIRException}
-    error::UnsupportedIRException
+    warnings::Vector{BadDAECompilerInputException}
+    error::BadDAECompilerInputException
 end
 
 function add_equation_row!(graph, solvable_graph, ieq::Int, inc::Incidence)
@@ -92,7 +92,9 @@ function make_structure_from_ipo(ipo::DAEIPOResult)
     graph = BipartiteGraph(neqs, nvars)
     solvable_graph = BipartiteGraph(neqs, nvars)
 
-    for (ieq, inc) in enumerate(ipo.total_incidence)
+    for ieq in 1:length(ipo.total_incidence)
+        isassigned(ipo.total_incidence, ieq) || continue
+        inc = ipo.total_incidence[ieq]
         add_equation_row!(graph, solvable_graph, ieq, inc)
     end
 

src/analysis/cache.jl

@@ -1,4 +1,3 @@
-# TODO: We should unify this function with _make_argument_lattice_elem to ensure consistency
 function _flatten_parameter!(𝕃, compact, argtypes, ntharg, line)
     list = Any[]
     for (argn, argt) in enumerate(argtypes)
@@ -8,6 +7,8 @@ function _flatten_parameter!(𝕃, compact, argtypes, ntharg, line)
             continue
         elseif Base.isprimitivetype(argt) || isa(argt, Incidence)
             push!(list, ntharg(argn))
+        elseif argt === equation || isa(argt, Eq)
+            continue
         elseif isa(argt, Type) && argt <: Intrinsics.AbstractScope
             continue
         elseif isabstracttype(argt) || ismutabletype(argt) || (!isa(argt, DataType) && !isa(argt, PartialStruct))
@@ -36,64 +37,79 @@ function flatten_parameter!(𝕃, compact, argtypes, ntharg, line)
 end
 
 # Needs to match flatten_arguments!
-function process_template_arg!(𝕃, coeffs, eq_mapping, applied_scopes, argt, template_argt, offset=0)
+function process_template_arg!(𝕃, coeffs, eq_mapping, applied_scopes, argt, template_argt, offset=0, eqoffset=0)::Pair{Int, Int}
     if isa(template_argt, Const)
         @assert isa(argt, Const) && argt.val === template_argt.val
-        return offset
+        return Pair{Int, Int}(offset, eqoffset)
     elseif Base.issingletontype(template_argt)
         @assert isa(template_argt, Type) && argt.instance === template_argt.instance
-        return offset
+        return Pair{Int, Int}(offset, eqoffset)
     elseif Base.isprimitivetype(template_argt)
         coeffs[offset+1] = argt
-        return offset + 1
+        return Pair{Int, Int}(offset + 1, eqoffset)
+    elseif template_argt === equation
+        eq_mapping[eqoffset+1] = argt.id
+        return Pair{Int, Int}(offset, eqoffset + 1)
     elseif isabstracttype(template_argt) || ismutabletype(template_argt) || (!isa(template_argt, DataType) && !isa(template_argt, PartialStruct))
-        return offset
+        return Pair{Int, Int}(offset, eqoffset)
     else
         if !isa(template_argt, PartialStruct) && Base.datatype_fieldcount(template_argt) === nothing
-            return offset
+            return Pair{Int, Int}(offset, eqoffset)
         end
         template_fields = isa(template_argt, PartialStruct) ? template_argt.fields : collect(fieldtypes(template_argt))
         return process_template!(𝕃, coeffs, eq_mapping, applied_scopes, Any[Compiler.getfield_tfunc(𝕃, argt, Const(i)) for i = 1:length(template_fields)], template_fields, offset)
     end
 end
 
-function process_template!(𝕃, coeffs, eq_mapping, applied_scopes, argtypes, template_argtypes, offset=0)
+function process_template!(𝕃, coeffs, eq_mapping, applied_scopes, argtypes, template_argtypes, offset=0, eqoffset=0)
     @assert length(argtypes) == length(template_argtypes)
     for (i, template_arg) in enumerate(template_argtypes)
-        offset = process_template_arg!(𝕃, coeffs, eq_mapping, applied_scopes, argtypes[i], template_arg, offset)
+        (offset, eqoffset) = process_template_arg!(𝕃, coeffs, eq_mapping, applied_scopes, argtypes[i], template_arg, offset)
     end
-    return offset
+    return Pair{Int, Int}(offset, eqoffset)
+end
+
+struct TransformedArg
+    ssa::Any
+    offset::Int
+    eqoffset::Int
+    TransformedArg(@nospecialize(arg), new_offset::Int, new_eqoffset::Int) = new(arg, new_offset, new_eqoffset)
 end
 
-function flatten_argument!(compact::Compiler.IncrementalCompact, argt, offset, argtypes::Vector{Any})::Pair{Any, Int}
-    @assert !isa(argt, Incidence)
+function flatten_argument!(compact::Compiler.IncrementalCompact, @nospecialize(argt), offset::Int, eqoffset::Int, argtypes::Vector{Any})::TransformedArg
+    @assert !isa(argt, Incidence) && !isa(argt, Eq)
     if isa(argt, Const)
-        return Pair{Any, Int}(argt.val, offset)
+        return TransformedArg(argt.val, offset, eqoffset)
     elseif Base.issingletontype(argt)
-        return Pair{Any, Int}(argt.instance, offset)
+        return TransformedArg(argt.instance, offset, eqoffset)
     elseif Base.isprimitivetype(argt)
         push!(argtypes, argt)
-        return Pair{Any, Int}(Argument(offset+1), offset+1)
+        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]))
+        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]))
-        return Pair{Any, Int}(ssa, offset)
+        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]))
-            return Pair{Any, Int}(ssa, offset)
+            return TransformedArg(ssa, -1, eqoffset)
         end
-        (args, _, offset) = flatten_arguments!(compact, isa(argt, PartialStruct) ? argt.fields : fieldtypes(argt), offset, argtypes)
+        (args, _, offset) = flatten_arguments!(compact, 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]))
-        return Pair{Any, Int}(ssa, offset)
+        return TransformedArg(ssa, offset, eqoffset)
     end
 end
 
-function flatten_arguments!(compact::Compiler.IncrementalCompact, argtypes, offset=0, new_argtypes::Vector{Any} = Any[])
+function flatten_arguments!(compact::Compiler.IncrementalCompact, argtypes::Vector{Any}, offset::Int=0, eqoffset::Int=0, new_argtypes::Vector{Any} = Any[])
     args = Any[]
     for argt in argtypes
-        (ssa, offset) = flatten_argument!(compact, argt, offset, new_argtypes)
+        (; ssa, offset, eqoffset) = flatten_argument!(compact, argt, offset, eqoffset, new_argtypes)
+        offset == -1 && break
         push!(args, ssa)
     end
-    return (args, new_argtypes, offset)
+    return (args, new_argtypes, offset, eqoffset)
 end

src/analysis/flattening.jl

@@ -6,35 +6,36 @@ struct CalleeMapping
 end
 
 struct CallerMappingState
-    result::DAEIPOResult
+    callee_result::DAEIPOResult
     caller_var_to_diff::DiffGraph
     caller_varclassification::Vector{VarEqClassification}
     caller_varkind::Union{Vector{Intrinsics.VarKind}, Nothing}
     caller_eqclassification::Vector{VarEqClassification}
+    caller_eqkinds::Union{Vector{Intrinsics.EqKind}, Nothing}
 end
 
-function compute_missing_coeff!(coeffs, (;result, caller_var_to_diff, caller_varclassification, caller_varkind)::CallerMappingState, v)
+function compute_missing_coeff!(coeffs, (;callee_result, caller_var_to_diff, caller_varclassification, caller_varkind)::CallerMappingState, v)
     # First find the rootvar, and if we already have a coeff for it
     # apply the derivatives.
     ndiffs = 0
-    calle_inv = invview(result.var_to_diff)
+    calle_inv = invview(callee_result.var_to_diff)
     while calle_inv[v] !== nothing && !isassigned(coeffs, v)
         ndiffs += 1
         v = calle_inv[v]
     end
 
     if !isassigned(coeffs, v)
-        @assert v > result.nexternalargvars # Arg vars should have already been mapped
+        @assert v > callee_result.nexternalargvars # Arg vars should have already been mapped
         # Reached the root and it's an internal variable. We need to allocate
         # it in the caller now
         coeffs[v] = Incidence(add_vertex!(caller_var_to_diff))
-        push!(caller_varclassification, result.varclassification[v] == External ? Owned : CalleeInternal)
-        push!(caller_varkind, result.varkinds[v])
+        push!(caller_varclassification, callee_result.varclassification[v] == External ? Owned : CalleeInternal)
+        push!(caller_varkind, callee_result.varkinds[v])
     end
     thisinc = coeffs[v]
 
     for _ = 1:ndiffs
-        dv = result.var_to_diff[v]
+        dv = callee_result.var_to_diff[v]
         coeffs[dv] = structural_inc_ddt(caller_var_to_diff, caller_varclassification, caller_varkind, thisinc)
         v = dv
     end
@@ -143,10 +144,9 @@ end
 function apply_linear_incidence(𝕃, ret::Eq, caller::CallerMappingState, mapping::CalleeMapping)
     eq_mapping = mapping.eqs[ret.id]
     if eq_mapping == 0
-        error("I removed these from StructuralRefiner for conceptual reasons - if we hit these, lets revisit")
-        #push!(caller_eqclassification, Owned)
-        #push!(caller_eqkinds, result.eqkinds[ret.id])
-        mapping.eqs[ret.id] = eq_mapping = length(caller_eqclassification)
+        push!(caller.caller_eqclassification, Owned)
+        push!(caller.caller_eqkinds, caller.callee_result.eqkinds[ret.id])
+        mapping.eqs[ret.id] = eq_mapping = length(caller.caller_eqclassification)
     end
     return Eq(eq_mapping)
 end
@@ -172,47 +172,3 @@ struct MappingInfo <: Compiler.CallInfo
     result::DAEIPOResult
     mapping::CalleeMapping
 end
-
-function _make_argument_lattice_elem(𝕃, which::Argument, @nospecialize(argt), add_variable!, add_equation!, add_scope!)
-    if isa(argt, Const)
-        #@assert !isa(argt.val, Scope) # Shouldn't have been forwarded
-        return argt
-    elseif isa(argt, Type) && argt <: Intrinsics.AbstractScope
-        return PartialScope(add_scope!(which))
-    elseif isa(argt, Type) && argt == equation
-        return Eq(add_equation!(which))
-    elseif is_non_incidence_type(argt)
-        return argt
-    elseif Compiler.isprimitivetype(argt)
-        inc = Incidence(add_variable!(which))
-        return argt === Float64 ? inc : Incidence(argt, inc.row)
-    elseif isa(argt, PartialStruct)
-        return PartialStruct(𝕃, argt.typ, Any[make_argument_lattice_elem(𝕃, which, f, add_variable!, add_equation!, add_scope!) for f in argt.fields])
-    elseif isabstracttype(argt) || ismutabletype(argt) || !isa(argt, DataType)
-        return nothing
-    else
-        fields = Any[]
-        any = false
-        # TODO: This doesn't handle recursion
-        if Base.datatype_fieldcount(argt) === nothing
-            return nothing
-        end
-        for i = 1:length(fieldtypes(argt))
-            # TODO: Can we make this lazy?
-            ft = fieldtype(argt, i)
-            mft = _make_argument_lattice_elem(𝕃, which, ft, add_variable!, add_equation!, add_scope!)
-            if mft === nothing
-                push!(fields, Incidence(ft))
-            else
-                any = true
-                push!(fields, mft)
-            end
-        end
-        return any ? PartialStruct(𝕃, argt, fields) : nothing
-    end
-end
-
-function make_argument_lattice_elem(𝕃, which::Argument, @nospecialize(argt), add_variable!, add_equation!, add_scope!)
-    mft = _make_argument_lattice_elem(𝕃, which, argt, add_variable!, add_equation!, add_scope!)
-    mft === nothing ? Incidence(argt) : mft
-end

src/analysis/ipoincidence.jl

@@ -81,6 +81,14 @@ const linear_time_and_state_dependent = Linearity(nonlinear = false)
 "The variable is used nonlinearly, with a possible dependence on time and other states."
 const nonlinear = Linearity()
 
+function Base.show(io::IO, lin::Linearity)
+    if lin === nonlinear
+        print(io, "nonlinear")
+    else
+        invoke(Base.show, Tuple{IO, Any}, io, lin)
+    end
+end
+
 join_linearity(a::Linearity, b::Real) = a
 join_linearity(a::Real, b::Linearity) = b
 join_linearity(a::Real, b::Real) = a == b ? a : linear

src/analysis/lattice.jl

@@ -10,6 +10,8 @@ struct StructuralRefiner <: Compiler.AbstractInterpreter
     var_to_diff::DiffGraph
     varkinds::Vector{Intrinsics.VarKind}
     varclassification::Vector{VarEqClassification}
+    eqkinds::Vector{Intrinsics.EqKind}
+    eqclassification::Vector{VarEqClassification}
 end
 
 struct StructureCache; end
@@ -59,7 +61,7 @@ Compiler.cache_owner(::StructuralRefiner) = StructureCache()
     argtypes = Compiler.collect_argtypes(interp, stmt.args, Compiler.StatementState(nothing, false), irsv)[2:end]
     mapping = CalleeMapping(Compiler.optimizer_lattice(interp), argtypes, callee_codeinst, callee_result, callee_codeinst.inferred.ir.argtypes)
     new_rt = apply_linear_incidence(Compiler.optimizer_lattice(interp), callee_result.extended_rt,
-        CallerMappingState(callee_result, interp.var_to_diff, interp.varclassification, interp.varkinds, VarEqClassification[]), mapping)
+        CallerMappingState(callee_result, interp.var_to_diff, interp.varclassification, interp.varkinds, interp.eqclassification, interp.eqkinds), mapping)
 
     # Remember this mapping, both for performance of not having to recompute it
     # and because we may have assigned caller variables to internal variables