Make unpacking a Tuple or NamedTuple Composite type-inferrable

Project.toml

@@ -1,13 +1,15 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.9.25"
+version = "0.9.26"
 
 [deps]
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 BenchmarkTools = "0.5"
+Compat = "2, 3"
 FiniteDifferences = "0.10"
 StaticArrays = "0.11, 0.12"
 julia = "1"

src/ChainRulesCore.jl

@@ -2,6 +2,7 @@ module ChainRulesCore
 using Base.Broadcast: broadcasted, Broadcasted, broadcastable, materialize, materialize!
 using LinearAlgebra: LinearAlgebra
 using SparseArrays: SparseVector, SparseMatrixCSC
+using Compat: hasfield
 
 export on_new_rule, refresh_rules # generation tools
 export frule, rrule # core function

src/differentials/composite.jl

@@ -73,10 +73,9 @@ Base.getindex(comp::Composite, idx) = getindex(backing(comp), idx)
 # for Tuple
 Base.getproperty(comp::Composite, idx::Int) = unthunk(getproperty(backing(comp), idx))
 function Base.getproperty(
- comp::Composite{P, <:NamedTuple{L}}, idx::Symbol
-) where {P, L}
- # Need to check L directly, or else this does not constant-fold
- idx ∈ L || return Zero()
+ comp::Composite{P, T}, idx::Symbol
+) where {P, T<:NamedTuple}
+ hasfield(T, idx) || return Zero()
  return unthunk(getproperty(backing(comp), idx))
 end
 
@@ -96,6 +95,10 @@ function Base.reverse(comp::Composite)
  Composite{typeof(rev_backing), typeof(rev_backing)}(rev_backing)
 end
 
+function Base.indexed_iterate(comp::Composite{P,<:Tuple}, i::Int, state=1) where {P}
+ return Base.indexed_iterate(backing(comp), i, state)
+end
+
 function Base.map(f, comp::Composite{P, <:Tuple}) where P
  vals::Tuple = map(f, backing(comp))
  return Composite{P, typeof(vals)}(vals)

test/differentials/composite.jl

@@ -18,7 +18,6 @@ struct StructWithInvariant
  StructWithInvariant(x) = new(x, 2x)
 end
 
-
 @testset "Composite" begin
  @testset "empty types" begin
  @test typeof(Composite{Tuple{}}()) == Composite{Tuple{}, Tuple{}}
@@ -77,6 +76,22 @@ end
  # Testing iterate via collect
  @test collect(Composite{Foo}(x=2.5)) == [2.5]
  @test collect(Composite{Tuple{Float64,}}(2.0)) == [2.0]
+
+ # Test indexed_iterate
+ ctup = Composite{Tuple{Float64,Int64}}(2.0, 3)
+ _unpack2tuple = function(comp)
+ a, b = comp
+ return (a, b)
+ end
+ @inferred _unpack2tuple(ctup)
+ @test _unpack2tuple(ctup) === (2.0, 3)
+
+ # Test getproperty is inferrable
+ _unpacknamedtuple = comp -> (comp.x, comp.y)
+ if VERSION ≥ v"1.2"
+ @inferred _unpacknamedtuple(Composite{Foo}(x=2, y=3.0))
+ @inferred _unpacknamedtuple(Composite{Foo}(y=3.0))
+ end
  end
 
  @testset "reverse" begin