Fix throw_undef_if_not error

HISTORY.md

@@ -1,3 +1,11 @@
+# 0.9.13
+
+Fix a bug where SSA registers in `throw_undef_if_not` expressions were not being correctly handled.
+This affected `TapedTask`s where the underlying function had variables that were conditionally defined.
+
+Also adds an internal function, `Libtask.generate_ir`, which is useful for debugging.
+This is not part of the public API; however if you are developing Libtask it is a useful function for quickly inspecting the IR that Libtask generates, both pre- and post-optimisation.
+
 # 0.9.12
 
 Added a new method, `Libtask.get_taped_globals(tt::TapedTask)`.

Project.toml

@@ -3,7 +3,7 @@ uuid = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f"
 license = "MIT"
 desc = "Tape based task copying in Turing"
 repo = "https://github.com/TuringLang/Libtask.jl.git"
-version = "0.9.12"
+version = "0.9.13"
 
 [deps]
 MistyClosures = "dbe65cb8-6be2-42dd-bbc5-4196aaced4f4"

docs/src/internals.md

@@ -1,6 +1,7 @@
 # Internals
 
 ```@docs; canonical=true
+Libtask.generate_ir
 Libtask.produce_value
 Libtask.is_produce_stmt
 Libtask.stmt_might_produce

src/copyable_task.jl

@@ -99,6 +99,41 @@ function _throw_ir_error(@nospecialize(sig::Type{<:Tuple}))
     throw(ArgumentError(msg))
 end
 
+"""
+    generate_ir(optimise::Bool, f, args...; kwargs...)
+
+Returns `(original_ir, transformed_ir)` for the call `f(args...; kwargs...)`.
+
+The first element is the original `IRCode` that Julia generates for the call. The second is
+the transformed `IRCode` that Libtask would use to implement the `produce`/`consume`
+interface in a `TapedTask`.
+
+`optimise` controls whether the transformed IR (i.e. for the `TapedTask`) is optimised or
+not. Apart from inspecting the effects of optimisation, setting `optimise` to `false` can
+also be useful for debugging, because the optimisation pass will also perform verification,
+and will error if the IR is malformed (which can happen if Libtask's transformation pass has
+bugs).
+
+This is intended purely as a debugging tool, and is not exported.
+"""
+function generate_ir(optimise::Bool, fargs...; kwargs...)
+    all_args = isempty(kwargs) ? fargs : (Core.kwcall, getfield(kwargs, :data), fargs...)
+    sig = typeof(all_args)
+    ir_results = Base.code_ircode_by_type(sig)
+    if isempty(ir_results)
+        _throw_ir_error(sig)
+    end
+    original_ir = ir_results[1][1]
+    seed_id!()
+    bb, _, _ = derive_copyable_task_ir(BBCode(original_ir))
+    transformed_ir = if optimise
+        optimise_ir!(IRCode(bb))
+    else
+        IRCode(bb)
+    end
+    return original_ir, transformed_ir
+end
+
 """
     build_callable(sig::Type{<:Tuple})
 
@@ -929,7 +964,8 @@ function derive_copyable_task_ir(ir::BBCode)::Tuple{BBCode,Tuple,Vector{Any}}
                 if Meta.isexpr(stmt, :invoke) ||
                     Meta.isexpr(stmt, :call) ||
                     Meta.isexpr(stmt, :new) ||
-                    Meta.isexpr(stmt, :foreigncall)
+                    Meta.isexpr(stmt, :foreigncall) ||
+                    Meta.isexpr(stmt, :throw_undef_if_not)
 
                     # Find any `ID`s and replace them with calls to read whatever is stored
                     # in the `Ref`s that they are associated to.
@@ -962,8 +998,6 @@ function derive_copyable_task_ir(ir::BBCode)::Tuple{BBCode,Tuple,Vector{Any}}
                     push!(inst_pairs, (id, inst))
                 elseif Meta.isexpr(stmt, :gc_preserve_end)
                     push!(inst_pairs, (id, inst))
-                elseif Meta.isexpr(stmt, :throw_undef_if_not)
-                    push!(inst_pairs, (id, inst))
                 elseif stmt isa Nothing
                     push!(inst_pairs, (id, inst))
                 elseif stmt isa GlobalRef

test/copyable_task.jl

@@ -373,6 +373,36 @@ using Test
         # This used to error
         @test (consume(t); true)
     end
+
+    # Regression test for https://github.com/TuringLang/Libtask.jl/issues/204
+    @testset "throw_undef_if_not handling" begin
+        function tuin_g(y)
+            return produce(y + 1)
+        end
+        Libtask.might_produce(::Type{Tuple{typeof(tuin_g),Int}}) = true
+
+        function tuin_f(x)
+            if x == 1
+                y = 2
+            end
+            # `g` must produce.
+            # Also, there must be at least two calls to `produce` in the function.
+            tuin_g(y)
+            tuin_g(y)
+            return 1
+        end
+
+        # TapedTask construction used to error.
+        t = TapedTask(nothing, tuin_f, 1)
+        # Check that it returns the right things.
+        @test consume(t) == 3
+        @test consume(t) == 3
+        @test consume(t) === nothing
+
+        # This is bogus, but we should check that it has the right behaviour.
+        t = TapedTask(nothing, tuin_f, 2)
+        @test_throws UndefVarError consume(t)
+    end
 end
 
 end # module