Fix a bug with using the return value of produce

Author: mhauru

src/copyable_task.jl

@@ -898,6 +898,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

@@ -220,4 +220,16 @@
         g() = produce(rand() > -1.0 ? 2 : 0.1)
         @test Libtask.consume(Libtask.TapedTask(nothing, g)) == 2
     end
+
+    @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.
+        produce_wrapper(x) = Libtask.produce(x)
+        Libtask.might_produce(::Type{<:Tuple{typeof(produce_wrapper),Any}}) = true
+        f(obs) = produce_wrapper(obs)
+        tt = Libtask.TapedTask(nothing, f, :a)
+        @test Libtask.consume(tt) === :a
+        @test Libtask.consume(tt) === nothing
+    end
 end