Handle circular references with-in mutable structs

format self-refrential (squash me into prev)

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

src/tangent_types/abstract_zero.jl

@@ -93,7 +93,7 @@ arguments.
 struct NoTangent <: AbstractZero end
 
 """
-    zero_tangent(primal)
+    zero_tangent(primal, _cache=nothing)
 
 This returns an appropriate zero tangent suitable for accumulating tangents of the primal.
 For mutable composites types this is a structural [`MutableTangent`](@ref)
@@ -107,55 +107,77 @@ In general, it is more likely to produce a structural tangent.
     `zero_tangent`is an experimental feature, and is part of the mutation support featureset.
     While this notice remains it may have changes in behavour, and interface in any _minor_ version of ChainRulesCore.
     Exactly how it should be used (e.g. is it forward-mode only?)
+
+The `_cache=nothing` is an internal implementation detail that the user should never need to set.
+(It is used to hold references to tangents for that might appear in self-referential structures)
 """
 function zero_tangent end
 
-zero_tangent(x::Number) = zero(x)
+zero_tangent(x::Number, _cache=nothing) = zero(x)
 
-zero_tangent(::Type) = NoTangent()
+zero_tangent(::Type, _cache=nothing) = NoTangent()
 
-function zero_tangent(x::MutableTangent{P}) where {P}
-    zb = backing(zero_tangent(backing(x)))
+function zero_tangent(x::MutableTangent{P}, _cache=nothing) where {P}
+    zb = backing(zero_tangent(backing(x), _cache))
     return MutableTangent{P}(zb)
 end
 
-function zero_tangent(x::Tangent{P}) where {P}
-    zb = backing(zero_tangent(backing(x)))
+function zero_tangent(x::Tangent{P}, _cache=nothing) where {P}
+    zb = backing(zero_tangent(backing(x), _cache))
     return Tangent{P,typeof(zb)}(zb)
 end
 
-@generated function zero_tangent(primal)
+@generated function zero_tangent(primal, _cache=nothing)
     fieldcount(primal) == 0 && return NoTangent()  # no tangent space at all, no need for structural zero.
     zfield_exprs = map(fieldnames(primal)) do fname
-        fval = :(
+        :(
             if isdefined(primal, $(QuoteNode(fname)))
-                zero_tangent(getfield(primal, $(QuoteNode(fname))))
+                zero_tangent(getfield(primal, $(QuoteNode(fname))), _cache)
             else
                 # This is going to be potentially bad, but that's what they get for not giving us a primal
                 # This will never me mutated inplace, rather it will alway be replaced with an actual value first
                 ZeroTangent()
             end
         )
-        Expr(:kw, fname, fval)
     end
     return if has_mutable_tangent(primal)
-        any_mask = map(fieldnames(primal), fieldtypes(primal)) do fname, ftype
-            # If it is is unassigned, or if it doesn't have a concrete type, let it take any value for its tangent
-            fdef = :(!isdefined(primal, $(QuoteNode(fname))) || !isconcretetype($ftype))
-            Expr(:kw, fname, fdef)
+        # This is a little complex because we need to support-self referential types
+        # So we need to:
+        # 1. create the tangent,
+        # 2. put it in the cache
+        # 3. Do all the calls to create the zeros for the fields giving them that cache)
+        # 4. put those zeros into the object
+        tangent_types = map(guess_zero_tangent_type, fieldtypes(primal))
+        is_defined_mask = Expr(:tuple, map(fieldnames(primal)) do fname
+            :(isdefined(primal, $(QuoteNode(fname))))
+        end...)
+
+        quote
+            isnothing(_cache) && (_cache = IdDict())
+            found_tangent = get(_cache, primal, nothing)
+            !isnothing(found_tangent) && return found_tangent
+
+            # Now we need to put into the cache a placeholder tangent so we can construct our fields using that cache
+            # then put those fields into the placeholder
+            tangent = $_MutableTangent(Val{$primal}(), $is_defined_mask, $tangent_types)
+            _cache[primal] = tangent
+            $(
+                map(fieldnames(primal), zfield_exprs) do fname, fval_expr
+                    :(setproperty!(tangent, $(QuoteNode(fname)), $fval_expr))
+                end...
+            )
+            return tangent            
         end
-        :($MutableTangent{$primal}(
-            $(Expr(:tuple, Expr(:parameters, any_mask...))),
-            $(Expr(:tuple, Expr(:parameters, zfield_exprs...))),
-        ))
     else
-        :($Tangent{$primal}($(Expr(:parameters, zfield_exprs...))))
+        :($Tangent{$primal}($(Expr(:parameters, Expr.(:kw, fieldnames(primal), zfield_exprs)...))))
     end
 end
 
-zero_tangent(primal::Tuple) = Tangent{typeof(primal)}(map(zero_tangent, primal)...)
+function zero_tangent(primal::Tuple, _cache=nothing)
+    return Tangent{typeof(primal)}(map(x -> zero_tangent(x, _cache), primal)...)
+end
 
-function zero_tangent(x::Array{P,N}) where {P,N}
+function zero_tangent(x::Array{P,N}, _cache=nothing) where {P,N}
     if (isbitstype(P) || all(i -> isassigned(x, i), eachindex(x)))
         return map(zero_tangent, x)
     end
@@ -165,16 +187,20 @@ function zero_tangent(x::Array{P,N}) where {P,N}
     y = Array{guess_zero_tangent_type(P),N}(undef, size(x)...)
     @inbounds for n in eachindex(y)
         if isassigned(x, n)
-            y[n] = zero_tangent(x[n])
+            y[n] = zero_tangent(x[n], _cache)
         end
     end
     return y
 end
 
-# Sad heauristic methods we need because of unassigned values
-guess_zero_tangent_type(::Type{T}) where {T<:Number} = T
-guess_zero_tangent_type(::Type{T}) where {T<:Integer} = typeof(float(zero(T)))
+# Sad heauristic methods
+#guess_zero_tangent_type(::Type{T}) where {T<:Number} = T
+#guess_zero_tangent_type(::Type{T}) where {T<:Integer} = typeof(float(zero(T)))
 function guess_zero_tangent_type(::Type{<:Array{T,N}}) where {T,N}
     return Array{guess_zero_tangent_type(T),N}
 end
-guess_zero_tangent_type(T::Type) = Any
+
+# The following will fall back to `Any` if it is hard to infer
+function guess_zero_tangent_type(::Type{T}) where {T}
+    return Core.Compiler.return_type(zero_tangent, Tuple{T})
+end

src/tangent_types/structural_tangent.jl

@@ -54,6 +54,7 @@ struct Tangent{P,T} <: StructuralTangent{P}
     end
 end
 
+function _MutableTangent end
 """
     MutableTangent{P}(fields) <: StructuralTangent{P} <: AbstractTangent
 
@@ -73,6 +74,23 @@ It itself is also mutable.
 struct MutableTangent{P,F} <: StructuralTangent{P}
     backing::F
 
+    # Uninitialized constructor
+    global function _MutableTangent(::Val{P}, is_defined_mask, tangent_types) where {P}
+        backing_vals = map(is_defined_mask, tangent_types) do is_def, tangent_type
+            ref = if !is_def
+                Ref{Union{ZeroTangent, tangent_type}}  # allow a Zero which will be used for uninitialized values
+            else
+                Ref{tangent_type}
+            end
+            return ref() # undefined, but it will be filled later
+        end
+        backing = NamedTuple{fieldnames(P)}(backing_vals)
+        return new{P, typeof(backing)}(backing)
+    end
+
+    # TODO: are the following two correct?
+    # Are they useful?
+    # The place they are used is just `map`, maybe we should instead just copy types the thing being mapped?
     function MutableTangent{P}(
         any_mask::NamedTuple{names, <:NTuple{<:Any, Bool}}, fvals::NamedTuple{names}
     ) where {names, P}
@@ -88,6 +106,7 @@ struct MutableTangent{P,F} <: StructuralTangent{P}
         return new{P, typeof(backing)}(backing)
     end
 
+
     function MutableTangent{P}(fvals) where P
         any_mask = NamedTuple{fieldnames(P)}((!isconcretetype).(fieldtypes(P)))
         return MutableTangent{P}(any_mask, fvals)

test/tangent_types/abstract_zero.jl

@@ -174,7 +174,7 @@ end
         @test zero_tangent(MutDemo(1.5)) isa MutableTangent{MutDemo}
         @test iszero(zero_tangent(MutDemo(1.5)))
 
-        @test zero_tangent((; a=1)) isa Tangent{typeof((; a = 1))}
+        @test zero_tangent((; a=1.3)) isa Tangent{typeof((; a = 1.3))}
         @test zero_tangent(Demo(1.2)) isa Tangent{Demo}
         @test zero_tangent(Demo(1.2)).x === 0.0
 
@@ -275,4 +275,41 @@ end
         @test d.z == [2.0, 3.0]
         @test d.z isa SubArray
     end
+
+    @testset "cyclic references" begin
+        mutable struct Link
+            data::Float64
+            next::Link
+            Link(data) = new(data)
+        end
+
+        lk = Link(1.5)
+        lk.next = lk
+
+        d = zero_tangent(lk)
+        @test d.data == 0.0
+        @test d.next === d
+
+        # The following two cases are broken
+        # We hope they are not too significant, because in general if you AD step by step they should work
+        # (as should the one above so maybe we should get rid of this extra complex logic)
+        # It's only a problem if you first do the multistep build then `zero_tangent` rather than `zero_tangent` at the constructor.
+
+        # Idea: check if `!isbitstype` only if so do we need to worry about caching etc
+        struct CarryingArray
+            x::Vector
+        end
+        ca = CarryingArray(Any[1.5])
+        push!(ca.x, ca)
+        @test_broken d_ca = zero_tangent(ca)
+        @test_broken d_ca[1] == 0.0
+        @test_broken d_ca[2] === _ca
+
+        # Idea: check if typeof(xs) <: eltype(xs), if so need to cache it before computing
+        xs = Any[1.5]
+        push!(xs, xs)
+        @test_broken d_xs = zero_tangent(xs)
+        @test_broken d_xs[1] == 0.0
+        @test_broken d_xs[2] == d_xs
+    end
 end