diff --git a/src/copyable_task.jl b/src/copyable_task.jl index 3dba030..ac54901 100644 --- a/src/copyable_task.jl +++ b/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. diff --git a/test/copyable_task.jl b/test/copyable_task.jl index def48b5..0557e27 100644 --- a/test/copyable_task.jl +++ b/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