Use flat Vector u0 for MIRK/FIRK/MIRKN nlprob to preserve inference

@wsmoses noted that NonlinearSolve's polyalg is type-stable (PR #870).
Verified: a `NonlinearProblem` with a plain `Vector{Float64}` u0 does
infer concretely under `solve(prob)`. The widening I was seeing on
LTS / 1.11 CI is specifically triggered by the **type** of `nlprob.u0`:
under RecursiveArrayTools v4, `vec(::VectorOfArray)` returns a
`Base.ReshapedArray{T, 1, VectorOfArray{...}, …}` (in v3 it returned a
plain `Vector`). NonlinearSolve's polyalg cannot infer `T, N, uType, R`
of the resulting `NonlinearSolution` when `u0` is that ReshapedArray,
so the inferred return type widens to `NonlinearSolution{_A,_B,_C,_D, …}`
and `@inferred solve(::BVProblem, ::MIRKN, …)` rejects it.

Minimal repro (no BVDE involved):

    using NonlinearSolve, RecursiveArrayTools
    f!(du, u, p) = (du .= u .^ 2 .- p)

    voa = VectorOfArray([[1.0], [1.0], [1.0]])
    prob_resh = NonlinearProblem(NonlinearFunction(f!), vec(voa), [2.0,2.0,2.0])
    Base.return_types(solve, Tuple{typeof(prob_resh)})
    # → NonlinearSolution{_A, _B, _C, _D, NonlinearProblem{ReshapedArray{…}, …}, …}

    prob_flat = NonlinearProblem(NonlinearFunction(f!), reduce(vcat, voa.u), [2.0,2.0,2.0])
    Base.return_types(solve, Tuple{typeof(prob_flat)})
    # → NonlinearSolution{Float64, 1, Vector{Float64}, Vector{Float64}, …}   (concrete)

So this is a downstream RAT-v4 inference cliff, not a polyalg bug.

Fix: in MIRK / FIRK / MIRKN's `__perform_*_iteration`, flatten the
VectorOfArray initial guess to a plain `Vector` before constructing
`nlprob`:

    nlprob = __construct_problem(cache, reduce(vcat, cache.y₀.u), copy(cache.y₀))

`recursive_unflatten!(cache.y₀, sol_nlprob.u)` afterwards still copies
the result back, so the loss of view-aliasing on the (always-copied-
internally-by-the-solver) `u0` is benign.

Also revert the previous workaround that commented out the MIRK
`@inferred solve` blocks in `test/misc/type_stability_tests.jl`; with
this fix they pass on Julia 1.10/1.11/1.12.

Bumps BoundaryValueDiffEqMIRKN patch version to 1.15.1.

Local verification on Julia 1.11.9:

    === Test 2: @inferred ===
    PASS

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

Author: ChrisRackauckas

lib/BoundaryValueDiffEqFIRK/src/firk.jl

@@ -552,7 +552,11 @@ function SciMLBase.solve!(
 end
 
 function __perform_firk_iteration(cache::Union{FIRKCacheExpand, FIRKCacheNested}, abstol, adaptive::Bool)
-    nlprob = __construct_problem(cache, vec(cache.y₀), copy(cache.y₀))
+    # See note in BoundaryValueDiffEqMIRK/src/mirk.jl: `vec(::VectorOfArray)`
+    # returns a `ReshapedArray` under RAT v4, which widens the inferred
+    # return type of NonlinearSolve's polyalg. Flatten via `.u` to keep
+    # `nlprob.u0` a plain `Vector` and preserve inference.
+    nlprob = __construct_problem(cache, reduce(vcat, cache.y₀.u), copy(cache.y₀))
     solve_alg = __concrete_solve_algorithm(nlprob, cache.alg.nlsolve, cache.alg.optimize)
     kwargs = __concrete_kwargs(
         cache.alg.nlsolve, cache.alg.optimize, cache.nlsolve_kwargs, cache.optimize_kwargs,

lib/BoundaryValueDiffEqMIRK/src/mirk.jl

@@ -308,7 +308,16 @@ function SciMLBase.solve!(
 end
 
 function __perform_mirk_iteration(cache::MIRKCache, abstol, adaptive::Bool, controller::AbstractErrorControl)
-    nlprob = __construct_problem(cache, vec(cache.y₀), copy(cache.y₀))
+    # Flatten via `reduce(vcat, .u)` rather than `vec(::VectorOfArray)`. Under
+    # RecursiveArrayTools v4, `vec` on a `VectorOfArray` returns a
+    # `ReshapedArray{T, 1, VectorOfArray{...}}`, whose presence as
+    # `nlprob.u0` widens the inferred return type of `solve(nlprob, polyalg)`
+    # — `T`, `N`, `uType`, `R` of the resulting `NonlinearSolution` come back
+    # as free type variables. A plain `Vector{Float64}` u0 lets NonlinearSolve's
+    # polyalg infer its return type concretely. (Behaviorally equivalent: the
+    # solver copies u0 internally and `recursive_unflatten!` writes the result
+    # back into `cache.y₀` afterwards.)
+    nlprob = __construct_problem(cache, reduce(vcat, cache.y₀.u), copy(cache.y₀))
     solve_alg = __concrete_solve_algorithm(nlprob, cache.alg.nlsolve, cache.alg.optimize)
     kwargs = __concrete_kwargs(
         cache.alg.nlsolve, cache.alg.optimize, cache.nlsolve_kwargs, cache.optimize_kwargs,

lib/BoundaryValueDiffEqMIRKN/Project.toml

@@ -1,7 +1,7 @@
 name = "BoundaryValueDiffEqMIRKN"
 uuid = "9255f1d6-53bf-473e-b6bd-23f1ff009da4"
 authors = ["Qingyu Qu <erikqqy123@gmail.com>"]
-version = "1.15.0"
+version = "1.15.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

lib/BoundaryValueDiffEqMIRKN/src/mirkn.jl

@@ -137,7 +137,11 @@ function SciMLBase.solve!(cache::MIRKNCache{iip, T}) where {iip, T}
 end
 
 function __perform_mirkn_iteration(cache::MIRKNCache)
-    nlprob = __construct_nlproblem(cache, vec(cache.y₀), copy(cache.y₀))
+    # See note in BoundaryValueDiffEqMIRK/src/mirk.jl: `vec(::VectorOfArray)`
+    # returns a `ReshapedArray` under RAT v4, which widens the inferred
+    # return type of NonlinearSolve's polyalg. Flatten via `.u` to keep
+    # `nlprob.u0` a plain `Vector` and preserve inference.
+    nlprob = __construct_nlproblem(cache, reduce(vcat, cache.y₀.u), copy(cache.y₀))
     solve_alg = __concrete_solve_algorithm(nlprob, cache.alg.nlsolve, cache.alg.optimize)
     kwargs = __concrete_kwargs(
         cache.alg.nlsolve, cache.alg.optimize, cache.nlsolve_kwargs, cache.optimize_kwargs,

test/misc/type_stability_tests.jl

@@ -51,20 +51,15 @@
         #     @inferred solve(mpbvp_oop, MultipleShooting(5, Tsit5(); jac_alg))
         # end
 
-        # MIRK methods can hit the same NonlinearSolve polyalg type instability
-        # as Shooting once adaptive mesh refinement is exercised — the inner
-        # NonlinearSolution type depends on which solver in the polyalg succeeds,
-        # which is not statically resolvable. Skipped pending the same upstream
-        # investigation noted above for Shooting.
-        # @testset "MIRK Methods" begin
-        #     for solver in (
-        #             MIRK2(; jac_alg), MIRK3(; jac_alg), MIRK4(; jac_alg),
-        #             MIRK5(; jac_alg), MIRK6(; jac_alg),
-        #         )
-        #         @inferred solve(mpbvp_iip, solver; dt = 0.2)
-        #         @inferred solve(mpbvp_oop, solver; dt = 0.2)
-        #     end
-        # end
+        @testset "MIRK Methods" begin
+            for solver in (
+                    MIRK2(; jac_alg), MIRK3(; jac_alg), MIRK4(; jac_alg),
+                    MIRK5(; jac_alg), MIRK6(; jac_alg),
+                )
+                @inferred solve(mpbvp_iip, solver; dt = 0.2)
+                @inferred solve(mpbvp_oop, solver; dt = 0.2)
+            end
+        end
     end
 
     # Two-Point BVP
@@ -84,15 +79,14 @@
         #     @inferred solve(tpbvp_oop, MultipleShooting(5, Tsit5(); jac_alg))
         # end
 
-        # See note above about MIRK + NonlinearSolve polyalg.
-        # @testset "MIRK Methods" begin
-        #     for solver in (
-        #             MIRK2(; jac_alg), MIRK3(; jac_alg), MIRK4(; jac_alg),
-        #             MIRK5(; jac_alg), MIRK6(; jac_alg),
-        #         )
-        #         @inferred solve(tpbvp_iip, solver; dt = 0.2)
-        #         @inferred solve(tpbvp_oop, solver; dt = 0.2)
-        #     end
-        # end
+        @testset "MIRK Methods" begin
+            for solver in (
+                    MIRK2(; jac_alg), MIRK3(; jac_alg), MIRK4(; jac_alg),
+                    MIRK5(; jac_alg), MIRK6(; jac_alg),
+                )
+                @inferred solve(tpbvp_iip, solver; dt = 0.2)
+                @inferred solve(tpbvp_oop, solver; dt = 0.2)
+            end
+        end
     end
 end