fix: fix subexpressions_not_involving_vars!

Now handles callable parameters and array variables correctly

Author: AayushSabharwal

lib/ModelingToolkitBase/src/utils.jl

@@ -1293,8 +1293,19 @@ function subexpressions_not_involving_vars!(
             # Use `Iterators.filter` instead of just checking `ir[i]` and returning `()` for
             # type-stability. The early-exit infers as a `Union`.
             is_atomic = SU.default_is_atomic(ir[i])
+            # If the operation is symbolic, it is the first `outneighbor`. Symbolic operations
+            # are already parameters, we don't want to cache it.
+            drop = Moshi.Match.@match ir[i] begin
+                BSImpl.Term(; f) && if f isa SymbolicT end => begin
+                1
+                end
+                _ => 0
+            end
             # We also don't want to descend into constants - those are not worth caching.
-            Iterators.filter(j -> !is_atomic && !SU.isconst(ir[j]), Graphs.outneighbors(graph, i))
+            return Iterators.filter(
+                j -> !is_atomic && !SU.isconst(ir[j]),
+                Iterators.drop(Graphs.outneighbors(graph, i), drop)
+            )
         end
     end
     rdfs = SU.RecursiveDFS(
@@ -1311,9 +1322,16 @@ function subexpressions_not_involving_vars!(
     # Walk through the usages of `banned_vars` that are in `unbanned_subexprs`, and remove
     # from the candidates any expression we encounter. The remaining vertices are
     # ones that do not use `banned_vars`.
-    unbanned_nbors_fn = let unbanned_subexprs = unbanned_subexprs
+    unbanned_nbors_fn = let unbanned_subexprs = unbanned_subexprs, ir = ir
         function __unbanned_nbors_fn(graph, i)
-            return Iterators.filter(in(unbanned_subexprs), Graphs.inneighbors(graph, i))
+            # If an `inneighbor` is atomic, it means `ir[i]` is an array variable and
+            # the neighbor is a scalarized element. We want to cache specific parts of
+            # symbolic arrays that are unbanned, and only ban usages of the full symbolic
+            # array.
+            return Iterators.filter(
+                j -> j in unbanned_subexprs && !SU.default_is_atomic(ir[j]),
+                Graphs.inneighbors(graph, i)
+            )
         end
     end
     rdfs = SU.RecursiveDFS(

test/scc_nonlinear_problem.jl

@@ -5,6 +5,9 @@ using OrdinaryDiffEqBDF
 using SciMLBase, Symbolics
 using StaticArrays
 using LinearAlgebra, Test
+using SymbolicUtils
+import ModelingToolkitBase
+using Symbolics: SymbolicT, VartypeT, unwrap
 using ModelingToolkit: t_nounits as t, D_nounits as D
 
 @testset "Trivial case" begin
@@ -388,3 +391,14 @@ end
     sol = solve(prob)
     @test SciMLBase.successful_retcode(sol)
 end
+
+@testset "`subexpressions_not_involving_vars!`" begin
+    @variables x[1:3]
+    @parameters (f::Function)(..)
+    ir = IRStructure{VartypeT}()
+    expr = f([x[1], 2x[2], x[2]^2 + x[3]]) + f(x) + f(x[2] + 1)
+    state = Dict{SymbolicT, SymbolicT}()
+    banned_vars = Set{SymbolicT}([x[3], x])
+    ModelingToolkitBase.subexpressions_not_involving_vars!(ir, [unwrap(expr)], banned_vars, state)
+    @test issetequal(keys(state), [x[1], 2x[2], x[2]^2, f(x[2] + 1)])
+end