Merge branch 'main' into mhauru/produce-wrap-loop

Author: yebai

src/copyable_task.jl

@@ -902,6 +902,15 @@ function derive_copyable_task_ir(ir::BBCode)::Tuple{BBCode,Tuple,Vector{Any}}
                     prod_val = deref_id
                 end
 
+                # Set the ref for this statement, as we would for any other call or invoke.
+                # The TapedTask may need to read this ref when it resumes, if the return
+                # value of `produce` is used within the original function.
+                if is_used_dict[id]
+                    out_ind = ssa_id_to_ref_index_map[id]
+                    set_ref = Expr(:call, set_ref_at!, refs_id, out_ind, prod_val)
+                    push!(inst_pairs, (ID(), new_inst(set_ref)))
+                end
+
                 # Construct a `ProducedValue`.
                 val_id = ID()
                 push!(inst_pairs, (val_id, new_inst(Expr(:call, ProducedValue, prod_val))))

test/copyable_task.jl

@@ -241,6 +241,13 @@
         @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
+        # https://github.com/TuringLang/Libtask.jl/issues/190.
+        f(obs) = produce(obs)
+        tt = Libtask.TapedTask(nothing, f, :a)
+        @test Libtask.consume(tt) === :a
         @test Libtask.consume(tt) === nothing
     end
 end