fix: improve inference for non-symbolic save_idxs

Author: AayushSabharwal

src/solutions/save_idxs.jl

@@ -105,6 +105,20 @@ struct SavedSubsystem{V, T, M, I, P, Q, C}
     partition_count::C
 end
 
+SavedSubsystem(indp, pobj, ::Nothing) = nothing
+
+function SavedSubsystem(indp, pobj, saved_idxs::Vector{Int})
+    isempty(saved_idxs) && return nothing
+    isempty(variable_symbols(indp)) && return nothing
+    state_map = Dict{Int, Int}(k => v for (k, v) in enumerate(saved_idxs))
+    return SavedSubsystem(state_map, nothing, nothing, nothing, nothing, nothing, nothing)
+end
+
+function SavedSubsystem(indp, pobj, idx::Int)
+    state_map = Dict(1 => idx)
+    return SavedSubsystem(state_map, nothing, nothing, nothing, nothing, nothing, nothing)
+end
+
 function SavedSubsystem(indp, pobj, saved_idxs)
     # nothing saved
     if saved_idxs === nothing || isempty(saved_idxs)
@@ -357,29 +371,32 @@ corresponding to the state variables and a `SavedSubsystem` to pass to `build_so
 The second return value (corresponding to the `SavedSubsystem`) may be `nothing` in case
 one is not required. `save_idxs` may be a scalar or `nothing`.
 """
+get_save_idxs_and_saved_subsystem(prob, ::Nothing) = nothing, nothing
+function get_save_idxs_and_saved_subsystem(prob, save_idxs::Vector{Int})
+    save_idxs, SavedSubsystem(prob, parameter_values(prob), save_idxs)
+end
+function get_save_idxs_and_saved_subsystem(prob, save_idx::Int)
+    save_idx, SavedSubsystem(prob, parameter_values(prob), save_idx)
+end
 function get_save_idxs_and_saved_subsystem(prob, save_idxs)
-    if save_idxs === nothing
-        saved_subsystem = nothing
+    if !(save_idxs isa AbstractArray) || symbolic_type(save_idxs) != NotSymbolic()
+        _save_idxs = [save_idxs]
     else
-        if !(save_idxs isa AbstractArray) || symbolic_type(save_idxs) != NotSymbolic()
-            _save_idxs = [save_idxs]
+        _save_idxs = save_idxs
+    end
+    saved_subsystem = SavedSubsystem(prob, parameter_values(prob), _save_idxs)
+    if saved_subsystem !== nothing
+        _save_idxs = get_saved_state_idxs(saved_subsystem)
+        if isempty(_save_idxs)
+            # no states to save
+            save_idxs = Int[]
+        elseif !(save_idxs isa AbstractArray) ||
+               symbolic_type(save_idxs) != NotSymbolic()
+            # only a single state to save, and save it as a scalar timeseries instead of
+            # single-element array
+            save_idxs = only(_save_idxs)
         else
-            _save_idxs = save_idxs
-        end
-        saved_subsystem = SavedSubsystem(prob, parameter_values(prob), _save_idxs)
-        if saved_subsystem !== nothing
-            _save_idxs = get_saved_state_idxs(saved_subsystem)
-            if isempty(_save_idxs)
-                # no states to save
-                save_idxs = Int[]
-            elseif !(save_idxs isa AbstractArray) ||
-                   symbolic_type(save_idxs) != NotSymbolic()
-                # only a single state to save, and save it as a scalar timeseries instead of
-                # single-element array
-                save_idxs = only(_save_idxs)
-            else
-                save_idxs = _save_idxs
-            end
+            save_idxs = _save_idxs
         end
     end