From a1bdba1955a8e281b1d57f8bdacbe4f0d4f08be9 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Wed, 23 Oct 2024 15:49:21 +0530
Subject: [PATCH] fixup! feat: support array variables in `linear_expansion`

---
 src/linear_algebra.jl | 21 +++++++++++++++++++++
 test/linear_solver.jl |  3 +++
 2 files changed, 24 insertions(+)

diff --git a/src/linear_algebra.jl b/src/linear_algebra.jl
index 61fc8639a..177a3571a 100644
--- a/src/linear_algebra.jl
+++ b/src/linear_algebra.jl
@@ -338,6 +338,11 @@ function _linear_expansion(t, x)
     else
         for (i, arg) in enumerate(args)
             isequal(arg, arrx) && return (0, 0, false)
+            if symbolic_type(arg) == NotSymbolic()
+                arg isa AbstractArray || continue
+                _occursin_array(x, arrx, arg) && return (0, 0, false)
+                continue
+            end
             a, b, islinear = linear_expansion(arg, x)
             (_iszero(a) && islinear) || return (0, 0, false)
         end
@@ -345,6 +350,22 @@ function _linear_expansion(t, x)
     end
 end
 
+"""
+    _occursin_array(sym, arrsym, arr)
+
+Check if `sym` (or, if `sym` is an element of an array symbolic, the array symbolic
+`arrsym`) occursin in the non-symbolic array `arr`.
+"""
+function _occursin_array(sym, arrsym, arr)
+    for el in arr
+        if symbolic_type(el) == NotSymbolic()
+            return el isa AbstractArray && _occursin_array(sym, arrsym, el)
+        else
+            return occursin(sym, el) || occursin(arrsym, el)
+        end
+    end
+end
+
 ###
 ### Utilities
 ###
diff --git a/test/linear_solver.jl b/test/linear_solver.jl
index 89b2a3371..1bd56ec04 100644
--- a/test/linear_solver.jl
+++ b/test/linear_solver.jl
@@ -70,4 +70,7 @@ a, b, islinear = Symbolics.linear_expansion(D(x) - x, x)
     @test islin && isequal(a, 2) && isequal(b, 0)
     a, b, islin = linear_expansion(y[1], x[1])
     @test islin && isequal(a, 0) && isequal(b, y[1])
+    @test !linear_expansion(z([x...]), x[1])[3]
+    @test !linear_expansion(z(collect(unwrap(x))), x[1])[3]
+    @test !linear_expansion(z([x, 2x]), x[1])[3]
 end