Fix a bug with indirect produce calls in a loop

Author: mhauru

src/copyable_task.jl

@@ -287,7 +287,7 @@ end
     fresh_copy(mc::T) where {T<:MistyClosure}
 
 Creates an independent copy of `mc` by (carefully) replacing the `Ref`s it
-contains it its `captures`. The resuting `MistyClosure` is safe to run.
+contains in its `captures`. The resuting `MistyClosure` is safe to run.
 
 This is achieved by replacing most `Ref`s with new `Ref`s of the same (el)type,
 but with nothing stored in them -- values will be stored in them when the
@@ -830,6 +830,10 @@ function derive_copyable_task_ir(ir::BBCode)::Tuple{BBCode,Tuple,Vector{Any}}
                 # terminator. We handle this in a similar way to the statements above.
 
                 if stmt isa ReturnNode
+                    # Reset the position counter to `-1`, so that if this function gets
+                    # called again, execution starts from the beginning.
+                    expr = Expr(:call, set_resume_block!, refs_id, Int32(-1))
+                    push!(inst_pairs, (ID(), new_inst(expr)))
                     # If returning an SSA, it might be one whose value was restored from
                     # before. Therefore, grab it out of storage, rather than assuming that
                     # it is def-ed.

test/copyable_task.jl

@@ -220,4 +220,27 @@
         g() = produce(rand() > -1.0 ? 2 : 0.1)
         @test Libtask.consume(Libtask.TapedTask(nothing, g)) == 2
     end
+
+    @testset "Indirect produce in a loop" begin
+        # Test that we can wrap a `produce` call in another function, and call that function
+        # in a loop. This used to only produce some of the values, see
+        # https://github.com/TuringLang/Libtask.jl/issues/190.
+        produce_wrapper(x) = (Libtask.produce(x); return nothing)
+        Libtask.might_produce(::Type{<:Tuple{typeof(produce_wrapper),Any}}) = true
+        function f(obs)
+            for o in obs
+                produce_wrapper(o)
+            end
+            return nothing
+        end
+
+        # That the eltype of vals is Any is significant for reproducing the original bug.
+        # Unclear why.
+        vals = Any[:a, :b, :c]
+        tt = Libtask.TapedTask(nothing, f, vals)
+        @test Libtask.consume(tt) === :a
+        @test Libtask.consume(tt) === :b
+        @test Libtask.consume(tt) === :c
+        @test Libtask.consume(tt) === nothing
+    end
 end