Handle Base.tail & friends (#567)

* handle Base.tail & friens

* empty cases

* tail on NamedTuples too

Author: mcabbott

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "1.15.2"
+version = "1.15.3"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

src/tangent_types/abstract_zero.jl

@@ -16,6 +16,10 @@ Base.iszero(::AbstractZero) = true
 Base.iterate(x::AbstractZero) = (x, nothing)
 Base.iterate(::AbstractZero, ::Any) = nothing
 
+Base.first(x::AbstractZero) = x
+Base.tail(x::AbstractZero) = x
+Base.last(x::AbstractZero) = x
+
 Base.Broadcast.broadcastable(x::AbstractZero) = Ref(x)
 Base.Broadcast.broadcasted(::Type{T}) where {T<:AbstractZero} = T()
 

src/tangent_types/tangent.jl

@@ -96,6 +96,19 @@ function Base.show(io::IO, tangent::Tangent{P}) where {P}
     end
 end
 
+Base.first(tangent::Tangent{P,T}) where {P,T<:Union{Tuple,NamedTuple}} = first(backing(canonicalize(tangent)))
+Base.last(tangent::Tangent{P,T}) where {P,T<:Union{Tuple,NamedTuple}} = last(backing(canonicalize(tangent)))
+
+Base.tail(t::Tangent{P}) where {P<:Tuple} = Tangent{_tailtype(P)}(Base.tail(backing(canonicalize(t)))...)
+@generated _tailtype(::Type{P}) where {P<:Tuple} = Tuple{P.parameters[2:end]...}
+Base.tail(t::Tangent{<:Tuple{Any}}) = NoTangent()
+Base.tail(t::Tangent{<:Tuple{}}) = NoTangent()
+
+Base.tail(t::Tangent{P}) where {P<:NamedTuple} = Tangent{_tailtype(P)}(; Base.tail(backing(canonicalize(t)))...)
+_tailtype(::Type{NamedTuple{S,P}}) where {S,P} = NamedTuple{Base.tail(S), _tailtype(P)}
+Base.tail(t::Tangent{<:NamedTuple{<:Any, <:Tuple{Any}}}) = NoTangent()
+Base.tail(t::Tangent{<:NamedTuple{<:Any, <:Tuple{}}}) = NoTangent()
+
 function Base.getindex(tangent::Tangent{P,T}, idx::Int) where {P,T<:Union{Tuple,NamedTuple}}
     back = backing(canonicalize(tangent))
     return unthunk(getfield(back, idx))
@@ -127,6 +140,7 @@ end
 
 Base.iterate(tangent::Tangent, args...) = iterate(backing(tangent), args...)
 Base.length(tangent::Tangent) = length(backing(tangent))
+
 Base.eltype(::Type{<:Tangent{<:Any,T}}) where {T} = eltype(T)
 function Base.reverse(tangent::Tangent)
     rev_backing = reverse(backing(tangent))

src/tangent_types/thunks.jl

@@ -24,6 +24,10 @@ end
     return element, (underlying_object, new_state)
 end
 
+Base.first(x::AbstractThunk) = first(unthunk(x))
+Base.last(x::AbstractThunk) = last(unthunk(x))
+Base.tail(x::AbstractThunk) = Base.tail(unthunk(x))
+
 Base.:(==)(a::AbstractThunk, b::AbstractThunk) = unthunk(a) == unthunk(b)
 
 Base.:(-)(a::AbstractThunk) = -unthunk(a)

test/tangent_types/abstract_zero.jl

@@ -86,6 +86,10 @@
         @test z[1:3] === z
         @test z[1, 2] === z
         @test getindex(z) === z
+        
+        @test first(z) === z
+        @test last(z) === z
+        @test Base.tail(z) === z
     end
 
     @testset "NoTangent" begin

test/tangent_types/tangent.jl

@@ -72,12 +72,22 @@ end
         end
         @test Tangent{Foo}(; x=2.5).x == 2.5
 
-        @test keys(Tangent{Tuple{Float64}}(2.0)) == Base.OneTo(1)
+        tang1 = Tangent{Tuple{Float64}}(2.0)
+        @test keys(tang1) == Base.OneTo(1)
         @test propertynames(Tangent{Tuple{Float64}}(2.0)) == (1,)
         @test getindex(Tangent{Tuple{Float64}}(2.0), 1) == 2.0
         @test getindex(Tangent{Tuple{Float64}}(@thunk 2.0^2), 1) == 4.0
         @test getproperty(Tangent{Tuple{Float64}}(2.0), 1) == 2.0
         @test getproperty(Tangent{Tuple{Float64}}(@thunk 2.0^2), 1) == 4.0
+        @test NoTangent() === @inferred Base.tail(tang1)
+        @test NoTangent() === @inferred Base.tail(Tangent{Tuple{}}())
+        
+        tang3 = Tangent{Tuple{Float64, String, Vector{Float64}}}(1.0, NoTangent(), @thunk [3.0] .+ 4)
+        @test @inferred(first(tang3)) === tang3[1] === 1.0
+        @test @inferred(last(tang3)) isa Thunk
+        @test unthunk(last(tang3)) == [7.0]
+        @test Tuple(@inferred Base.tail(tang3))[1] === NoTangent()
+        @test Tuple(Base.tail(tang3))[end] isa Thunk
 
         NT = NamedTuple{(:a, :b),Tuple{Float64,Float64}}
         @test getindex(Tangent{NT}(; a=(@thunk 2.0^2)), :a) == 4.0
@@ -89,6 +99,14 @@ end
         @test getproperty(Tangent{NT}(; a=(@thunk 2.0^2)), :b) == ZeroTangent()
         @test getproperty(Tangent{NT}(; b=(@thunk 2.0^2)), 1) == ZeroTangent()
         @test getproperty(Tangent{NT}(; b=(@thunk 2.0^2)), 2) == 4.0
+        
+        @test first(Tangent{NT}(; a=(@thunk 2.0^2))) isa Thunk
+        @test unthunk(first(Tangent{NT}(; a=(@thunk 2.0^2)))) == 4.0
+        @test last(Tangent{NT}(; a=(@thunk 2.0^2))) isa ZeroTangent
+        
+        ntang1 = @inferred Base.tail(Tangent{NT}(; b=(@thunk 2.0^2)))
+        @test ntang1 isa Tangent{<:NamedTuple{(:b,)}}
+        @test NoTangent() === @inferred Base.tail(ntang1)
 
         # TODO: uncomment this once https://github.com/JuliaLang/julia/issues/35516
         if VERSION >= v"1.8-"

test/tangent_types/thunks.jl

@@ -16,6 +16,13 @@
 
         @test nothing === iterate(@thunk ()) == iterate(())
     end
+    
+    @testset "first, last, tail" begin
+        @test first(@thunk (1,2,3) .+ 4) === 5
+        @test last(@thunk (1,2,3) .+ 4) === 7
+        @test Base.tail(@thunk (1,2,3) .+ 4) === (6, 7)
+        @test Base.tail(@thunk NoTangent() * 5) === NoTangent()
+    end
 
     @testset "show" begin
         rep = repr(Thunk(rand))