Improve inferability of forward mode data structures (#104)

None of this is particularly, pretty, but these changes help the
compiler constprop things better, so let's do this for now.
We can revisit in the future when the compiler learns more tricks.

Author: Keno

src/extra_rules.jl

@@ -172,7 +172,7 @@ function ChainRules.rrule(::DiffractorRuleConfig, ::Type{SArray{S, T, N, L}}, x:
 end
 
 function ChainRules.frule((_, ∂x), ::Type{SArray{S, T, N, L}}, x::NTuple{L,T}) where {S, T, N, L}
-    SArray{S, T, N, L}(x), SArray{S, T, N, L}(∂x)
+    SArray{S, T, N, L}(x), SArray{S, T, N, L}(∂x.backing)
 end
 
 function ChainRules.frule((_, ∂x), ::Type{SArray{S, T, N, L}}, x::NTuple{L,Any}) where {S, T, N, L}

src/stage1/forward.jl

@@ -145,12 +145,24 @@ end
         map(x->lifted_getfield(x, s), x.tangent.partials))
 end
 
+@Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::TangentBundle{N}, s::ATB{N}, inbounds::ATB{N}) where {N}
+    s = primal(s)
+    ExplicitTangentBundle{N}(getfield(primal(x), s, primal(inbounds)),
+        map(x->lifted_getfield(x, s), x.tangent.partials))
+end
+
 @Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::TaylorBundle{N}, s::AbstractTangentBundle{N}) where {N}
     s = primal(s)
     TaylorBundle{N}(getfield(primal(x), s),
         map(y->lifted_getfield(y, s), x.tangent.coeffs))
 end
 
+@Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::TaylorBundle{N}, s::AbstractTangentBundle{N}, inbounds::ATB{N}) where {N}
+    s = primal(s)
+    TaylorBundle{N}(getfield(primal(x), s, primal(inbounds)),
+        map(y->lifted_getfield(y, s), x.tangent.coeffs))
+end
+
 @Base.constprop :aggressive function (::∂☆{N})(::ATB{N, typeof(getfield)}, x::CompositeBundle{N}, s::AbstractTangentBundle{N, Int}) where {N}
     x.tup[primal(s)]
 end
@@ -159,12 +171,6 @@ end
     x.tup[Base.fieldindex(B, primal(s))]
 end
 
-@Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::ATB{N}, s::ATB{N}, inbounds::ATB{N}) where {N}
-    s = primal(s)
-    ExplicitTangentBundle{N}(getfield(primal(x), s, primal(inbounds)),
-        map(x->lifted_getfield(x, s), x.tangent.partials))
-end
-
 @Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::UniformBundle{N, <:Any, U}, s::AbstractTangentBundle{N}) where {N, U}
     UniformBundle{N,<:Any,U}(getfield(primal(x), primal(s)), x.tangent.val)
 end

src/tangent.jl

@@ -90,6 +90,11 @@ struct ExplicitTangent{P <: Tuple} <: AbstractTangentSpace
     partials::P
 end
 
+@eval struct TaylorTangent{C <: Tuple} <: AbstractTangentSpace
+    coeffs::C
+    TaylorTangent(coeffs) = $(Expr(:new, :(TaylorTangent{typeof(coeffs)}), :coeffs))
+end
+
 """
     struct TaylorTangent{C}
 
@@ -113,9 +118,7 @@ by analogy with the (truncated) Taylor series
 
     c₀ + c₁ x + 1/2 c₂ x² + 1/3! c₃ x³ + O(x⁴)
 """
-struct TaylorTangent{C <: Tuple} <: AbstractTangentSpace
-    coeffs::C
-end
+TaylorTangent
 
 """
     struct ProductTangent{T <: Tuple{Vararg{AbstractTangentSpace}}}
@@ -137,20 +140,24 @@ struct UniformTangent{U} <: AbstractTangentSpace
     val::U
 end
 
+function _TangentBundle end
+
+@eval struct TangentBundle{N, B, P <: AbstractTangentSpace} <: AbstractTangentBundle{N, B}
+    primal::B
+    tangent::P
+    global _TangentBundle(::Val{N}, primal::B, tangent::P) where {N, B, P} = $(Expr(:new, :(TangentBundle{N, Core.Typeof(primal), typeof(tangent)}), :primal, :tangent))
+end
+
 """
     struct TangentBundle{N, B, P}
 
 Represents a tangent bundle as an explicit primal together
 with some representation of (potentially a product of) the tangent space.
 """
-struct TangentBundle{N, B, P <: AbstractTangentSpace} <: AbstractTangentBundle{N, B}
-    primal::B
-    tangent::P
-    TangentBundle{N, B, P}(primal::B, tangent::P) where {N, B, P} = new{N, B, P}(primal, tangent)
-end
+TangentBundle
 
 TangentBundle{N}(primal::B, tangent::P) where {N, B, P<:AbstractTangentSpace} =
-    TangentBundle{N, B, P}(primal, tangent)
+    _TangentBundle(Val{N}(), primal, tangent)
 
 const ExplicitTangentBundle{N, B, P} = TangentBundle{N, B, ExplicitTangent{P}}
 
@@ -159,17 +166,17 @@ check_tangent_invariant(lp, N) = @assert lp == 2^N - 1
 
 function ExplicitTangentBundle{N}(primal::B, partials::P) where {N, B, P}
     check_tangent_invariant(length(partials), N)
-    TangentBundle{N, Core.Typeof(primal), ExplicitTangent{P}}(primal, ExplicitTangent{P}(partials))
+    _TangentBundle(Val{N}(), primal, ExplicitTangent{P}(partials))
 end
 
 function ExplicitTangentBundle{N,B}(primal::B, partials::P) where {N, B, P}
     check_tangent_invariant(length(partials), N)
-    TangentBundle{N, B, ExplicitTangent{P}}(primal, ExplicitTangent{P}(partials))
+    _TangentBundle(Val{N}(), primal, ExplicitTangent{P}(partials))
 end
 
 function ExplicitTangentBundle{N,B,P}(primal::B, partials::P) where {N, B, P}
     check_tangent_invariant(length(partials), N)
-    TangentBundle{N, B, ExplicitTangent{P}}(primal, ExplicitTangent{P}(partials))
+    _TangentBundle(Val{N}(), primal, ExplicitTangent{P}(partials))
 end
 
 function Base.show(io::IO, x::ExplicitTangentBundle)
@@ -194,9 +201,9 @@ end
 
 const TaylorBundle{N, B, P} = TangentBundle{N, B, TaylorTangent{P}}
 
-function TaylorBundle{N, B}(primal::B, coeffs::P) where {N, B, P}
+function TaylorBundle{N, B}(primal::B, coeffs) where {N, B}
     check_taylor_invariants(coeffs, primal, N)
-    TangentBundle{N, B, TaylorTangent{P}}(primal, TaylorTangent{P}(coeffs))
+    _TangentBundle(Val{N}(), primal, TaylorTangent(coeffs))
 end
 
 function check_taylor_invariants(coeffs, primal, N)
@@ -208,7 +215,7 @@ end
 @ChainRulesCore.non_differentiable check_taylor_invariants(coeffs, primal, N)
 
 function TaylorBundle{N}(primal, coeffs) where {N}
-    TaylorBundle{N, Core.Typeof(primal)}(primal, coeffs)
+    _TangentBundle(Val{N}(), primal, TaylorTangent(coeffs))
 end
 
 function Base.show(io::IO, x::TaylorBundle{1})
@@ -224,12 +231,12 @@ function Base.getindex(tb::TaylorBundle, tti::CanonicalTangentIndex)
 end
 
 const UniformBundle{N, B, U} = TangentBundle{N, B, UniformTangent{U}}
-UniformBundle{N, B, U}(primal::B, partial::U) where {N,B,U} = UniformBundle{N, B, U}(primal, UniformTangent{U}(partial))
-UniformBundle{N, B, U}(primal::B) where {N,B,U} = UniformBundle{N, B, U}(primal, UniformTangent{U}(U.instance))
-UniformBundle{N, B}(primal::B, partial::U) where {N,B,U} = UniformBundle{N, Core.Typeof(primal), U}(primal, UniformTangent{U}(partial))
-UniformBundle{N}(primal, partial::U) where {N,U} = UniformBundle{N, Core.Typeof(primal), U}(primal, UniformTangent{U}(partial))
-UniformBundle{N, <:Any, U}(primal, partial::U) where {N, U} = UniformBundle{N, Core.Typeof(primal), U}(primal, UniformTangent{U}(U.instance))
-UniformBundle{N, <:Any, U}(primal) where {N, U} = UniformBundle{N, Core.Typeof(primal), U}(primal, UniformTangent{U}(U.instance))
+UniformBundle{N, B, U}(primal::B, partial::U) where {N,B,U} = _TangentBundle(Val{N}(), primal, UniformTangent{U}(partial))
+UniformBundle{N, B, U}(primal::B) where {N,B,U} = _TangentBundle(Val{N}(), primal, UniformTangent{U}(U.instance))
+UniformBundle{N, B}(primal::B, partial::U) where {N,B,U} = _TangentBundle(Val{N}(),primal, UniformTangent{U}(partial))
+UniformBundle{N}(primal, partial::U) where {N,U} = _TangentBundle(Val{N}(), primal, UniformTangent{U}(partial))
+UniformBundle{N, <:Any, U}(primal, partial::U) where {N, U} = _TangentBundle(Val{N}(), primal, UniformTangent{U}(U.instance))
+UniformBundle{N, <:Any, U}(primal) where {N, U} = _TangentBundle(Val{N}(), primal, UniformTangent{U}(U.instance))
 
 const ZeroBundle{N, B} = UniformBundle{N, B, ZeroTangent}
 const DNEBundle{N, B} = UniformBundle{N, B, NoTangent}