Handle first order constructors

Author: oxinabox

src/codegen/forward.jl

@@ -69,7 +69,7 @@ function fwd_transform!(ci, mi, nargs, N)
             if isa(stmt, Expr)
                 error("Unexprected statement encountered. This is a bug in Diffractor. stmt=$stmt")
             end
-            return Expr(:call, zero_bundle{order}(), stmt)
+            return Expr(:call, zero_bundle{N}(), stmt)
         end
     end
 

src/codegen/forward_demand.jl

@@ -105,6 +105,7 @@ function forward_diff_uncached!(ir::IRCode, interp::AbstractInterpreter, irsv::I
             Δbacking = insert_node!(ir, ssa, NewInstruction(Expr(:splatnew, widenconst(tup_typ), Δbacking), tup_typ_typ.val))
         end
         tangentT = Core.Compiler.apply_type_tfunc(Const(ChainRulesCore.Tangent), newT, tup_typ_typ).val
+        # TODO do we need to make sure this inserts right 
         Δtangent = insert_node!(ir, ssa, NewInstruction(Expr(:new, tangentT, Δbacking), tangentT))
         return Δtangent
     else # general frule handling

src/stage1/generated.jl

@@ -390,7 +390,7 @@ end
 lifted_getfield(x::ZeroTangent, s) = ZeroTangent()
 lifted_getfield(x::NoTangent, s) = NoTangent()
 
-lifted_getfield(x::Tangent, s) = getproperty(x, s)
+lifted_getfield(x::StructuralTangent, s) = getproperty(x, s)
 
 function lifted_getfield(x::Tangent{<:Tangent{T}}, s) where T
     bb = getfield(x.backing, 1)

src/stage1/recurse_fwd.jl

@@ -15,28 +15,29 @@ struct ∂☆new{N}; end
 function (::∂☆new{1})(B::Type, xs::AbstractTangentBundle{1}...)
     primal_args = map(primal, xs)
     the_primal = _construct(B, primal_args)
-
+    @info "∂☆new{1}"
     tangent_tup = map(first_partial, xs)
     the_partial = if B<:Tuple
         Tangent{B, typeof(tangent_tup)}(tangent_tup)
     else
         names = fieldnames(B)
         tangent_nt = NamedTuple{names}(tangent_tup)
-        Tangent{B, typeof(tangent_nt)}(tangent_nt)
+        StructuralTangent{B}(tangent_nt)
     end
     return TaylorBundle{1, B}(the_primal, (the_partial,))
 end
 
 function (::∂☆new{N})(B::Type, xs::AbstractTangentBundle{N}...) where {N}
     primal_args = map(primal, xs)
     the_primal = _construct(B, primal_args)
-        
+    @info "∂☆new{N}"
     the_partials = ntuple(Val{N}()) do ii
         iith_order_type = ii==1 ? B : Any  # the type of the higher order tangents isn't worth tracking
         tangent_tup = map(x->partial(x, ii), xs)
         tangent = if B<:Tuple
             Tangent{iith_order_type, typeof(tangent_tup)}(tangent_tup)
         else
+            # TODO support mutation
             names = fieldnames(B)
             tangent_nt = NamedTuple{names}(tangent_tup)
             Tangent{iith_order_type, typeof(tangent_nt)}(tangent_nt)

test/forward_mutation.jl

@@ -1,10 +1,13 @@
 using Diffractor
-using Diffractor: ∂☆
+using Diffractor: ∂☆, ZeroBundle, TaylorBundle
 using Diffractor: bundle
 
 mutable struct MDemo1
     x::Float64
 end
+
+∂☆{1}()(ZeroBundle{1}(MDemo1), TaylorBundle{1}(1.5, (1.0,)))
+
 function double!(val::MDemo1)
     val.x *= 2.0
     return val