Support type constructors for IPO

Author: serenity4

src/analysis/flattening.jl

@@ -90,6 +90,8 @@ function flatten_argument!(state::FlatteningState, @nospecialize(argt))
         line = compact[Compiler.OldSSAValue(1)][:line]
         ssa = @insert_instruction_here(compact, line, settings, (:invoke)(nothing, InternalIntrinsics.external_equation)::Eq(eq))
         return ssa
+    elseif isa(argt, Type)
+        return argt.parameters[1]
     elseif isabstracttype(argt) || ismutabletype(argt) || (!isa(argt, DataType) && !isa(argt, PartialStruct))
         line = compact[Compiler.OldSSAValue(1)][:line]
         ssa = @insert_instruction_here(compact, line, settings, error("Cannot IPO model arg type $argt")::Union{})
@@ -245,6 +247,7 @@ function annotate_variables_and_equations(argtypes::Vector{Any}, map::ArgumentMa
 end
 
 init_partialstruct(@nospecialize(T)) = PartialStruct(T, collect(Any, fieldtypes(T)))
+init_partialstruct(pstruct::PartialStruct) = pstruct
 
 function find_base(dict::Dict{CompositeIndex}, index::CompositeIndex)
     for i in reverse(eachindex(index))

src/transform/reconstruct.jl

@@ -104,23 +104,26 @@ function compute_residual_vectors(f, u, du; t = 1.0, mode=DAE, world=Base.tls_wo
     torn_ir = torn_ci.inferred
     removed_states = extract_removed_states(state, key, torn_ir, u, du, t)
 
-    residuals = zeros(length(u))
-    p = SciMLBase.NullParameters()
-    indices = filter(!in(removed_states), eachindex(u))
-    u_compressed = u[indices]
-    du_compressed = du[indices]
-    residuals_compressed = zeros(length(residuals) - length(removed_states))
-
     our_prob = DAECProblem(f, (1,) .=> 1., insert_stmt_debuginfo = true)
     sciml_prob = DiffEqBase.get_concrete_problem(our_prob, true)
     f_compressed! = sciml_prob.f.f
-    f_compressed!(residuals_compressed, du_compressed, u_compressed, p, t)
 
     our_prob = DAECProblem(f, (1,) .=> 1., insert_stmt_debuginfo = true, skip_optimizations = true)
     sciml_prob = DiffEqBase.get_concrete_problem(our_prob, true)
     f_original! = sciml_prob.f.f
-    f_original!(residuals, du, u, p, t)
 
+    residuals = zeros(length(u))
+    p = SciMLBase.NullParameters()
+    indices = filter(!in(removed_states), eachindex(u))
+    u_compressed = u[indices]
+    du_compressed = du[indices]
+
+    n = length(residuals) - length(removed_states)
+    @assert n ≥ 1
+    residuals_compressed = zeros(n)
+    f_compressed!(residuals_compressed, du_compressed, u_compressed, p, t)
+    f_original!(residuals, du, u, p, t)
     expanded = expand_residuals(f, residuals_compressed, u, du, t)
+
     return residuals, expanded
 end

test/validation.jl

@@ -79,6 +79,12 @@ function equation_used_multiple_times!()
     apply_equation_on_ddtx_minus_one(eq, x)
 end
 
+function equation_with_callable!()
+    x = continuous()
+    callable = @noinline Returns(x)
+    always!(ddt(callable()) - 3.0)
+end
+
 @testset "Validation" begin
     refresh() # TODO: remove before merge
 
@@ -149,4 +155,10 @@ end
     residuals, expanded_residuals = compute_residual_vectors(equation_used_multiple_times!, u, du)
     @test residuals ≈ [6.0]
     @test residuals ≈ expanded_residuals
+
+    u = [2.0]
+    du = [4.0]
+    residuals, expanded_residuals = compute_residual_vectors(equation_with_callable!, u, du)
+    @test residuals ≈ [1.0]
+    @test residuals ≈ expanded_residuals
 end;