feat: add assertions support to debug_system

Author: AayushSabharwal

src/debugging.jl

@@ -42,3 +42,21 @@ function debug_sub(ex, funcs; kw...)
     f in funcs ? logged_fun(f, args...; kw...) :
     maketerm(typeof(ex), f, args, metadata(ex))
 end
+
+function _debug_assertion(expr::Bool, message::String, log::Bool)
+    expr && return 0.0
+    log && @error message
+    return NaN
+end
+
+@register_symbolic _debug_assertion(expr::Bool, message::String, log::Bool)
+
+const ASSERTION_LOG_VARIABLE = only(@parameters __log_assertions_ₘₜₖ::Bool = false)
+
+function get_assertions_expr(assertions::Dict{BasicSymbolic, String})
+    term = 0
+    for (k, v) in assertions
+        term += _debug_assertion(k, "Assertion $k failed:\n$v", ASSERTION_LOG_VARIABLE)
+    end
+    return term
+end

src/systems/abstractsystem.jl

@@ -2312,7 +2312,15 @@ function debug_system(
         error("debug_system(sys) only works on systems with no sub-systems! Consider flattening it with flatten(sys) or structural_simplify(sys) first.")
     end
     if has_eqs(sys)
-        @set! sys.eqs = debug_sub.(equations(sys), Ref(functions); kw...)
+        eqs = debug_sub.(equations(sys), Ref(functions); kw...)
+        expr = get_assertions_expr(assertions(sys))
+        eqs[end] = eqs[end].lhs ~ eqs[end].rhs + expr
+        @set! sys.eqs = eqs
+        @set! sys.ps = unique!([get_ps(sys); ASSERTION_LOG_VARIABLE])
+        @set! sys.defaults = merge(get_defaults(sys), Dict(ASSERTION_LOG_VARIABLE => true))
+        if iscomplete(sys)
+            sys = complete(sys; split = is_split(sys))
+        end
     end
     if has_observed(sys)
         @set! sys.observed = debug_sub.(observed(sys), Ref(functions); kw...)