fix: fix incorrect indexes of array symbolics

Author: AayushSabharwal

src/systems/index_cache.jl

@@ -17,7 +17,8 @@ struct ParameterIndex{P, I}
 end
 
 const ParamIndexMap = Dict{Union{Symbol, BasicSymbolic}, Tuple{Int, Int}}
-const UnknownIndexMap = Dict{Union{Symbol, BasicSymbolic}, Union{Int, UnitRange{Int}}}
+const UnknownIndexMap = Dict{
+    Union{Symbol, BasicSymbolic}, Union{Int, UnitRange{Int}, Array{Int}}}
 
 struct IndexCache
     unknown_idx::UnknownIndexMap
@@ -46,11 +47,26 @@ function IndexCache(sys::AbstractSystem)
             end
             unk_idxs[usym] = sym_idx
 
-            if hasname(sym)
+            if hasname(sym) && (!istree(sym) || operation(sym) !== getindex)
                 unk_idxs[getname(usym)] = sym_idx
             end
             idx += length(sym)
         end
+        for sym in unks
+            usym = unwrap(sym)
+            istree(sym) && operation(sym) === getindex || continue
+            arrsym = arguments(sym)[1]
+            all(haskey(unk_idxs, arrsym[i]) for i in eachindex(arrsym)) || continue
+
+            idxs = [unk_idxs[arrsym[i]] for i in eachindex(arrsym)]
+            if idxs == idxs[begin]:idxs[end]
+                idxs = idxs[begin]:idxs[end]
+            end
+            unk_idxs[arrsym] = idxs
+            if hasname(arrsym)
+                unk_idxs[getname(arrsym)] = idxs
+            end
+        end
     end
 
     disc_buffers = Dict{Any, Set{BasicSymbolic}}()

test/odesystem.jl

@@ -1120,3 +1120,16 @@ tearing_state = TearingState(expand_connections(sys))
 ts_vars = tearing_state.fullvars
 orig_vars = unknowns(sys)
 @test isempty(setdiff(ts_vars, orig_vars))
+
+# Ensure indexes of array symbolics are cached appropriately
+@variables x(t)[1:2]
+@named sys = ODESystem(Equation[], t, [x], [])
+sys1 = complete(sys)
+@named sys = ODESystem(Equation[], t, [x...], [])
+sys2 = complete(sys)
+for sys in [sys1, sys2]
+    for (sym, idx) in [(x, 1:2), (x[1], 1), (x[2], 2)]
+        @test is_variable(sys, sym)
+        @test variable_index(sys, sym) == idx
+    end
+end