fix: refactor IndexCache for non-scalarized unknowns

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -192,7 +192,9 @@ function SymbolicIndexingInterface.is_variable(sys::AbstractSystem, sym)
         ic = get_index_cache(sys)
         h = getsymbolhash(sym)
         return haskey(ic.unknown_idx, h) ||
-               haskey(ic.unknown_idx, getsymbolhash(default_toterm(sym)))
+               haskey(ic.unknown_idx, getsymbolhash(default_toterm(sym))) ||
+               (istree(sym) && operation(sym) === getindex &&
+                is_variable(sys, first(arguments(sym))))
     else
         return any(isequal(sym), variable_symbols(sys)) ||
                hasname(sym) && is_variable(sys, getname(sym))
@@ -213,16 +215,15 @@ function SymbolicIndexingInterface.variable_index(sys::AbstractSystem, sym)
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
         ic = get_index_cache(sys)
         h = getsymbolhash(sym)
-        return if haskey(ic.unknown_idx, h)
-            ic.unknown_idx[h]
-        else
-            h = getsymbolhash(default_toterm(sym))
-            if haskey(ic.unknown_idx, h)
-                ic.unknown_idx[h]
-            else
-                nothing
-            end
-        end
+        haskey(ic.unknown_idx, h) && return ic.unknown_idx[h]
+
+        h = getsymbolhash(default_toterm(sym))
+        haskey(ic.unknown_idx, h) && return ic.unknown_idx[h]
+        sym = unwrap(sym)
+        istree(sym) && operation(sym) === getindex || return nothing
+        idx = variable_index(sys, first(arguments(sym)))
+        idx === nothing && return nothing
+        return idx[arguments(sym)[(begin + 1):end]...]
     end
     idx = findfirst(isequal(sym), variable_symbols(sys))
     if idx === nothing && hasname(sym)

src/systems/index_cache.jl

@@ -21,7 +21,7 @@ end
 const IndexMap = Dict{UInt, Tuple{Int, Int}}
 
 struct IndexCache
-    unknown_idx::Dict{UInt, Int}
+    unknown_idx::Dict{UInt, Union{Int, UnitRange{Int}}}
     discrete_idx::IndexMap
     param_idx::IndexMap
     constant_idx::IndexMap
@@ -36,10 +36,17 @@ end
 
 function IndexCache(sys::AbstractSystem)
     unks = solved_unknowns(sys)
-    unk_idxs = Dict{UInt, Int}()
-    for (i, sym) in enumerate(unks)
-        h = getsymbolhash(sym)
-        unk_idxs[h] = i
+    unk_idxs = Dict{UInt, Union{Int, UnitRange{Int}}}()
+    let idx = 1
+        for sym in unks
+            h = getsymbolhash(sym)
+            if Symbolics.isarraysymbolic(sym)
+                unk_idxs[h] = idx:(idx + length(sym) - 1)
+            else
+                unk_idxs[h] = idx
+            end
+            idx += length(sym)
+        end
     end
 
     disc_buffers = Dict{DataType, Set{BasicSymbolic}}()