feat: support indexing with Symbols in IndexCache

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -346,6 +346,9 @@ function SymbolicIndexingInterface.is_variable(sys::AbstractSystem, sym)
 end
 
 function SymbolicIndexingInterface.is_variable(sys::AbstractSystem, sym::Symbol)
+    if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
+        return haskey(ic.unknown_idx, hash(sym))
+    end
     return any(isequal(sym), getname.(variable_symbols(sys))) ||
            count('₊', string(sym)) == 1 &&
            count(isequal(sym), Symbol.(nameof(sys), :₊, getname.(variable_symbols(sys)))) ==
@@ -377,6 +380,9 @@ function SymbolicIndexingInterface.variable_index(sys::AbstractSystem, sym)
 end
 
 function SymbolicIndexingInterface.variable_index(sys::AbstractSystem, sym::Symbol)
+    if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
+        return get(ic.unknown_idx, h, nothing)
+    end
     idx = findfirst(isequal(sym), getname.(variable_symbols(sys)))
     if idx !== nothing
         return idx
@@ -401,19 +407,24 @@ function SymbolicIndexingInterface.is_parameter(sys::AbstractSystem, sym)
         ic = get_index_cache(sys)
         h = getsymbolhash(sym)
         return if haskey(ic.param_idx, h) || haskey(ic.discrete_idx, h) ||
-                  haskey(ic.constant_idx, h) || haskey(ic.dependent_idx, h)
+                  haskey(ic.constant_idx, h) || haskey(ic.dependent_idx, h) ||
+                  haskey(ic.nonnumeric_idx, h)
             true
         else
             h = getsymbolhash(default_toterm(sym))
             haskey(ic.param_idx, h) || haskey(ic.discrete_idx, h) ||
-                haskey(ic.constant_idx, h) || haskey(ic.dependent_idx, h)
+                haskey(ic.constant_idx, h) || haskey(ic.dependent_idx, h) ||
+                haskey(ic.nonnumeric_idx, h)
         end
     end
     return any(isequal(sym), parameter_symbols(sys)) ||
            hasname(sym) && is_parameter(sys, getname(sym))
 end
 
 function SymbolicIndexingInterface.is_parameter(sys::AbstractSystem, sym::Symbol)
+    if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
+        return ParameterIndex(ic, sym) !== nothing
+    end
     return any(isequal(sym), getname.(parameter_symbols(sys))) ||
            count('₊', string(sym)) == 1 &&
            count(isequal(sym),
@@ -426,7 +437,6 @@ function SymbolicIndexingInterface.parameter_index(sys::AbstractSystem, sym)
     end
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
         ic = get_index_cache(sys)
-        h = getsymbolhash(sym)
         return if (idx = ParameterIndex(ic, sym)) !== nothing
             idx
         elseif (idx = ParameterIndex(ic, default_toterm(sym))) !== nothing
@@ -444,6 +454,9 @@ function SymbolicIndexingInterface.parameter_index(sys::AbstractSystem, sym)
 end
 
 function SymbolicIndexingInterface.parameter_index(sys::AbstractSystem, sym::Symbol)
+    if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
+        return ParameterIndex(ic, sym)
+    end
     idx = findfirst(isequal(sym), getname.(parameter_symbols(sys)))
     if idx !== nothing
         return idx

src/systems/index_cache.jl

@@ -40,10 +40,16 @@ function IndexCache(sys::AbstractSystem)
     let idx = 1
         for sym in unks
             h = getsymbolhash(sym)
-            if Symbolics.isarraysymbolic(sym)
-                unk_idxs[h] = idx:(idx + length(sym) - 1)
+            sym_idx = if Symbolics.isarraysymbolic(sym)
+                idx:(idx + length(sym) - 1)
             else
-                unk_idxs[h] = idx
+                idx
+            end
+            unk_idxs[h] = sym_idx
+
+            if hasname(sym)
+                h = hash(getname(sym))
+                unk_idxs[h] = sym_idx
             end
             idx += length(sym)
         end
@@ -122,6 +128,10 @@ function IndexCache(sys::AbstractSystem)
                 idxs[h] = (i, j)
                 h = getsymbolhash(default_toterm(p))
                 idxs[h] = (i, j)
+                if hasname(p)
+                    h = hash(getname(p))
+                    idxs[h] = (i, j)
+                end
             end
             push!(buffer_sizes, BufferTemplate(T, length(buf)))
         end
@@ -151,7 +161,7 @@ end
 
 function ParameterIndex(ic::IndexCache, p, sub_idx = ())
     p = unwrap(p)
-    h = getsymbolhash(p)
+    h = p isa Symbol ? hash(p) : getsymbolhash(p)
     return if haskey(ic.param_idx, h)
         ParameterIndex(SciMLStructures.Tunable(), (ic.param_idx[h]..., sub_idx...))
     elseif haskey(ic.discrete_idx, h)