Preserve FIRK EvalSol cache for compatible states

Keep the cached EvalSol fast path when the unflattened state can be copied safely, and rebuild it only for incompatible element types such as vector-of-vector ForwardDiff states.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Author: ChrisRackauckas

lib/BoundaryValueDiffEqFIRK/src/firk.jl

@@ -705,6 +705,14 @@ function __construct_problem(
     return __construct_problem(cache, y, loss_bc, loss_collocation, loss, pt, constraint)
 end
 
+@inline function __firk_eval_sol!(eval_sol, y, mesh, cache)
+    if eltype(first(y)) <: eltype(first(eval_sol.u))
+        eval_sol.u[1:end] .= y
+        return eval_sol
+    end
+    return EvalSol(__restructure_sol(y, cache.in_size), mesh, cache)
+end
+
 function __construct_problem(
         cache::FIRKCacheExpand{iip, T, DC, tune_parameters}, y, loss_bc::BC, loss_collocation::C,
         loss::LF, ::StandardBVProblem, ::Val{true}
@@ -1395,25 +1403,25 @@ end
 
 @views function __firk_loss!(
         resid, u, p, y, pt::StandardBVProblem, bc!::BC, residual, mesh,
-        cache, _, trait::DiffCacheNeeded, constraint
+        cache, eval_sol, trait::DiffCacheNeeded, constraint
     ) where {BC}
     y_ = recursive_unflatten!(y, u)
     resids = [get_tmp(r, u) for r in residual]
     Φ!(resids[2:end], cache, y_, u, trait, constraint)
-    eval_sol = EvalSol(__restructure_sol(y_, cache.in_size), mesh, cache)
+    eval_sol = __firk_eval_sol!(eval_sol, y_, mesh, cache)
     eval_bc_residual!(resids[1], pt, bc!, eval_sol, p, mesh)
     recursive_flatten!(resid, resids)
     return nothing
 end
 
 @views function __firk_loss!(
         resid, u, p, y, pt::StandardBVProblem, bc!::BC, residual, mesh,
-        cache, _, trait::NoDiffCacheNeeded, constraint
+        cache, eval_sol, trait::NoDiffCacheNeeded, constraint
     ) where {BC}
     y_ = recursive_unflatten!(y, u)
     resids = [r for r in residual]
     Φ!(resids[2:end], cache, y_, u, trait, constraint)
-    eval_sol = EvalSol(__restructure_sol(y_, cache.in_size), mesh, cache)
+    eval_sol = __firk_eval_sol!(eval_sol, y_, mesh, cache)
     eval_bc_residual!(resids[1], pt, bc!, eval_sol, p, mesh)
     recursive_flatten!(resid, resids)
     return nothing
@@ -1470,10 +1478,10 @@ end
 end
 
 @views function __firk_loss(
-        u, p, y, pt::StandardBVProblem, bc::BC, mesh, cache, _, trait
+        u, p, y, pt::StandardBVProblem, bc::BC, mesh, cache, eval_sol, trait
     ) where {BC}
     y_ = recursive_unflatten!(y, u)
-    eval_sol = EvalSol(__restructure_sol(y_, cache.in_size), mesh, cache)
+    eval_sol = __firk_eval_sol!(eval_sol, y_, mesh, cache)
     resid_bc = eval_bc_residual(pt, bc, eval_sol, p, mesh)
     resid_co = Φ(cache, y_, u, trait)
     return vcat(resid_bc, mapreduce(vec, vcat, resid_co))