Implement rewrite validator

benchmark/main.jl

@@ -0,0 +1,10 @@
+using Test
+using SciMLBase, Sundials
+
+include("thermalfluid.jl")
+
+Benchmark{3}()()
+@test isa(code_lowered(DAECompiler.factory, Tuple{Val{DAECompiler.Settings(mode=DAECompiler.DAENoInit)}, Benchmark{3}})[1], Core.CodeInfo)
+let sol = solve(DAECProblem(Benchmark{3}(), [1:9;] .=> 0.), IDA())
+    @test_broken sol.retcode == ReturnCode.Success
+end

benchmark/thermalfluid.jl

@@ -2,8 +2,6 @@ using DAECompiler
 using DAECompiler.Intrinsics
 using Polynomials: fit
 using XSteam: my_pT, rho_pT, Cp_pT, tc_pT
-using SciMLBase, Sundials
-using Test
 
 struct Sink; end
 function (::Sink)(inlet)
@@ -282,9 +280,3 @@ function (::Benchmark{N})() where {N}
     PreinsulatedPipe{N}()(in[2], out[1])
     Sink()(out[2])
 end
-
-Benchmark{3}()()
-@test isa(code_lowered(DAECompiler.factory, Tuple{Val{DAECompiler.Settings(mode=DAECompiler.DAENoInit)}, Benchmark{3}})[1], Core.CodeInfo)
-let sol = solve(DAECProblem(Benchmark{3}(), [1:9;] .=> 0.), IDA())
-    @test_broken sol.retcode == ReturnCode.Success
-end

src/DAECompiler.jl

@@ -6,7 +6,7 @@ module DAECompiler
     using Diffractor
     using OrderedCollections
     using Compiler
-    using Compiler: IRCode, IncrementalCompact, DebugInfoStream, NewInstruction, argextype, singleton_type, isexpr, widenconst
+    using Compiler: AbstractLattice, IRCode, IncrementalCompact, DebugInfoStream, NewInstruction, argextype, singleton_type, isexpr, widenconst
     using Core.IR
     using SciMLBase
     using AutoHashEquals
@@ -16,18 +16,21 @@ module DAECompiler
     include("settings.jl")
     include("utils.jl")
     include("intrinsics.jl")
+    include("reflection.jl")
     include("analysis/utils.jl")
     include("analysis/lattice.jl")
     include("analysis/ADAnalyzer.jl")
     include("analysis/scopes.jl")
+    include("analysis/flattening.jl")
     include("analysis/cache.jl")
     include("analysis/refiner.jl")
     include("analysis/ipoincidence.jl")
     include("analysis/structural.jl")
-    include("analysis/flattening.jl")
     include("transform/state_selection.jl")
     include("transform/common.jl")
     include("transform/runtime.jl")
+    include("transform/unoptimized.jl")
+    include("transform/reconstruct.jl")
     include("transform/tearing/schedule.jl")
     include("transform/codegen/dae_factory.jl")
     include("transform/codegen/ode_factory.jl")
@@ -40,5 +43,4 @@ module DAECompiler
     include("analysis/consistency.jl")
     include("interface.jl")
     include("problem_interface.jl")
-    include("reflection.jl")
 end

src/analysis/ADAnalyzer.jl

@@ -62,6 +62,7 @@ end
 
 struct AnalyzedSource
     ir::Compiler.IRCode
+    slotnames::Vector{Any}
     inline_cost::Compiler.InlineCostType
     nargs::UInt
     isva::Bool
@@ -72,19 +73,26 @@ end
     Core.svec(edges..., interp.edges...)
 end
 
+function get_slotnames(def::Method)
+    names = split(def.slot_syms, '\0')
+    return map(Symbol, names)
+end
+
 @override function Compiler.transform_result_for_cache(interp::ADAnalyzer, result::InferenceResult, edges::SimpleVector)
     ir = result.src.optresult.ir
+    slotnames = get_slotnames(result.linfo.def)
     params = Compiler.OptimizationParams(interp)
-    return AnalyzedSource(ir, Compiler.compute_inlining_cost(interp, result), result.src.src.nargs, result.src.src.isva)
+    return AnalyzedSource(ir, slotnames, Compiler.compute_inlining_cost(interp, result), result.src.src.nargs, result.src.src.isva)
 end
 
 @override function Compiler.transform_result_for_local_cache(interp::ADAnalyzer, result::InferenceResult)
     if Compiler.result_is_constabi(interp, result)
         return nothing
     end
     ir = result.src.optresult.ir
+    slotnames = get_slotnames(result.linfo.def)
     params = Compiler.OptimizationParams(interp)
-    return AnalyzedSource(ir, Compiler.compute_inlining_cost(interp, result), result.src.src.nargs, result.src.src.isva)
+    return AnalyzedSource(ir, slotnames, Compiler.compute_inlining_cost(interp, result), result.src.src.nargs, result.src.src.isva)
 end
 
 function Compiler.retrieve_ir_for_inlining(ci::CodeInstance, result::AnalyzedSource)

src/analysis/cache.jl

@@ -21,6 +21,19 @@ end
     var_schedule::Vector{Pair{BitSet, BitSet}}
 end
 
+struct VariableReplacement
+    replaced::Incidence
+    by::Int
+    equation::Int
+end
+
+mutable struct NonlinearReplacementMap
+    const variables::Vector{VariableReplacement}
+    variable_counter::Int
+    equation_counter::Int
+end
+NonlinearReplacementMap() = NonlinearReplacementMap(VariableReplacement[], 0, 0)
+
 """
     StructuralSSARef
 
@@ -35,6 +48,7 @@ struct DAEIPOResult
     opaque_eligible::Bool
     extended_rt::Any
     argtypes
+    argmap::ArgumentMap
     nexternalargvars::Int # total vars is length(var_to_diff)
     nsysmscopes::Int
     nexternaleqs::Int
@@ -44,6 +58,7 @@ struct DAEIPOResult
     total_incidence::Vector{Any}
     eqclassification::Vector{VarEqClassification}
     eq_callee_mapping::Vector{Union{Nothing, Vector{Pair{StructuralSSARef, Int}}}}
+    replacement_map::NonlinearReplacementMap
     names::OrderedDict{Any, ScopeDictEntry} # TODO: OrderedIdDict
     varkinds::Vector{Union{Intrinsics.VarKind, Nothing}}
     eqkinds::Vector{Union{Intrinsics.EqKind, Nothing}}