Merge pull request #4510 from SciML/as/faster-odeproblem

refactor: use `IRStructure` in `calculate_A_b`

Author: AayushSabharwal

lib/ModelingToolkitBase/src/systems/codegen.jl

@@ -1443,6 +1443,7 @@ function calculate_A_b(sys::System; sparse = false, throw = true)
     V = SymbolicT[]
     b = Vector{SymbolicT}(undef, length(rhss))
     query_predicate = in(Set{SymbolicT}(dvs))
+    ir = get_irstructure(sys)
     # `linear_expansion` caches values based on `var`. This loop ordering helps
     # avoid invalidating the cache frequently.
     for (j, var) in enumerate(dvs)
@@ -1451,7 +1452,8 @@ function calculate_A_b(sys::System; sparse = false, throw = true)
             p, q, islinear = lex(resid)
             # An equation such as `0 ~ 1 + x * y` will return `(x, 1, true)` for `y`.
             # Check that `p` doesn't contain unknowns to avoid this.
-            if !islinear || SU.query(query_predicate, p)
+            SU.populate_ir!(ir, p)
+            if !islinear || SU.query(query_predicate, ir, p)
                 err = NotAffineError(fulleqs[i].rhs, var)
                 store_to_mutable_cache!(sys, CachedLinearAb, err)
                 throw || return nothing