Add simpler type aliases and constructors for IIP FunctionWrappers

Add named type aliases for the FunctionWrappersWrapper types used in
the VF64 pattern (DifferentialEquations.jl#1128):

- `IIPFunctionWrapper{duType, uType, pType, tType}`: 1-wrapper alias
  for solvers without ForwardDiff (e.g. Tsit5, Verner)
- `IIPFunctionWrapperVF64{pType}`: VF64-specialized version
- `wrapfun_iip_simple(ff, du, u, p, t)`: constructor that avoids
  ForwardDiff extension backedges
- `IIPFunctionWrapperForwardDiff{T1,T2,T3,T4,dT1,dT2,dT4}`: 4-wrapper
  alias for ForwardDiff-aware solvers (Rosenbrock, implicit methods)
- `IIPFunctionWrapperForwardDiffVF64{pType}`: VF64-specialized version
- `ODEDualTag`, `ODEDualType`: named aliases for the ODE ForwardDiff
  tag and dual number types

These aliases enable the VF64 pattern in downstream packages
(OrdinaryDiffEq, StochasticDiffEq, etc.) to reference these types
by name for struct field annotations and pattern matching, instead
of spelling out the full FunctionWrappersWrapper type parameters.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: ChrisRackauckas

ext/DiffEqBaseForwardDiffExt.jl

@@ -5,13 +5,100 @@ using DiffEqBase.ArrayInterface
 using DiffEqBase: Void, FunctionWrappersWrappers, OrdinaryDiffEqTag,
     AbstractTimeseriesSolution,
     RecursiveArrayTools, reduce_tup, _promote_tspan, has_continuous_callback
+using FunctionWrappers: FunctionWrapper
 import DiffEqBase: hasdualpromote, wrapfun_oop, wrapfun_iip, prob2dtmin,
     promote_tspan, ODE_DEFAULT_NORM
 import SciMLBase: isdualtype, DualEltypeChecker, sse, __sum
 
 const dualT = ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEqTag, Float64}, Float64, 1}
 dualgen(::Type{T}) where {T} = ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEqTag, T}, T, 1}
 
+"""
+    ODEDualTag
+
+Type alias for the ForwardDiff tag used by ODE solvers:
+`ForwardDiff.Tag{OrdinaryDiffEqTag, Float64}`.
+
+Access from downstream packages via:
+```julia
+ext = Base.get_extension(DiffEqBase, :DiffEqBaseForwardDiffExt)
+ext.ODEDualTag
+```
+"""
+const ODEDualTag = ForwardDiff.Tag{OrdinaryDiffEqTag, Float64}
+
+"""
+    ODEDualType
+
+Type alias for the ForwardDiff dual number used by ODE solvers:
+`ForwardDiff.Dual{ODEDualTag, Float64, 1}`.
+
+Access from downstream packages via:
+```julia
+ext = Base.get_extension(DiffEqBase, :DiffEqBaseForwardDiffExt)
+ext.ODEDualType
+```
+"""
+const ODEDualType = ForwardDiff.Dual{ODEDualTag, Float64, 1}
+
+"""
+    IIPFunctionWrapperForwardDiff{T1, T2, T3, T4, dT1, dT2, dT4}
+
+Type alias for a `FunctionWrappersWrapper` containing 4 `FunctionWrapper` entries
+for an in-place function `f!(du, u, p, t) -> Nothing` with ForwardDiff support.
+
+The 4 wrappers cover:
+1. Base types: `(T1, T2, T3, T4)`
+2. Dual state: `(dT1, dT2, T3, T4)`
+3. Dual time: `(dT1, T2, T3, dT4)`
+4. Dual state+time: `(dT1, dT2, T3, dT4)`
+
+Used by solvers that require ForwardDiff internally (e.g. Rosenbrock, implicit methods).
+See also [`DiffEqBase.IIPFunctionWrapper`](@ref) for the simpler non-ForwardDiff variant.
+
+Access from downstream packages via:
+```julia
+ext = Base.get_extension(DiffEqBase, :DiffEqBaseForwardDiffExt)
+ext.IIPFunctionWrapperForwardDiff
+```
+"""
+const IIPFunctionWrapperForwardDiff{T1, T2, T3, T4, dT1, dT2, dT4} =
+    FunctionWrappersWrappers.FunctionWrappersWrapper{
+        Tuple{
+            FunctionWrapper{Nothing, Tuple{T1, T2, T3, T4}},
+            FunctionWrapper{Nothing, Tuple{dT1, dT2, T3, T4}},
+            FunctionWrapper{Nothing, Tuple{dT1, T2, T3, dT4}},
+            FunctionWrapper{Nothing, Tuple{dT1, dT2, T3, dT4}},
+        },
+        false,
+    }
+
+"""
+    IIPFunctionWrapperForwardDiffVF64{pType}
+
+VF64-specialized alias for `IIPFunctionWrapperForwardDiff` matching the common
+in-place `Vector{Float64}` ODE case with ForwardDiff support.
+
+Equivalent to:
+```julia
+IIPFunctionWrapperForwardDiff{
+    Vector{Float64}, Vector{Float64}, pType, Float64,
+    Vector{ODEDualType}, Vector{ODEDualType}, ODEDualType,
+}
+```
+
+Access from downstream packages via:
+```julia
+ext = Base.get_extension(DiffEqBase, :DiffEqBaseForwardDiffExt)
+ext.IIPFunctionWrapperForwardDiffVF64
+```
+"""
+const IIPFunctionWrapperForwardDiffVF64{pType} =
+    IIPFunctionWrapperForwardDiff{
+        Vector{Float64}, Vector{Float64}, pType, Float64,
+        Vector{dualT}, Vector{dualT}, dualT,
+    }
+
 const NORECOMPILE_IIP_SUPPORTED_ARGS = (
     Tuple{
         Vector{Float64}, Vector{Float64},

src/DiffEqBase.jl

@@ -175,6 +175,9 @@ export initialize!, finalize!
 
 export SensitivityADPassThrough
 
+# FunctionWrapper type aliases and constructors for the VF64 pattern
+export IIPFunctionWrapper, IIPFunctionWrapperVF64, wrapfun_iip_simple
+
 include("precompilation.jl")
 
 end # module

src/norecompile.jl

@@ -21,6 +21,47 @@ function unwrap_fw(fw::FunctionWrapper)
     return fw.obj[]
 end
 
+"""
+    IIPFunctionWrapper{duType, uType, pType, tType}
+
+Type alias for a `FunctionWrappersWrapper` containing a single `FunctionWrapper`
+for an in-place function `f!(du, u, p, t) -> Nothing`.
+
+Used by solvers that do **not** require ForwardDiff internally (e.g. Tsit5, Verner).
+See also [`IIPFunctionWrapperForwardDiff`](@ref) for the ForwardDiff-aware variant.
+"""
+const IIPFunctionWrapper{duType, uType, pType, tType} =
+    FunctionWrappersWrappers.FunctionWrappersWrapper{
+        Tuple{FunctionWrapper{Nothing, Tuple{duType, uType, pType, tType}}},
+        false,
+    }
+
+"""
+    IIPFunctionWrapperVF64{pType}
+
+VF64-specialized alias: `IIPFunctionWrapper{Vector{Float64}, Vector{Float64}, pType, Float64}`.
+Matches the wrapper produced for the common in-place `Vector{Float64}` ODE case
+when ForwardDiff is **not** used by the solver.
+"""
+const IIPFunctionWrapperVF64{pType} =
+    IIPFunctionWrapper{Vector{Float64}, Vector{Float64}, pType, Float64}
+
+"""
+    wrapfun_iip_simple(ff, du, u, p, t)
+
+Wrap an in-place function `ff(du, u, p, t) -> Nothing` into an [`IIPFunctionWrapper`](@ref)
+(single `FunctionWrapper`, no ForwardDiff support).
+
+Unlike [`wrapfun_iip`](@ref), this function is **not** overridden by the ForwardDiff
+extension, so it avoids creating method-table backedges that cause invalidation.
+Use this when the solver does not need ForwardDiff internally.
+"""
+function wrapfun_iip_simple(ff, du, u, p, t)
+    return FunctionWrappersWrappers.FunctionWrappersWrapper(
+        Void(ff), (typeof((du, u, p, t)),), (Nothing,)
+    )
+end
+
 # Default dispatch assumes no ForwardDiff, gets added in the new dispatch
 function wrapfun_iip(ff, inputs)
     return FunctionWrappersWrappers.FunctionWrappersWrapper(

src/solve.jl

@@ -836,12 +836,7 @@ function promote_f(f::F, ::Val{specialize}, u0, p, t, ::Val{false}) where {F, sp
                     !(f.f isa FunctionWrappersWrappers.FunctionWrappersWrapper)
             )
         )
-        return unwrapped_f(
-            f,
-            FunctionWrappersWrappers.FunctionWrappersWrapper(
-                Void(f.f), (typeof((u0, u0, p, t)),), (Nothing,)
-            )
-        )
+        return unwrapped_f(f, wrapfun_iip_simple(f.f, u0, u0, p, t))
     else
         return f
     end

test/function_wrapper_aliases.jl

@@ -0,0 +1,106 @@
+using DiffEqBase, ForwardDiff, Test
+using DiffEqBase: Void, FunctionWrappersWrappers, OrdinaryDiffEqTag,
+    wrapfun_iip, wrapfun_iip_simple
+using FunctionWrappers: FunctionWrapper
+
+# Get the ForwardDiff extension module
+const FDExt = Base.get_extension(DiffEqBase, :DiffEqBaseForwardDiffExt)
+
+@testset "IIPFunctionWrapper type alias (no ForwardDiff)" begin
+    ff = (du, u, p, t) -> (du .= u; nothing)
+    du = zeros(3)
+    u = ones(3)
+    p = [1.0, 2.0]
+    t = 0.0
+
+    # wrapfun_iip_simple should produce an IIPFunctionWrapper
+    wrapped = wrapfun_iip_simple(ff, du, u, p, t)
+    @test wrapped isa DiffEqBase.IIPFunctionWrapper
+    @test wrapped isa DiffEqBase.IIPFunctionWrapper{
+        Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64}
+
+    # VF64 alias should match
+    @test wrapped isa DiffEqBase.IIPFunctionWrapperVF64{Vector{Float64}}
+
+    # The wrapped function should be callable and produce correct results
+    du_test = zeros(3)
+    wrapped.fw[1](du_test, u, p, t)
+    @test du_test == u
+end
+
+@testset "IIPFunctionWrapperVF64 with NullParameters" begin
+    ff = (du, u, p, t) -> (du .= u; nothing)
+    du = zeros(3)
+    u = ones(3)
+    p = SciMLBase.NullParameters()
+    t = 0.0
+
+    wrapped = wrapfun_iip_simple(ff, du, u, p, t)
+    @test wrapped isa DiffEqBase.IIPFunctionWrapper
+    @test wrapped isa DiffEqBase.IIPFunctionWrapperVF64{SciMLBase.NullParameters}
+end
+
+@testset "ODEDualTag and ODEDualType (ForwardDiff extension)" begin
+    @test FDExt.ODEDualTag === ForwardDiff.Tag{OrdinaryDiffEqTag, Float64}
+    @test FDExt.ODEDualType === ForwardDiff.Dual{
+        ForwardDiff.Tag{OrdinaryDiffEqTag, Float64}, Float64, 1}
+end
+
+@testset "IIPFunctionWrapperForwardDiff type alias" begin
+    ff = (du, u, p, t) -> (du .= u; nothing)
+    du = zeros(3)
+    u = ones(3)
+    p = [1.0, 2.0]
+    t = 0.0
+
+    # wrapfun_iip with ForwardDiff loaded should produce a 4-wrapper
+    wrapped = wrapfun_iip(ff, (du, u, p, t))
+    @test wrapped isa FDExt.IIPFunctionWrapperForwardDiff
+
+    # VF64 alias should match
+    @test wrapped isa FDExt.IIPFunctionWrapperForwardDiffVF64{Vector{Float64}}
+
+    # The wrapped function should be callable
+    du_test = zeros(3)
+    wrapped.fw[1](du_test, u, p, t)
+    @test du_test == u
+end
+
+@testset "IIPFunctionWrapperForwardDiffVF64 with NullParameters" begin
+    ff = (du, u, p, t) -> (du .= u; nothing)
+
+    # Default wrapfun_iip (no args) produces the 7-wrapper variant, not 4-wrapper
+    wrapped_default = wrapfun_iip(ff)
+    # The default 7-wrapper has a different structure (7 entries, not 4),
+    # so it should NOT match the 4-wrapper alias
+    @test !(wrapped_default isa FDExt.IIPFunctionWrapperForwardDiff)
+
+    # But with explicit 4-tuple args and NullParameters, it should match
+    du = zeros(3)
+    u = ones(3)
+    p = SciMLBase.NullParameters()
+    t = 0.0
+    wrapped = wrapfun_iip(ff, (du, u, p, t))
+    @test wrapped isa FDExt.IIPFunctionWrapperForwardDiff
+    @test wrapped isa FDExt.IIPFunctionWrapperForwardDiffVF64{SciMLBase.NullParameters}
+end
+
+@testset "wrapfun_iip_simple does not change behavior with ForwardDiff loaded" begin
+    ff = (du, u, p, t) -> (du .= u; nothing)
+    du = zeros(3)
+    u = ones(3)
+    p = [1.0, 2.0]
+    t = 0.0
+
+    # wrapfun_iip_simple should ALWAYS produce a 1-wrapper, even with ForwardDiff loaded
+    wrapped_simple = wrapfun_iip_simple(ff, du, u, p, t)
+    @test wrapped_simple isa DiffEqBase.IIPFunctionWrapper
+
+    # It should NOT match the ForwardDiff 4-wrapper alias
+    @test !(wrapped_simple isa FDExt.IIPFunctionWrapperForwardDiff)
+
+    # wrapfun_iip should produce a 4-wrapper (ForwardDiff ext overrides)
+    wrapped_fd = wrapfun_iip(ff, (du, u, p, t))
+    @test wrapped_fd isa FDExt.IIPFunctionWrapperForwardDiff
+    @test !(wrapped_fd isa DiffEqBase.IIPFunctionWrapper)
+end

test/runtests.jl

@@ -36,6 +36,7 @@ end
         @time @safetestset "Norm" include("norm.jl")
         @time @safetestset "Utils" include("utils.jl")
         @time @safetestset "ForwardDiff Dual Detection" include("forwarddiff_dual_detection.jl")
+        @time @safetestset "FunctionWrapper Aliases" include("function_wrapper_aliases.jl")
         @time @safetestset "ODE default norm" include("ode_default_norm.jl")
         @time @safetestset "ODE default unstable check" include("ode_default_unstable_check.jl")
         @time @safetestset "Problem Kwargs Merging" include("problem_kwargs_merging.jl")