Fix a bug with indirect produce calls in a loop (#192)

* Bump patch version to 0.9.3

* Fix a bug with indirect produce calls in a loop

* Fix a typo in a docstring

---------

Co-authored-by: Hong Ge <[email protected]>

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 resulting `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

@@ -221,6 +221,26 @@
         @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
     @testset "Return produce" begin
         # Test calling a function that does something with the return value of `produce`.
         # In this case it just returns it. This used to error, see