feat: propagate state machines in structural simplification

Author: AayushSabharwal

src/systems/systems.jl

@@ -75,8 +75,10 @@ function __mtkcompile(sys::AbstractSystem; simplify = false,
         return simplify_optimization_system(sys; kwargs..., sort_eqs, simplify)
     end
 
+    sys, statemachines = extract_top_level_statemachines(sys)
     sys = expand_connections(sys)
-    state = TearingState(sys; sort_eqs)
+    state = TearingState(sys)
+    append!(state.statemachines, statemachines)
 
     @unpack structure, fullvars = state
     @unpack graph, var_to_diff, var_types = structure

src/systems/systemstructure.jl

@@ -204,7 +204,6 @@ end
 mutable struct TearingState{T <: AbstractSystem} <: AbstractTearingState{T}
     """The system of equations."""
     sys::T
-    original_eqs::Vector{Equation}
     """The set of variables of the system."""
     fullvars::Vector{BasicSymbolic}
     structure::SystemStructure
@@ -216,6 +215,7 @@ mutable struct TearingState{T <: AbstractSystem} <: AbstractTearingState{T}
     are not used in the rest of the system.
     """
     additional_observed::Vector{Equation}
+    statemachines::Vector{T}
 end
 
 TransformationState(sys::AbstractSystem) = TearingState(sys)
@@ -224,6 +224,22 @@ function system_subset(ts::TearingState, ieqs::Vector{Int})
     @set! ts.sys.eqs = eqs[ieqs]
     @set! ts.original_eqs = ts.original_eqs[ieqs]
     @set! ts.structure = system_subset(ts.structure, ieqs)
+    if all(eq -> eq.rhs isa StateMachineOperator, get_eqs(ts.sys))
+        names = Symbol[]
+        for eq in get_eqs(ts.sys)
+            if eq.lhs isa Transition
+                push!(names, first(namespace_hierarchy(nameof(eq.rhs.from))))
+                push!(names, first(namespace_hierarchy(nameof(eq.rhs.to))))
+            elseif eq.lhs isa InitialState
+                push!(names, first(namespace_hierarchy(nameof(eq.rhs.s))))
+            else
+                error("Unhandled state machine operator")
+            end
+        end
+        @set! ts.statemachines = filter(x -> nameof(x) in names, ts.statemachines)
+    else
+        @set! ts.statemachines = eltype(ts.statemachines)[]
+    end
     ts
 end
 
@@ -277,6 +293,49 @@ function symbolic_contains(var, set)
         all(x -> x in set, Symbolics.scalarize(var))
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Descend through the system hierarchy and look for statemachines. Remove equations from
+the inner statemachine systems. Return the new `sys` and an array of top-level
+statemachines.
+"""
+function extract_top_level_statemachines(sys::AbstractSystem)
+    eqs = get_eqs(sys)
+
+    if !isempty(eqs) && all(eq -> eq.lhs isa StateMachineOperator, eqs)
+        # top-level statemachine
+        with_removed = @set sys.systems = map(remove_child_equations, get_systems(sys))
+        return with_removed, [sys]
+    elseif !isempty(eqs) && any(eq -> eq.lhs isa StateMachineOperator, eqs)
+        # error: can't mix
+        error("Mixing statemachine equations and standard equations in a top-level statemachine is not allowed.")
+    else
+        # descend
+        subsystems = get_systems(sys)
+        newsubsystems = eltype(subsystems)[]
+        statemachines = eltype(subsystems)[]
+        for subsys in subsystems
+            newsubsys, sub_statemachines = extract_top_level_statemachines(subsys)
+            push!(newsubsystems, newsubsys)
+            append!(statemachines, sub_statemachines)
+        end
+        @set! sys.systems = newsubsystems
+        return sys, statemachines
+    end
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Return `sys` with all equations (including those in subsystems) removed.
+"""
+function remove_child_equations(sys::AbstractSystem)
+    @set! sys.eqs = eltype(get_eqs(sys))[]
+    @set! sys.systems = map(remove_child_equations, get_systems(sys))
+    return sys
+end
+
 function TearingState(sys; quick_cancel = false, check = true, sort_eqs = true)
     # flatten system
     sys = flatten(sys)
@@ -342,9 +401,16 @@ function TearingState(sys; quick_cancel = false, check = true, sort_eqs = true)
             # change the equation if the RHS is `missing` so the rest of this loop works
             eq = 0.0 ~ coalesce(eq.rhs, 0.0)
         end
-        rhs = quick_cancel ? quick_cancel_expr(eq.rhs) : eq.rhs
-        if !_iszero(eq.lhs)
+        is_statemachine_equation = false
+        if eq.lhs isa StateMachineOperator
+            is_statemachine_equation = true
+            eq = eq
+            rhs = eq.rhs
+        elseif _iszero(eq.lhs)
+            rhs = quick_cancel ? quick_cancel_expr(eq.rhs) : eq.rhs
+        else
             lhs = quick_cancel ? quick_cancel_expr(eq.lhs) : eq.lhs
+            rhs = quick_cancel ? quick_cancel_expr(eq.rhs) : eq.rhs
             eq = 0 ~ rhs - lhs
         end
         empty!(varsbuf)
@@ -409,8 +475,7 @@ function TearingState(sys; quick_cancel = false, check = true, sort_eqs = true)
                 addvar!(v, VARIABLE)
             end
         end
-
-        if isalgeq
+        if isalgeq || is_statemachine_equation
             eqs[i] = eq
         else
             eqs[i] = eqs[i].lhs ~ rhs
@@ -528,11 +593,10 @@ function TearingState(sys; quick_cancel = false, check = true, sort_eqs = true)
 
     eq_to_diff = DiffGraph(nsrcs(graph))
 
-    ts = TearingState(sys, original_eqs, fullvars,
+    ts = TearingState(sys, fullvars,
         SystemStructure(complete(var_to_diff), complete(eq_to_diff),
             complete(graph), nothing, var_types, false),
-        Any[], param_derivative_map, original_eqs, Equation[])
-
+        Any[], param_derivative_map, original_eqs, Equation[], typeof(sys)[])
     return ts
 end
 
@@ -862,6 +926,7 @@ function mtkcompile!(state::TearingState; simplify = false,
             inputs = [inputs; clocked_inputs[continuous_id]], outputs, disturbance_inputs,
             check_consistency, fully_determined,
             kwargs...)
+        additional_passes = get(kwargs, :additional_passes, nothing)
         if !isnothing(additional_passes) && any(discrete_compile_pass, additional_passes)
             discrete_pass_idx = findfirst(discrete_compile_pass, additional_passes)
             discrete_compile = additional_passes[discrete_pass_idx]