Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jul 9, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion src/copyable_task.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down
20 changes: 20 additions & 0 deletions test/copyable_task.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading