Merge branch 'master' into perf

Author: yebai

.github/workflows/Testing.yaml

@@ -42,5 +42,3 @@ jobs:
             ${{ runner.os }}-
       - uses: julia-actions/julia-buildpkg@latest
       - uses: julia-actions/julia-runtest@latest
-        with:
-          coverage: false

Project.toml

@@ -3,17 +3,16 @@ 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.7.1"
+version = "0.7.3"
 
 [deps]
-CodeInfoTools = "bc773b8a-8374-437a-b9f2-0e9785855863"
 FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-CodeInfoTools = "0.3.4"
+FunctionWrappers = "1.1"
 LRUCache = "1.3"
 julia = "1.7"
 

src/Libtask.jl

@@ -1,6 +1,5 @@
 module Libtask
 
-using CodeInfoTools
 using FunctionWrappers: FunctionWrapper
 using LRUCache
 

src/tapedfunction.jl

@@ -3,6 +3,14 @@
 abstract type AbstractInstruction end
 const RawTape = Vector{AbstractInstruction}
 
+function _infer(f, args_type)
+    # `code_typed` returns a vector: [Pair{Core.CodeInfo, DataType}]
+    ir0 = code_typed(f, Tuple{args_type...}, optimize=false)[1][1]
+    # ir1 = CodeInfoTools.code_inferred(f, args_type...)
+    # ir1.ssavaluetypes = ir0.ssavaluetypes
+    return ir0
+end
+
 mutable struct TapedFunction{F, TapeType}
     func::F # maybe a function, a constructor, or a callable object
     arity::Int
@@ -22,8 +30,7 @@ mutable struct TapedFunction{F, TapeType}
             tf.counter = 1
             return tf
         end
-
-        ir = CodeInfoTools.code_inferred(f, args_type...)
+        ir = _infer(f, args_type)
         bindings, tape = translate!(RawTape(), ir)
 
         tf = new{F, T}(f, length(args), ir, tape, 1, bindings, :none)

src/tapedtask.jl

@@ -162,12 +162,24 @@ Base.IteratorEltype(::Type{<:TapedTask}) = Base.EltypeUnknown()
 # copy the task
 
 function Base.copy(t::TapedTask; args=())
+    length(args) > 0 && t.tf.counter >1 &&
+        error("can't copy started task with new arguments")
     tf = copy(t.tf)
     task_args = if length(args) > 0
+        # this cond implies t.tf.counter == 0, i.e., the task is not started yet
         typeof(args) == typeof(t.args) || error("bad arguments")
         args
     else
-        tape_copy.(t.args)
+        if t.tf.counter > 1
+            # the task is running, we find the
+            # real args from the copied bindings
+            map(1:length(t.args)) do i
+                get(tf.bindings, Symbol("_", i + 1), t.args[i])
+            end
+        else
+            # the task is not started yet, but no args is given
+            tape_copy.(t.args)
+        end
     end
     new_t = TapedTask(tf, task_args...)
     storage = t.task.storage::IdDict{Any,Any}