fix iteration of thunks (#370)

Author: mzgubic

Project.toml

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

src/differentials/thunks.jl

@@ -8,9 +8,11 @@ Base.Broadcast.broadcastable(x::AbstractThunk) = broadcastable(unthunk(x))
     return element, (val, state)
 end
 
-@inline function Base.iterate(::AbstractThunk, (val, state))
-    element, new_state = iterate(val, state)
-    return element, (val, new_state)
+@inline function Base.iterate(::AbstractThunk, (underlying_object, state))
+    next = iterate(underlying_object, state)
+    next === nothing && return nothing
+    element, new_state = next
+    return element, (underlying_object, new_state)
 end
 
 Base.:(==)(a::AbstractThunk, b::AbstractThunk) = unthunk(a) == unthunk(b)

test/differentials/thunks.jl

@@ -7,6 +7,14 @@
         @test 3.2 == InplaceableThunk(@thunk(3.2), x -> x + 3.2)
     end
 
+    @testset "iterate" begin
+        a = [1.0, 2.0, 3.0]
+        t = @thunk(a)
+        for (i, j) in zip(a, t)
+            @test i == j
+        end
+    end
+
     @testset "show" begin
         rep = repr(Thunk(rand))
         @test occursin(r"Thunk\(.*rand.*\)", rep)