Update args at copy (#141)

* Update mappingx

* Revert "Update mappingx"

This reverts commit 0355eab.

* store and copy args of tapedtask

* check type of args

* Update src/tapedtask.jl

* Update src/tapedtask.jl

Co-authored-by: KDr2 <zhuo.dev@gmail.com>
Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>

Author: 3 people

src/tapedtask.jl

@@ -3,20 +3,22 @@ struct TapedTaskException
     backtrace::Vector{Any}
 end
 
-struct TapedTask{F}
+struct TapedTask{F, AT<:Tuple}
     task::Task
     tf::TapedFunction{F}
+    args::AT
     produce_ch::Channel{Any}
     consume_ch::Channel{Int}
     produced_val::Vector{Any}
 
     function TapedTask(
         t::Task,
         tf::TapedFunction{F},
+        args::AT,
         produce_ch::Channel{Any},
         consume_ch::Channel{Int}
-    ) where {F}
-        new{F}(t, tf, produce_ch, consume_ch, Any[])
+    ) where {F, AT<:Tuple}
+        new{F, AT}(t, tf, args, produce_ch, consume_ch, Any[])
     end
 end
 
@@ -55,7 +57,7 @@ function TapedTask(tf::TapedFunction, args...)
     produce_ch = Channel()
     consume_ch = Channel{Int}()
     task = @task wrap_task(tf, produce_ch, consume_ch, args...)
-    t = TapedTask(task, tf, produce_ch, consume_ch)
+    t = TapedTask(task, tf, args, produce_ch, consume_ch)
     task.storage === nothing && (task.storage = IdDict())
     task.storage[:tapedtask] = t
     return t
@@ -159,9 +161,15 @@ Base.IteratorEltype(::Type{<:TapedTask}) = Base.EltypeUnknown()
 
 # copy the task
 
-function Base.copy(t::TapedTask)
+function Base.copy(t::TapedTask; args=())
     tf = copy(t.tf)
-    new_t = TapedTask(tf)
+    task_args = if length(args) > 0
+        typeof(args) == typeof(t.args) || error("bad arguments")
+        args
+    else
+        tape_copy.(t.args)
+    end
+    new_t = TapedTask(tf, task_args...)
     storage = t.task.storage::IdDict{Any,Any}
     new_t.task.storage = copy(storage)
     new_t.task.storage[:tapedtask] = new_t

test/tapedtask.jl

@@ -184,4 +184,20 @@
             end
         end
     end
+
+    @testset "Issues" begin
+        @testset "Issue-140, copy unstarted task" begin
+            function f(x)
+                for i in 1:3
+                    produce(i + x)
+                end
+            end
+
+            ttask = TapedTask(f, 3)
+            ttask2 = copy(ttask)
+            @test consume(ttask2) == 4
+            ttask3 = copy(ttask; args=(4,))
+            @test consume(ttask3) == 5
+        end
+    end
 end