compile one observed function for all CopyParamsByTemplate fallback batches

lib/ModelingToolkitBase/src/systems/problem_utils.jl

@@ -768,7 +768,7 @@ Callable struct designed for use by `MTKParametersReconstructor`. Uses a fixed s
 act as a very dynamic (and limited) observed function returning an array. See `__apply_copy_template`
 for the supported templates.
 """
-struct CopyParamsByTemplate{IsRoot, T, N}
+struct CopyParamsByTemplate{IsRoot, T, N, G}
     """
     List of templates.
     """
@@ -777,10 +777,27 @@ struct CopyParamsByTemplate{IsRoot, T, N}
     Size of the returned buffer.
     """
     size::NTuple{N, Int}
+    """
+    Merged getter for all symbolic-fallback batches, or `nothing`.
+    """
+    fallback_getter::G
 end
 
-function CopyParamsByTemplate{IR}(temp::T, size::NTuple{N, Int}) where {IR, T, N}
-    return CopyParamsByTemplate{IR, T, N}(temp, size)
+function CopyParamsByTemplate{IR}(
+        temp::T, size::NTuple{N, Int}, fallback_getter::G = nothing
+    ) where {IR, T, N, G}
+    return CopyParamsByTemplate{IR, T, N, G}(temp, size, fallback_getter)
+end
+
+"""
+    $TYPEDEF
+
+Template entry indexing into the result of a `CopyParamsByTemplate`'s merged
+`fallback_getter`. `range` is the contiguous slice of that result corresponding to this
+fallback batch's symbols.
+"""
+struct FallbackSlice
+    range::UnitRange{Int}
 end
 
 function __apply_copy_template(valp, template)
@@ -826,9 +843,24 @@ function __apply_copy_template(valp, template)
     end
 end
 
+@inline function __apply_root_template(src, template, fb)
+    if template isa FallbackSlice
+        return fb[template.range]
+    else
+        return __apply_copy_template(src, template)
+    end
+end
+
 function (cp::CopyParamsByTemplate{IsRoot})(src) where {IsRoot}
     return if IsRoot
-        reshape(mapreduce(Base.Fix1(__apply_copy_template, src), vcat, cp.template), cp.size)
+        if cp.fallback_getter === nothing
+            reshape(mapreduce(Base.Fix1(__apply_copy_template, src), vcat, cp.template), cp.size)
+        else
+            fb = cp.fallback_getter(src)
+            reshape(
+                mapreduce(t -> __apply_root_template(src, t, fb), vcat, cp.template), cp.size
+            )
+        end
     else
         buffers = map(Base.Fix1(__apply_copy_template, src), cp.template)
         if cp.template isa Tuple
@@ -998,12 +1030,30 @@ function CopyParamsByTemplate(srcsys::AbstractSystem, syms::AbstractArray{Symbol
         end
     end
 
+    fallback_idxs = Int[]
     for i in eachindex(template)
-        if template[i] isa Vector{SymbolicT}
-            template[i] = cached_template_getu(srcsys, template[i]; kws...)
-            delete!(elem_types, Vector{SymbolicT})
-            push!(elem_types, typeof(template[i]))
+        template[i] isa Vector{SymbolicT} && push!(fallback_idxs, i)
+    end
+    fallback_getter = nothing
+    if length(fallback_idxs) == 1
+        i = fallback_idxs[1]
+        template[i] = cached_template_getu(srcsys, template[i]::Vector{SymbolicT}; kws...)
+        delete!(elem_types, Vector{SymbolicT})
+        push!(elem_types, typeof(template[i]))
+    elseif length(fallback_idxs) > 1
+        all_fallback = SymbolicT[]
+        for i in fallback_idxs
+            append!(all_fallback, template[i]::Vector{SymbolicT})
         end
+        fallback_getter = cached_template_getu(srcsys, all_fallback; kws...)
+        offset = 0
+        for i in fallback_idxs
+            len = length(template[i]::Vector{SymbolicT})
+            template[i] = FallbackSlice((offset + 1):(offset + len))
+            offset += len
+        end
+        delete!(elem_types, Vector{SymbolicT})
+        push!(elem_types, FallbackSlice)
     end
 
     # Lift buffer indices of multi-buffer portions (discrete/constants/nonnumeric) into
@@ -1019,7 +1069,9 @@ function CopyParamsByTemplate(srcsys::AbstractSystem, syms::AbstractArray{Symbol
         end
     end
 
-    return CopyParamsByTemplate{true}(__specialize_templates(template, elem_types), size(syms))
+    return CopyParamsByTemplate{true}(
+        __specialize_templates(template, elem_types), size(syms), fallback_getter
+    )
 end
 
 function CopyParamsByTemplate(srcsys::AbstractSystem, syms::AbstractArray; kws...)