From c8af0f4c28cef51b3e02296e6a1167a850de16a5 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Tue, 5 Jan 2021 15:57:58 +0000
Subject: [PATCH 01/14] improve TArray benchmarks

---
 test/benchmarks.jl | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/test/benchmarks.jl b/test/benchmarks.jl
index ea9fd025..d8a63793 100644
--- a/test/benchmarks.jl
+++ b/test/benchmarks.jl
@@ -1,21 +1,23 @@
 using BenchmarkTools
 using Libtask
 
+INTENSITY = 3
+
 println("= Benchmarks on Arrays =")
-A = rand(100, 100)
-x, y =  abs.(rand(Int, 2) .% 100)
+A = rand(1000, 1000)
+x, y =  abs.(rand(Int, 2) .% 999) .+ 1
 print("indexing: ")
-@btime $A[$x, $y] + $A[$x, $y]
+@btime for _ in 1:INTENSITY; $A[$x, $y]; end
 print("set indexing: ")
-@btime $A[$x, $y] = 1
+@btime for _ in 1:INTENSITY; $A[$x, $y] = 1; end
 print("broadcast: ")
-@btime $A .+ $A
+@btime for _ in 1:INTENSITY; $A .+ $A; end
 
 println("= Benchmarks on TArrays =")
 TA = Libtask.localize(deepcopy(A))
 print("indexing: ")
-@btime $TA[$x, $y] + $TA[$x, $y]
+@btime for _ in 1:INTENSITY; $TA[$x, $y]; end
 print("set indexing: ")
-@btime $TA[$x, $y] = 1
+@btime for _ in 1:INTENSITY; $TA[$x, $y] = 1; end
 print("broadcast: ")
-@btime $TA .+ $TA
+@btime for _ in 1:INTENSITY; $TA .+ $TA; end

From fd4abd6148bee0d9a30311e173f9a8a6da609872 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Thu, 7 Jan 2021 07:04:50 +0000
Subject: [PATCH 02/14] update benchmarks

---
 test/benchmarks.jl                  | 23 ++++++++++++++++-------
 {deps => utils}/methods_of_array.jl |  0
 2 files changed, 16 insertions(+), 7 deletions(-)
 rename {deps => utils}/methods_of_array.jl (100%)

diff --git a/test/benchmarks.jl b/test/benchmarks.jl
index d8a63793..efefdfd1 100644
--- a/test/benchmarks.jl
+++ b/test/benchmarks.jl
@@ -1,23 +1,32 @@
 using BenchmarkTools
 using Libtask
 
-INTENSITY = 3
+
+ macro rep(cnt, exp)
+     blk =:(begin end)
+     for _ in 1:eval(cnt)
+         push!(blk.args, esc(exp))
+     end
+     blk
+ end
+
+INTENSITY = 2
 
 println("= Benchmarks on Arrays =")
 A = rand(1000, 1000)
 x, y =  abs.(rand(Int, 2) .% 999) .+ 1
 print("indexing: ")
-@btime for _ in 1:INTENSITY; $A[$x, $y]; end
+@btime @rep INTENSITY $A[$x, $y]
 print("set indexing: ")
-@btime for _ in 1:INTENSITY; $A[$x, $y] = 1; end
+@btime @rep INTENSITY $A[$x, $y] = 1
 print("broadcast: ")
-@btime for _ in 1:INTENSITY; $A .+ $A; end
+@btime @rep INTENSITY $A .+ $A
 
 println("= Benchmarks on TArrays =")
 TA = Libtask.localize(deepcopy(A))
 print("indexing: ")
-@btime for _ in 1:INTENSITY; $TA[$x, $y]; end
+@btime @rep INTENSITY $TA[$x, $y]
 print("set indexing: ")
-@btime for _ in 1:INTENSITY; $TA[$x, $y] = 1; end
+@btime @rep INTENSITY $TA[$x, $y] = 1
 print("broadcast: ")
-@btime for _ in 1:INTENSITY; $TA .+ $TA; end
+@btime @rep INTENSITY $TA .+ $TA
diff --git a/deps/methods_of_array.jl b/utils/methods_of_array.jl
similarity index 100%
rename from deps/methods_of_array.jl
rename to utils/methods_of_array.jl

From ba7d533156f27732b20001b3602c00ee32db7552 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Thu, 7 Jan 2021 10:43:11 +0000
Subject: [PATCH 03/14] 1. Base.@propagate_inbounds, 2. use function in
 benchmarks

---
 src/tarray.jl      |  7 +++----
 test/benchmarks.jl | 18 +++++++++++-------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/src/tarray.jl b/src/tarray.jl
index 67df1f23..6363e00d 100644
--- a/src/tarray.jl
+++ b/src/tarray.jl
@@ -174,12 +174,11 @@ end
 #
 
 # Indexing Interface
-function Base.getindex(x::TArray{T, N}, I::Vararg{Int,N}) where {T, N}
-    t, d = task_local_storage(x.ref)
-    return d[I...]
+Base.@propagate_inbounds function Base.getindex(x::TArray{T, N}, I::Vararg{Int,N}) where {T, N}
+    return task_local_storage(x.ref)[2][I...]
 end
 
-function Base.setindex!(x::TArray{T, N}, e, I::Vararg{Int,N}) where {T, N}
+Base.@propagate_inbounds function Base.setindex!(x::TArray{T, N}, e, I::Vararg{Int,N}) where {T, N}
     n, d = task_local_storage(x.ref)
     cn   = n_copies()
     newd = d
diff --git a/test/benchmarks.jl b/test/benchmarks.jl
index efefdfd1..e3bb3a85 100644
--- a/test/benchmarks.jl
+++ b/test/benchmarks.jl
@@ -10,23 +10,27 @@ using Libtask
      blk
  end
 
-INTENSITY = 2
+INTENSITY = 6
+
+indexing(a, x, y) = @rep INTENSITY a[x, y]
+setindexing(a, x, y) = @rep INTENSITY a[x, y] = 1
+broadcasting(a) = @rep INTENSITY a .+ a
 
 println("= Benchmarks on Arrays =")
 A = rand(1000, 1000)
 x, y =  abs.(rand(Int, 2) .% 999) .+ 1
 print("indexing: ")
-@btime @rep INTENSITY $A[$x, $y]
+@btime indexing($A, $x, $y)
 print("set indexing: ")
-@btime @rep INTENSITY $A[$x, $y] = 1
+@btime setindexing($A, $x, $y)
 print("broadcast: ")
-@btime @rep INTENSITY $A .+ $A
+@btime broadcasting($A)
 
 println("= Benchmarks on TArrays =")
 TA = Libtask.localize(deepcopy(A))
 print("indexing: ")
-@btime @rep INTENSITY $TA[$x, $y]
+@btime indexing($TA, $x, $y)
 print("set indexing: ")
-@btime @rep INTENSITY $TA[$x, $y] = 1
+@btime setindexing($TA, $x, $y)
 print("broadcast: ")
-@btime @rep INTENSITY $TA .+ $TA
+@btime broadcasting($TA)

From cf3bbaf1e891a8e10c74c9271266bddda37ff38a Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Tue, 12 Jan 2021 07:20:22 +0000
Subject: [PATCH 04/14] optimize indexing for Matrix{Float64}

---
 src/ctask.jl       | 18 +++++++++++++----
 src/tarray.jl      | 50 +++++++++++++++++++++++++++++++---------------
 test/benchmarks.jl | 20 ++++++++++++-------
 3 files changed, 61 insertions(+), 27 deletions(-)

diff --git a/src/ctask.jl b/src/ctask.jl
index 06958d2b..b3e8bcec 100644
--- a/src/ctask.jl
+++ b/src/ctask.jl
@@ -5,12 +5,20 @@ Wrapper of a [`Task`](@ref) for which deep copying of stack allocated objects is
 """
 struct CTask
     task::Task
+    data_f2::IdDict{Symbol, Tuple{Int, Array{Float64, 2}}}
 
-    function CTask(task::Task)
-        new(enable_stack_copying(task))
+    function CTask(task::Task, df2::IdDict{Symbol, Tuple{Int, Array{Float64, 2}}})
+        ret = new(enable_stack_copying(task), df2)
+        task.storage === nothing && (task.storage = IdDict())
+        task.storage[:ctask] = ret
+        ret
     end
 end
 
+function CTask(task::Task)
+    df2 = IdDict{Symbol, Tuple{Int, Array{Float64, 2}}}()
+    CTask(enable_stack_copying(task), df2)
+end
 CTask(f) = CTask(Task(task_wrapper(f)))
 
 # Iteration interface.
@@ -24,12 +32,13 @@ end
 
 function Base.showerror(io::IO, ex::CTaskException)
     println(io, "CTaskException:")
+    ct = ex.task
     bt = @static if VERSION < v"1.6.0-DEV.1145"
         ct.backtrace
     else
         ct.storage[:_libtask_bt]
     end
-    showerror(io, ex.task.exception, bt)
+    showerror(io, ct.exception, bt)
 end
 
 # Utility function for self-copying mechanism
@@ -107,7 +116,8 @@ function Base.copy(ctask::CTask)
     setstate!(newtask, getstate(task))
     newtask.result = task.result
 
-    return CTask(newtask)
+    df2 = copy(ctask.data_f2)
+    return CTask(newtask, df2)
 end
 
 function produce(v)
diff --git a/src/tarray.jl b/src/tarray.jl
index 6363e00d..72cb5dea 100644
--- a/src/tarray.jl
+++ b/src/tarray.jl
@@ -40,7 +40,7 @@ function TArray(T::Type, dim)
     res = TArray{T,length(dim)}();
     n = n_copies()
     d = Array{T}(undef, dim)
-    task_local_storage(res.ref, (n,d))
+    _local_storage(res.ref, (n,d))
     res
 end
 
@@ -49,6 +49,23 @@ TArray(x::AbstractArray) = convert(TArray, x)
 localize(x) = x
 localize(x::AbstractArray) = TArray(x)
 
+# _local_storage
+_local_storage(k) = task_local_storage(k)
+_local_storage(k, v) = task_local_storage(k, v)
+function _local_storage(k::Symbol, v::Tuple{Int, Array{Float64, 2}})
+    ctask::CTask = task_local_storage(:ctask)
+    ctask.data_f2[k] = v
+end
+function _local_storage(k::Symbol, ::Type{Array{T, N}}) where {T, N}
+    n::Int, d::Array{T, N} = task_local_storage(k)
+    return n, d
+end
+
+function _local_storage(k::Symbol, ::Type{Array{Float64, 2}})
+    ctask::CTask = task_local_storage(:ctask)
+    return ctask.data_f2[k]
+end
+
 # Constructors
 """
      tzeros(dims, ...)
@@ -71,7 +88,7 @@ function tzeros(T::Type, dim)
     res = TArray{T,length(dim)}();
     n = n_copies()
     d = zeros(T,dim)
-    task_local_storage(res.ref, (n,d))
+    _local_storage(res.ref, (n,d))
     return res
 end
 
@@ -99,7 +116,7 @@ function tfill(val::Real, dim)
     res = TArray{typeof(val),length(dim)}();
     n = n_copies()
     d = fill(val,dim)
-    task_local_storage(res.ref, (n,d))
+    _local_storage(res.ref, (n,d))
     return res
 end
 
@@ -107,8 +124,8 @@ end
 # Conversion between TArray and Array
 #
 _get(x) = x
-function _get(x::TArray)
-    n, d = task_local_storage(x.ref)
+function _get(x::TArray{T, N}) where {T, N}
+    n, d = _local_storage(x.ref, Array{T, N})
     return d
 end
 
@@ -126,7 +143,7 @@ end
 function Base.convert(::Type{TArray{T,N}}, x::AbstractArray{T,N}) where {T,N}
     res = TArray{T,N}()
     n   = n_copies()
-    task_local_storage(res.ref, (n,x))
+    _local_storage(res.ref, (n,x))
     return res
 end
 
@@ -141,7 +158,8 @@ function Base.show(io::IO, ::MIME"text/plain", x::TArray)
     show(io,  MIME("text/plain"), arr)
 end
 
-Base.show(io::IO, x::TArray) = Base.show(io::IO, task_local_storage(x.ref)[2])
+Base.show(io::IO, x::TArray{T, N}) where {T, N} =
+    Base.show(io::IO, _local_storage(x.ref, Array{T, N})[2])
 
 function Base.summary(io::IO, x::TArray)
   print(io, "Task Local Array: ")
@@ -175,39 +193,39 @@ end
 
 # Indexing Interface
 Base.@propagate_inbounds function Base.getindex(x::TArray{T, N}, I::Vararg{Int,N}) where {T, N}
-    return task_local_storage(x.ref)[2][I...]
+    return _local_storage(x.ref, Array{T, N})[2][I...]
 end
 
 Base.@propagate_inbounds function Base.setindex!(x::TArray{T, N}, e, I::Vararg{Int,N}) where {T, N}
-    n, d = task_local_storage(x.ref)
+    n, d = _local_storage(x.ref, Array{T, N})
     cn   = n_copies()
     newd = d
     if cn > n
         # println("[setindex!]: $(x.ref) copying data")
         newd = deepcopy(d)
-        task_local_storage(x.ref, (cn, newd))
+        _local_storage(x.ref, (cn, newd))
     end
     newd[I...] = e
 end
 
-function Base.push!(x::TArray{T}, e) where T
-    n, d = task_local_storage(x.ref)
+function Base.push!(x::TArray{T, N}, e) where {T, N}
+    n, d = _local_storage(x.ref, Array{T, N})
     cn   = n_copies()
     newd = d
     if cn > n
         newd = deepcopy(d)
-        task_local_storage(x.ref, (cn, newd))
+        _local_storage(x.ref, (cn, newd))
     end
     push!(newd, e)
 end
 
-function Base.pop!(x::TArray)
-    n, d = task_local_storage(x.ref)
+function Base.pop!(x::TArray{T, N}) where {T, N}
+    n, d = _local_storage(x.ref, Array{T, N})
     cn   = n_copies()
     newd = d
     if cn > n
         newd = deepcopy(d)
-        task_local_storage(x.ref, (cn, newd))
+        _local_storage(x.ref, (cn, newd))
     end
     pop!(d)
 end
diff --git a/test/benchmarks.jl b/test/benchmarks.jl
index e3bb3a85..f221814b 100644
--- a/test/benchmarks.jl
+++ b/test/benchmarks.jl
@@ -27,10 +27,16 @@ print("broadcast: ")
 @btime broadcasting($A)
 
 println("= Benchmarks on TArrays =")
-TA = Libtask.localize(deepcopy(A))
-print("indexing: ")
-@btime indexing($TA, $x, $y)
-print("set indexing: ")
-@btime setindexing($TA, $x, $y)
-print("broadcast: ")
-@btime broadcasting($TA)
+function ctask_test()
+    TA = Libtask.localize(deepcopy(A))
+    print("indexing: ")
+    @btime indexing($TA, $x, $y)
+    print("set indexing: ")
+    @btime setindexing($TA, $x, $y)
+    print("broadcast: ")
+    @btime broadcasting($TA)
+end
+
+ct = CTask(ctask_test)
+schedule(ct.task)
+wait(ct.task)

From 3f9c6e77972e4bcd0a09705125fccdb3b9ef69b3 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Tue, 12 Jan 2021 10:48:54 +0000
Subject: [PATCH 05/14] indent code

---
 test/benchmarks.jl | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/test/benchmarks.jl b/test/benchmarks.jl
index f221814b..30460494 100644
--- a/test/benchmarks.jl
+++ b/test/benchmarks.jl
@@ -2,13 +2,13 @@ using BenchmarkTools
 using Libtask
 
 
- macro rep(cnt, exp)
-     blk =:(begin end)
-     for _ in 1:eval(cnt)
-         push!(blk.args, esc(exp))
-     end
-     blk
- end
+macro rep(cnt, exp)
+    blk =:(begin end)
+    for _ in 1:eval(cnt)
+        push!(blk.args, esc(exp))
+    end
+    blk
+end
 
 INTENSITY = 6
 

From f1f9e6db6683ddedf0f2a5920afe6aad70a35e7f Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Wed, 13 Jan 2021 09:48:37 +0000
Subject: [PATCH 06/14] move TArray data into TArray

---
 src/ctask.jl  |  13 ++----
 src/tarray.jl | 121 ++++++++++++++++++++++++++------------------------
 2 files changed, 67 insertions(+), 67 deletions(-)

diff --git a/src/ctask.jl b/src/ctask.jl
index b3e8bcec..8665c7b5 100644
--- a/src/ctask.jl
+++ b/src/ctask.jl
@@ -5,20 +5,15 @@ Wrapper of a [`Task`](@ref) for which deep copying of stack allocated objects is
 """
 struct CTask
     task::Task
-    data_f2::IdDict{Symbol, Tuple{Int, Array{Float64, 2}}}
 
-    function CTask(task::Task, df2::IdDict{Symbol, Tuple{Int, Array{Float64, 2}}})
-        ret = new(enable_stack_copying(task), df2)
+    function CTask(task::Task)
+        ret = new(enable_stack_copying(task))
         task.storage === nothing && (task.storage = IdDict())
         task.storage[:ctask] = ret
         ret
     end
 end
 
-function CTask(task::Task)
-    df2 = IdDict{Symbol, Tuple{Int, Array{Float64, 2}}}()
-    CTask(enable_stack_copying(task), df2)
-end
 CTask(f) = CTask(Task(task_wrapper(f)))
 
 # Iteration interface.
@@ -116,8 +111,8 @@ function Base.copy(ctask::CTask)
     setstate!(newtask, getstate(task))
     newtask.result = task.result
 
-    df2 = copy(ctask.data_f2)
-    return CTask(newtask, df2)
+    copy_tarrays(task, newtask)
+    return CTask(newtask)
 end
 
 function produce(v)
diff --git a/src/tarray.jl b/src/tarray.jl
index 72cb5dea..d228fb5f 100644
--- a/src/tarray.jl
+++ b/src/tarray.jl
@@ -23,10 +23,15 @@ for i in 1:4 ta[i] = i end  # assign
 Array(ta)                   # convert to 4-element Array{Int64,1}: [1, 2, 3, 4]
 ```
 """
-struct TArray{T,N} <: AbstractArray{T,N}
-    ref :: Symbol  # object_id
+struct TArray{T, N} <: AbstractArray{T, N}
     orig_task :: Task
-    TArray{T,N}() where {T,N} = new(gensym(), current_task())
+    data::Dict{Task, Tuple{Int, AbstractArray{T, N}}}
+    function TArray{T,N}() where {T,N}
+        d = Dict{Task, Tuple{Int, AbstractArray{T, N}}}()
+        res = new(current_task(), d)
+        keep(res)
+        return res
+    end
 end
 
 TArray{T}(d::Integer...) where T = TArray(T, d)
@@ -37,35 +42,55 @@ TArray{T,N}(::UndefInitializer, d::Vararg{<:Integer,N}) where {T,N} = TArray{T,N
 TArray{T,N}(dim::NTuple{N,Int}) where {T,N} = TArray(T, dim)
 
 function TArray(T::Type, dim)
-    res = TArray{T,length(dim)}();
+    res = TArray{T, length(dim)}();
     n = n_copies()
     d = Array{T}(undef, dim)
-    _local_storage(res.ref, (n,d))
+    _local_storage(res, (n, d))
     res
 end
 
 TArray(x::AbstractArray) = convert(TArray, x)
 
-localize(x) = x
-localize(x::AbstractArray) = TArray(x)
+# TArray House-Keeper
+
+const TArrayKeeper = Vector{WeakRef}()
+keep(x::TArray) = push!(TArrayKeeper, WeakRef(x))
+function copy_tarrays(task1::Task, task2::Task)
+    deleteat!(TArrayKeeper, map(x -> x == nothing, TArrayKeeper))
+    for wref in TArrayKeeper
+        ta = wref.value
+        if haskey(ta.data, task1) && !haskey(ta.data, task2)
+            ta.data[task2] = ta.data[task1]
+        end
+    end
+end
 
 # _local_storage
-_local_storage(k) = task_local_storage(k)
-_local_storage(k, v) = task_local_storage(k, v)
-function _local_storage(k::Symbol, v::Tuple{Int, Array{Float64, 2}})
-    ctask::CTask = task_local_storage(:ctask)
-    ctask.data_f2[k] = v
-end
-function _local_storage(k::Symbol, ::Type{Array{T, N}}) where {T, N}
-    n::Int, d::Array{T, N} = task_local_storage(k)
-    return n, d
+_local_storage(x::TArray{T, N}) where {T, N} = x.data[current_task()]
+function _local_storage(x::TArray{T, N}, v::Tuple{Int, AbstractArray{T, N}}) where {T, N}
+    x.data[current_task()] = v
 end
 
-function _local_storage(k::Symbol, ::Type{Array{Float64, 2}})
-    ctask::CTask = task_local_storage(:ctask)
-    return ctask.data_f2[k]
+_get(x) = x
+function _get(x::TArray{T, N}) where {T, N}
+    n, d = x.data[current_task()]
+    return d
+end
+function _get_for_write(x::TArray{T, N}) where {T, N}
+    n, d = x.data[current_task()]
+    cn   = n_copies()
+    newd = d
+    if cn > n
+        # println("[setindex!]: $(x.ref) copying data")
+        newd = deepcopy(d)
+        x.data[current_task()] = (cn, newd)
+    end
+    return newd
 end
 
+localize(x) = x
+localize(x::AbstractArray) = TArray(x)
+
 # Constructors
 """
      tzeros(dims, ...)
@@ -88,7 +113,7 @@ function tzeros(T::Type, dim)
     res = TArray{T,length(dim)}();
     n = n_copies()
     d = zeros(T,dim)
-    _local_storage(res.ref, (n,d))
+    _local_storage(res, (n, d))
     return res
 end
 
@@ -116,18 +141,13 @@ function tfill(val::Real, dim)
     res = TArray{typeof(val),length(dim)}();
     n = n_copies()
     d = fill(val,dim)
-    _local_storage(res.ref, (n,d))
+    _local_storage(res, (n, d))
     return res
 end
 
 #
 # Conversion between TArray and Array
 #
-_get(x) = x
-function _get(x::TArray{T, N}) where {T, N}
-    n, d = _local_storage(x.ref, Array{T, N})
-    return d
-end
 
 function Base.convert(::Type{Array}, x::TArray)
     return convert(Array{eltype(x), ndims(x)}, x)
@@ -138,12 +158,12 @@ function Base.convert(::Type{Array{T,N}}, x::TArray{T,N}) where {T,N}
 end
 
 function Base.convert(::Type{TArray}, x::AbstractArray)
-    return convert(TArray{eltype(x),ndims(x)}, x)
+    return convert(TArray{eltype(x), ndims(x)}, x)
 end
-function Base.convert(::Type{TArray{T,N}}, x::AbstractArray{T,N}) where {T,N}
-    res = TArray{T,N}()
+function Base.convert(::Type{TArray{T,N}}, x::AbstractArray{T, N}) where {T, N}
+    res = TArray{T, N}()
     n   = n_copies()
-    _local_storage(res.ref, (n,x))
+    _local_storage(res, (n, x))
     return res
 end
 
@@ -151,15 +171,19 @@ end
 # Representation
 #
 function Base.show(io::IO, ::MIME"text/plain", x::TArray)
-    arr = x.orig_task.storage[x.ref][2]
+    arr = x.data[x.orig_task][2]
     @warn "Here shows the originating task's storage, " *
         "not the current task's storage. " *
         "Please explicitly call show(::TArray) to display the current task's version of a TArray."
     show(io,  MIME("text/plain"), arr)
 end
 
-Base.show(io::IO, x::TArray{T, N}) where {T, N} =
-    Base.show(io::IO, _local_storage(x.ref, Array{T, N})[2])
+function Base.show(io::IO, x::TArray)
+    if haskey(x.data, current_task())
+        return Base.show(io::IO, _get(x))
+    end
+    show(io, MIME("text/plain"), x)
+end
 
 function Base.summary(io::IO, x::TArray)
   print(io, "Task Local Array: ")
@@ -193,40 +217,21 @@ end
 
 # Indexing Interface
 Base.@propagate_inbounds function Base.getindex(x::TArray{T, N}, I::Vararg{Int,N}) where {T, N}
-    return _local_storage(x.ref, Array{T, N})[2][I...]
+    return _get(x)[I...]
 end
 
 Base.@propagate_inbounds function Base.setindex!(x::TArray{T, N}, e, I::Vararg{Int,N}) where {T, N}
-    n, d = _local_storage(x.ref, Array{T, N})
-    cn   = n_copies()
-    newd = d
-    if cn > n
-        # println("[setindex!]: $(x.ref) copying data")
-        newd = deepcopy(d)
-        _local_storage(x.ref, (cn, newd))
-    end
-    newd[I...] = e
+    d = _get_for_write(x)
+    d[I...] = e
 end
 
 function Base.push!(x::TArray{T, N}, e) where {T, N}
-    n, d = _local_storage(x.ref, Array{T, N})
-    cn   = n_copies()
-    newd = d
-    if cn > n
-        newd = deepcopy(d)
-        _local_storage(x.ref, (cn, newd))
-    end
-    push!(newd, e)
+    d = _get_for_write(x)
+    push!(d, e)
 end
 
 function Base.pop!(x::TArray{T, N}) where {T, N}
-    n, d = _local_storage(x.ref, Array{T, N})
-    cn   = n_copies()
-    newd = d
-    if cn > n
-        newd = deepcopy(d)
-        _local_storage(x.ref, (cn, newd))
-    end
+    d = _get_for_write(x)
     pop!(d)
 end
 

From 694c23e53a63bcb5efb6b6985377c38e06c41c01 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Wed, 13 Jan 2021 10:15:56 +0000
Subject: [PATCH 07/14] update benchmarks

---
 test/benchmarks.jl | 20 +++++++-------------
 1 file changed, 7 insertions(+), 13 deletions(-)

diff --git a/test/benchmarks.jl b/test/benchmarks.jl
index 30460494..25327808 100644
--- a/test/benchmarks.jl
+++ b/test/benchmarks.jl
@@ -27,16 +27,10 @@ print("broadcast: ")
 @btime broadcasting($A)
 
 println("= Benchmarks on TArrays =")
-function ctask_test()
-    TA = Libtask.localize(deepcopy(A))
-    print("indexing: ")
-    @btime indexing($TA, $x, $y)
-    print("set indexing: ")
-    @btime setindexing($TA, $x, $y)
-    print("broadcast: ")
-    @btime broadcasting($TA)
-end
-
-ct = CTask(ctask_test)
-schedule(ct.task)
-wait(ct.task)
+TA = Libtask.localize(deepcopy(A))
+print("indexing: ")
+@btime indexing($TA, $x, $y)
+print("set indexing: ")
+@btime setindexing($TA, $x, $y)
+print("broadcast: ")
+@btime broadcasting($TA)

From f0a8179c1a03f55acc890167ea704eadf882de66 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Thu, 14 Jan 2021 00:22:06 +0000
Subject: [PATCH 08/14] update comments and fix other issues

---
 src/tarray.jl | 27 ++++++++++++++-------------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/src/tarray.jl b/src/tarray.jl
index d228fb5f..9c63867a 100644
--- a/src/tarray.jl
+++ b/src/tarray.jl
@@ -7,7 +7,13 @@
 
 Implementation of data structures that automatically perform copy-on-write after task copying.
 
-If current_task is an existing key in `s`, then return `s[current_task]`. Otherwise, return `s[current_task] = s[last_task]`.
+Underlying `TArray`, we have an concrete array object for each task, i,e, if you access (read
+or write) a same `TArray` is different tasks, you may get different result because you are
+eventaully accessing different arrays wrapped by the `TArray` object.
+
+Now we store the underlying arrays in the field `TArray.data`, a dict whose keys are tasks,
+instead of storing them in the task local storage, because the latter way can't provide
+assistance for type inference thus may bring a perfomance penalty.
 
 Usage:
 
@@ -56,7 +62,7 @@ TArray(x::AbstractArray) = convert(TArray, x)
 const TArrayKeeper = Vector{WeakRef}()
 keep(x::TArray) = push!(TArrayKeeper, WeakRef(x))
 function copy_tarrays(task1::Task, task2::Task)
-    deleteat!(TArrayKeeper, map(x -> x == nothing, TArrayKeeper))
+    filter!(x -> x != nothing, TArrayKeeper)
     for wref in TArrayKeeper
         ta = wref.value
         if haskey(ta.data, task1) && !haskey(ta.data, task2)
@@ -66,17 +72,17 @@ function copy_tarrays(task1::Task, task2::Task)
 end
 
 # _local_storage
-_local_storage(x::TArray{T, N}) where {T, N} = x.data[current_task()]
+_local_storage(x::TArray) = x.data[current_task()]
 function _local_storage(x::TArray{T, N}, v::Tuple{Int, AbstractArray{T, N}}) where {T, N}
     x.data[current_task()] = v
 end
 
 _get(x) = x
-function _get(x::TArray{T, N}) where {T, N}
+function _get(x::TArray)
     n, d = x.data[current_task()]
     return d
 end
-function _get_for_write(x::TArray{T, N}) where {T, N}
+function _get_for_write(x::TArray)
     n, d = x.data[current_task()]
     cn   = n_copies()
     newd = d
@@ -178,12 +184,7 @@ function Base.show(io::IO, ::MIME"text/plain", x::TArray)
     show(io,  MIME("text/plain"), arr)
 end
 
-function Base.show(io::IO, x::TArray)
-    if haskey(x.data, current_task())
-        return Base.show(io::IO, _get(x))
-    end
-    show(io, MIME("text/plain"), x)
-end
+Base.show(io::IO, x::TArray) = show(io::IO, _get(x))
 
 function Base.summary(io::IO, x::TArray)
   print(io, "Task Local Array: ")
@@ -225,12 +226,12 @@ Base.@propagate_inbounds function Base.setindex!(x::TArray{T, N}, e, I::Vararg{I
     d[I...] = e
 end
 
-function Base.push!(x::TArray{T, N}, e) where {T, N}
+function Base.push!(x::TArray, e)
     d = _get_for_write(x)
     push!(d, e)
 end
 
-function Base.pop!(x::TArray{T, N}) where {T, N}
+function Base.pop!(x::TArray)
     d = _get_for_write(x)
     pop!(d)
 end

From 13489c8b17a9b7e363ecc016147720002a5ec030 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Thu, 14 Jan 2021 10:38:14 +0000
Subject: [PATCH 09/14] more performant way to perdicate if a weakref is empty

---
 src/tarray.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/tarray.jl b/src/tarray.jl
index 9c63867a..3b62d33f 100644
--- a/src/tarray.jl
+++ b/src/tarray.jl
@@ -62,7 +62,7 @@ TArray(x::AbstractArray) = convert(TArray, x)
 const TArrayKeeper = Vector{WeakRef}()
 keep(x::TArray) = push!(TArrayKeeper, WeakRef(x))
 function copy_tarrays(task1::Task, task2::Task)
-    filter!(x -> x != nothing, TArrayKeeper)
+    filter!(x -> x.value !== nothing, TArrayKeeper)
     for wref in TArrayKeeper
         ta = wref.value
         if haskey(ta.data, task1) && !haskey(ta.data, task2)

From ea7fbc545f432c00e349415f63d1573d800089b1 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Thu, 14 Jan 2021 13:48:46 +0000
Subject: [PATCH 10/14] underlying array type as a parametric type of TArray

---
 src/tarray.jl | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/src/tarray.jl b/src/tarray.jl
index 3b62d33f..a2cf2350 100644
--- a/src/tarray.jl
+++ b/src/tarray.jl
@@ -29,11 +29,11 @@ for i in 1:4 ta[i] = i end  # assign
 Array(ta)                   # convert to 4-element Array{Int64,1}: [1, 2, 3, 4]
 ```
 """
-struct TArray{T, N} <: AbstractArray{T, N}
+struct TArray{T, N, A <: AbstractArray{T, N}} <: AbstractArray{T, N}
     orig_task :: Task
-    data::Dict{Task, Tuple{Int, AbstractArray{T, N}}}
-    function TArray{T,N}() where {T,N}
-        d = Dict{Task, Tuple{Int, AbstractArray{T, N}}}()
+    data::Dict{Task, Tuple{Int, A}}
+    function TArray{T, N, A}() where {T, N, A <: AbstractArray{T, N}}
+        d = Dict{Task, Tuple{Int, A}}()
         res = new(current_task(), d)
         keep(res)
         return res
@@ -48,7 +48,8 @@ TArray{T,N}(::UndefInitializer, d::Vararg{<:Integer,N}) where {T,N} = TArray{T,N
 TArray{T,N}(dim::NTuple{N,Int}) where {T,N} = TArray(T, dim)
 
 function TArray(T::Type, dim)
-    res = TArray{T, length(dim)}();
+    N_dim = length(dim)
+    res = TArray{T, N_dim, Array{T, N_dim}}()
     n = n_copies()
     d = Array{T}(undef, dim)
     _local_storage(res, (n, d))
@@ -116,7 +117,7 @@ Array(tz)                   # convert to 4-element Array{Int64,1}: [0, 0, 0, 0]
 ```
 """
 function tzeros(T::Type, dim)
-    res = TArray{T,length(dim)}();
+    res = TArray{T,length(dim), Array{T, length(dim)}}();
     n = n_copies()
     d = zeros(T,dim)
     _local_storage(res, (n, d))
@@ -144,7 +145,7 @@ Array(tz)                     # convert to 4-element Array{Float64,1}:  [9.0  9.
 ```
 """
 function tfill(val::Real, dim)
-    res = TArray{typeof(val),length(dim)}();
+    res = TArray{typeof(val), length(dim), Array{typeof(val), length(dim)}}();
     n = n_copies()
     d = fill(val,dim)
     _local_storage(res, (n, d))
@@ -158,7 +159,7 @@ end
 function Base.convert(::Type{Array}, x::TArray)
     return convert(Array{eltype(x), ndims(x)}, x)
 end
-function Base.convert(::Type{Array{T,N}}, x::TArray{T,N}) where {T,N}
+function Base.convert(::Type{Array{T, N}}, x::TArray{T, N}) where {T, N}
     c = convert(Array{T, N}, deepcopy(_get(x)))
     return c
 end
@@ -166,8 +167,8 @@ end
 function Base.convert(::Type{TArray}, x::AbstractArray)
     return convert(TArray{eltype(x), ndims(x)}, x)
 end
-function Base.convert(::Type{TArray{T,N}}, x::AbstractArray{T, N}) where {T, N}
-    res = TArray{T, N}()
+function Base.convert(::Type{TArray{T, N}}, x::AbstractArray{T, N}) where {T, N}
+    res = TArray{T, N, typeof(x)}()
     n   = n_copies()
     _local_storage(res, (n, x))
     return res
@@ -282,7 +283,7 @@ Base.:*(x::AbstractArray, y::TArray) = x * _get(y) |> localize
 Base.:*(x::TArray, y::AbstractArray) = _get(x) * y |> localize
 
 # broadcast
-Base.BroadcastStyle(::Type{TArray{T, N}}) where {T, N} = Broadcast.ArrayStyle{TArray}()
+Base.BroadcastStyle(::Type{TArray{T, N, A}}) where {T, N, A} = Broadcast.ArrayStyle{TArray}()
 Broadcast.broadcasted(::Broadcast.ArrayStyle{TArray}, f, args...) = f.(_get.(args)...) |> localize
 
 import LinearAlgebra

From 6e192dba733253e74ceba8e71d320c5abb1ea2d6 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Thu, 14 Jan 2021 14:51:08 +0000
Subject: [PATCH 11/14] minor update

---
 src/tarray.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/tarray.jl b/src/tarray.jl
index a2cf2350..ac573890 100644
--- a/src/tarray.jl
+++ b/src/tarray.jl
@@ -185,7 +185,7 @@ function Base.show(io::IO, ::MIME"text/plain", x::TArray)
     show(io,  MIME("text/plain"), arr)
 end
 
-Base.show(io::IO, x::TArray) = show(io::IO, _get(x))
+Base.show(io::IO, x::TArray) = show(io, _get(x))
 
 function Base.summary(io::IO, x::TArray)
   print(io, "Task Local Array: ")

From aaf00d2f568f06286acd495b6ef6c53bc5cc0ab1 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Tue, 26 Jan 2021 22:24:09 +0000
Subject: [PATCH 12/14] update according to code review

---
 src/tarray.jl  | 165 +++++++++++++++++++++++++------------------------
 test/tarray.jl |   4 +-
 2 files changed, 85 insertions(+), 84 deletions(-)

diff --git a/src/tarray.jl b/src/tarray.jl
index ac573890..3ddbb64e 100644
--- a/src/tarray.jl
+++ b/src/tarray.jl
@@ -7,13 +7,15 @@
 
 Implementation of data structures that automatically perform copy-on-write after task copying.
 
-Underlying `TArray`, we have an concrete array object for each task, i,e, if you access (read
-or write) a same `TArray` is different tasks, you may get different result because you are
-eventaully accessing different arrays wrapped by the `TArray` object.
+For each `TArray` object, we have a task-local array object, i.e. if
+you access (read or write) the same `TArray` object in different
+tasks, you may be dealing with different array objects each belonging
+to a different task.  These task-specific array objects, however,
+share the same parent `TArray` object.
 
-Now we store the underlying arrays in the field `TArray.data`, a dict whose keys are tasks,
-instead of storing them in the task local storage, because the latter way can't provide
-assistance for type inference thus may bring a perfomance penalty.
+More specifically, we store the underlying arrays in the field
+`TArray.data`, a dictionary whose keys are tasks, instead of storing
+them in the task local storage, for performance considerations.
 
 Usage:
 
@@ -35,7 +37,7 @@ struct TArray{T, N, A <: AbstractArray{T, N}} <: AbstractArray{T, N}
     function TArray{T, N, A}() where {T, N, A <: AbstractArray{T, N}}
         d = Dict{Task, Tuple{Int, A}}()
         res = new(current_task(), d)
-        keep(res)
+        register_to_keeper(res)
         return res
     end
 end
@@ -46,13 +48,14 @@ TArray{T}(::UndefInitializer, dim::NTuple{N,Int}) where {T,N} = TArray(T, dim)
 TArray{T,N}(d::Vararg{<:Integer,N}) where {T,N} = TArray(T, d)
 TArray{T,N}(::UndefInitializer, d::Vararg{<:Integer,N}) where {T,N} = TArray{T,N}(d)
 TArray{T,N}(dim::NTuple{N,Int}) where {T,N} = TArray(T, dim)
+TArray{T, N, Array{T, N}}(::UndefInitializer, dim::NTuple{N,Int}) where {T,N} = TArray(T, dim)
 
 function TArray(T::Type, dim)
     N_dim = length(dim)
     res = TArray{T, N_dim, Array{T, N_dim}}()
     n = n_copies()
     d = Array{T}(undef, dim)
-    _local_storage(res, (n, d))
+    _set_local_storage(res, n, d)
     res
 end
 
@@ -61,7 +64,7 @@ TArray(x::AbstractArray) = convert(TArray, x)
 # TArray House-Keeper
 
 const TArrayKeeper = Vector{WeakRef}()
-keep(x::TArray) = push!(TArrayKeeper, WeakRef(x))
+register_to_keeper(x::TArray) = push!(TArrayKeeper, WeakRef(x))
 function copy_tarrays(task1::Task, task2::Task)
     filter!(x -> x.value !== nothing, TArrayKeeper)
     for wref in TArrayKeeper
@@ -73,28 +76,24 @@ function copy_tarrays(task1::Task, task2::Task)
 end
 
 # _local_storage
-_local_storage(x::TArray) = x.data[current_task()]
-function _local_storage(x::TArray{T, N}, v::Tuple{Int, AbstractArray{T, N}}) where {T, N}
-    x.data[current_task()] = v
-end
 
-_get(x) = x
-function _get(x::TArray)
-    n, d = x.data[current_task()]
-    return d
-end
-function _get_for_write(x::TArray)
+_get_local_storage(x) = x
+function _get_local_storage(x::TArray; copydata=false)
     n, d = x.data[current_task()]
+    copydata || return d
     cn   = n_copies()
     newd = d
     if cn > n
-        # println("[setindex!]: $(x.ref) copying data")
         newd = deepcopy(d)
-        x.data[current_task()] = (cn, newd)
+        _set_local_storage(x, cn, newd)
     end
     return newd
 end
 
+function _set_local_storage(x::TArray{T, N}, n_copies::Int, d::AbstractArray{T, N}) where {T, N}
+    x.data[current_task()] = (n_copies, d)
+end
+
 localize(x) = x
 localize(x::AbstractArray) = TArray(x)
 
@@ -120,7 +119,7 @@ function tzeros(T::Type, dim)
     res = TArray{T,length(dim), Array{T, length(dim)}}();
     n = n_copies()
     d = zeros(T,dim)
-    _local_storage(res, (n, d))
+    _set_local_storage(res, n, d)
     return res
 end
 
@@ -148,7 +147,7 @@ function tfill(val::Real, dim)
     res = TArray{typeof(val), length(dim), Array{typeof(val), length(dim)}}();
     n = n_copies()
     d = fill(val,dim)
-    _local_storage(res, (n, d))
+    _set_local_storage(res, n, d)
     return res
 end
 
@@ -160,7 +159,7 @@ function Base.convert(::Type{Array}, x::TArray)
     return convert(Array{eltype(x), ndims(x)}, x)
 end
 function Base.convert(::Type{Array{T, N}}, x::TArray{T, N}) where {T, N}
-    c = convert(Array{T, N}, deepcopy(_get(x)))
+    c = convert(Array{T, N}, deepcopy(_get_local_storage(x)))
     return c
 end
 
@@ -170,7 +169,7 @@ end
 function Base.convert(::Type{TArray{T, N}}, x::AbstractArray{T, N}) where {T, N}
     res = TArray{T, N, typeof(x)}()
     n   = n_copies()
-    _local_storage(res, (n, x))
+    _set_local_storage(res, n, x)
     return res
 end
 
@@ -185,11 +184,11 @@ function Base.show(io::IO, ::MIME"text/plain", x::TArray)
     show(io,  MIME("text/plain"), arr)
 end
 
-Base.show(io::IO, x::TArray) = show(io, _get(x))
+Base.show(io::IO, x::TArray) = show(io, _get_local_storage(x))
 
 function Base.summary(io::IO, x::TArray)
   print(io, "Task Local Array: ")
-  summary(io, _get(x))
+  summary(io, _get_local_storage(x))
 end
 
 #
@@ -198,19 +197,19 @@ end
 for F in (:size,
           :iterate,
           :firstindex, :lastindex, :axes)
-    @eval Base.$F(a::TArray, args...) = $F(_get(a), args...)
+    @eval Base.$F(a::TArray, args...) = $F(_get_local_storage(a), args...)
 end
 
 #
 # Similarity implementation
 #
 
-Base.similar(x::TArray, ::Type{T}, dims::Dims) where T = TArray(similar(_get(x), T, dims))
+Base.similar(x::TArray, ::Type{T}, dims::Dims) where T = TArray(similar(_get_local_storage(x), T, dims))
 
 for op in [:(==), :≈]
-    @eval Base.$op(x::TArray, y::AbstractArray) = Base.$op(_get(x), y)
-    @eval Base.$op(x::AbstractArray, y::TArray) = Base.$op(x, _get(y))
-    @eval Base.$op(x::TArray, y::TArray) = Base.$op(_get(x), _get(y))
+    @eval Base.$op(x::TArray, y::AbstractArray) = Base.$op(_get_local_storage(x), y)
+    @eval Base.$op(x::AbstractArray, y::TArray) = Base.$op(x, _get_local_storage(y))
+    @eval Base.$op(x::TArray, y::TArray) = Base.$op(_get_local_storage(x), _get_local_storage(y))
 end
 
 #
@@ -219,91 +218,93 @@ end
 
 # Indexing Interface
 Base.@propagate_inbounds function Base.getindex(x::TArray{T, N}, I::Vararg{Int,N}) where {T, N}
-    return _get(x)[I...]
+    return _get_local_storage(x)[I...]
 end
 
 Base.@propagate_inbounds function Base.setindex!(x::TArray{T, N}, e, I::Vararg{Int,N}) where {T, N}
-    d = _get_for_write(x)
+    d = _get_local_storage(x; copydata=true)
     d[I...] = e
 end
 
 function Base.push!(x::TArray, e)
-    d = _get_for_write(x)
+    d = _get_local_storage(x; copydata=true)
     push!(d, e)
 end
 
 function Base.pop!(x::TArray)
-    d = _get_for_write(x)
+    d = _get_local_storage(x; copydata=true)
     pop!(d)
 end
 
 # Other methods from stdlib
 
 Base.view(x::TArray, inds...; kwargs...) =
-    Base.view(_get(x), inds...; kwargs...) |> localize
-Base.:-(x::TArray) = (- _get(x)) |> localize
-Base.transpose(x::TArray) = transpose(_get(x)) |> localize
-Base.adjoint(x::TArray) = adjoint(_get(x)) |> localize
-Base.repeat(x::TArray; kw...) = repeat(_get(x); kw...) |> localize
+    Base.view(_get_local_storage(x), inds...; kwargs...) |> localize
+Base.:-(x::TArray) = (- _get_local_storage(x)) |> localize
+Base.transpose(x::TArray) = transpose(_get_local_storage(x)) |> localize
+Base.adjoint(x::TArray) = adjoint(_get_local_storage(x)) |> localize
+Base.repeat(x::TArray; kw...) = repeat(_get_local_storage(x); kw...) |> localize
 
 Base.hcat(xs::Union{TArray{T,1}, TArray{T,2}}...) where T =
-    hcat(_get.(xs)...) |> localize
+    hcat(_get_local_storage.(xs)...) |> localize
 Base.vcat(xs::Union{TArray{T,1}, TArray{T,2}}...) where T =
-    vcat(_get.(xs)...) |> localize
+    vcat(_get_local_storage.(xs)...) |> localize
 Base.cat(xs::Union{TArray{T,1}, TArray{T,2}}...; dims) where T =
-    cat(_get.(xs)...; dims = dims) |> localize
+    cat(_get_local_storage.(xs)...; dims = dims) |> localize
 
 
-Base.reshape(x::TArray, dims::Union{Colon,Int}...) = reshape(_get(x), dims) |> localize
+Base.reshape(x::TArray, dims::Union{Colon,Int}...) = reshape(_get_local_storage(x), dims) |> localize
 Base.reshape(x::TArray, dims::Tuple{Vararg{Union{Int,Colon}}}) =
-    reshape(_get(x), Base._reshape_uncolon(_get(x), dims)) |> localize
-Base.reshape(x::TArray, dims::Tuple{Vararg{Int}}) = reshape(_get(x), dims) |> localize
-
-Base.permutedims(x::TArray, perm) = permutedims(_get(x), perm) |> localize
-Base.PermutedDimsArray(x::TArray, perm) = PermutedDimsArray(_get(x), perm) |> localize
-Base.reverse(x::TArray; dims) = reverse(_get(x), dims = dims) |> localize
-
-Base.sum(x::TArray; dims = :) = sum(_get(x), dims = dims) |> localize
-Base.sum(f::Union{Function,Type},x::TArray) = sum(f.(_get(x))) |> localize
-Base.prod(x::TArray; dims=:) = prod(_get(x); dims=dims) |> localize
-Base.prod(f::Union{Function, Type}, x::TArray) = prod(f.(_get(x))) |> localize
-
-Base.findfirst(x::TArray, args...) = findfirst(_get(x), args...) |> localize
-Base.maximum(x::TArray; dims = :) = maximum(_get(x), dims = dims) |> localize
-Base.minimum(x::TArray; dims = :) = minimum(_get(x), dims = dims) |> localize
-
-Base.:/(x::TArray, y::TArray) = _get(x) / _get(y) |> localize
-Base.:/(x::AbstractArray, y::TArray) = x / _get(y) |> localize
-Base.:/(x::TArray, y::AbstractArray) = _get(x) / y |> localize
-Base.:\(x::TArray, y::TArray) = _get(x) \ _get(y) |> localize
-Base.:\(x::AbstractArray, y::TArray) = x \ _get(y) |> localize
-Base.:\(x::TArray, y::AbstractArray) = _get(x) \ y |> localize
-Base.:*(x::TArray, y::TArray) = _get(x) * _get(y) |> localize
-Base.:*(x::AbstractArray, y::TArray) = x * _get(y) |> localize
-Base.:*(x::TArray, y::AbstractArray) = _get(x) * y |> localize
+    reshape(_get_local_storage(x), Base._reshape_uncolon(_get_local_storage(x), dims)) |> localize
+Base.reshape(x::TArray, dims::Tuple{Vararg{Int}}) = reshape(_get_local_storage(x), dims) |> localize
+
+Base.permutedims(x::TArray, perm) = permutedims(_get_local_storage(x), perm) |> localize
+Base.PermutedDimsArray(x::TArray, perm) = PermutedDimsArray(_get_local_storage(x), perm) |> localize
+Base.reverse(x::TArray; dims) = reverse(_get_local_storage(x), dims = dims) |> localize
+
+Base.sum(x::TArray; dims = :) = sum(_get_local_storage(x), dims = dims) |> localize
+Base.sum(f::Union{Function,Type},x::TArray) = sum(f.(_get_local_storage(x))) |> localize
+Base.prod(x::TArray; dims=:) = prod(_get_local_storage(x); dims=dims) |> localize
+Base.prod(f::Union{Function, Type}, x::TArray) = prod(f.(_get_local_storage(x))) |> localize
+
+Base.findfirst(x::TArray, args...) = findfirst(_get_local_storage(x), args...) |> localize
+Base.maximum(x::TArray; dims = :) = maximum(_get_local_storage(x), dims = dims) |> localize
+Base.minimum(x::TArray; dims = :) = minimum(_get_local_storage(x), dims = dims) |> localize
+
+Base.:/(x::TArray, y::TArray) = _get_local_storage(x) / _get_local_storage(y) |> localize
+Base.:/(x::AbstractArray, y::TArray) = x / _get_local_storage(y) |> localize
+Base.:/(x::TArray, y::AbstractArray) = _get_local_storage(x) / y |> localize
+Base.:\(x::TArray, y::TArray) = _get_local_storage(x) \ _get_local_storage(y) |> localize
+Base.:\(x::AbstractArray, y::TArray) = x \ _get_local_storage(y) |> localize
+Base.:\(x::TArray, y::AbstractArray) = _get_local_storage(x) \ y |> localize
+Base.:*(x::TArray, y::TArray) = _get_local_storage(x) * _get_local_storage(y) |> localize
+Base.:*(x::AbstractArray, y::TArray) = x * _get_local_storage(y) |> localize
+Base.:*(x::TArray, y::AbstractArray) = _get_local_storage(x) * y |> localize
 
 # broadcast
-Base.BroadcastStyle(::Type{TArray{T, N, A}}) where {T, N, A} = Broadcast.ArrayStyle{TArray}()
-Broadcast.broadcasted(::Broadcast.ArrayStyle{TArray}, f, args...) = f.(_get.(args)...) |> localize
+Base.BroadcastStyle(::Type{TArray{T, N, A}}) where {T, N, A} = Broadcast.ArrayStyle{TArray{T, N, A}}()
+# Broadcast.broadcasted(::Broadcast.ArrayStyle{TArray}, f, args...) = f.(_get_local_storage.(args)...) |> localize
+Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{TArray{T, N, A}}}, ::Type{T}) where {T, N, A} =
+    similar(TArray{T, N, Array{T, N}}, axes(bc))
 
 import LinearAlgebra
 import LinearAlgebra:  \, /, inv, det, logdet, logabsdet, norm
 
-LinearAlgebra.inv(x::TArray) = inv(_get(x)) |> localize
-LinearAlgebra.det(x::TArray) = det(_get(x)) |> localize
-LinearAlgebra.logdet(x::TArray) = logdet(_get(x)) |> localize
-LinearAlgebra.logabsdet(x::TArray) = logabsdet(_get(x)) |> localize
+LinearAlgebra.inv(x::TArray) = inv(_get_local_storage(x)) |> localize
+LinearAlgebra.det(x::TArray) = det(_get_local_storage(x)) |> localize
+LinearAlgebra.logdet(x::TArray) = logdet(_get_local_storage(x)) |> localize
+LinearAlgebra.logabsdet(x::TArray) = logabsdet(_get_local_storage(x)) |> localize
 LinearAlgebra.norm(x::TArray, p::Real = 2) =
-    LinearAlgebra.norm(_get(x), p) |> localize
+    LinearAlgebra.norm(_get_local_storage(x), p) |> localize
 
 import LinearAlgebra: dot
-dot(x::TArray, ys::TArray) = dot(_get(x), _get(ys)) |> localize
-dot(x::AbstractArray, ys::TArray) = dot(x, _get(ys)) |> localize
-dot(x::TArray, ys::AbstractArray) = dot(_get(x), ys) |> localize
+dot(x::TArray, ys::TArray) = dot(_get_local_storage(x), _get_local_storage(ys)) |> localize
+dot(x::AbstractArray, ys::TArray) = dot(x, _get_local_storage(ys)) |> localize
+dot(x::TArray, ys::AbstractArray) = dot(_get_local_storage(x), ys) |> localize
 
 using Statistics
-Statistics.mean(x::TArray; dims = :) = mean(_get(x), dims = dims) |> localize
-Statistics.std(x::TArray; kw...) = std(_get(x), kw...) |> localize
+Statistics.mean(x::TArray; dims = :) = mean(_get_local_storage(x), dims = dims) |> localize
+Statistics.std(x::TArray; kw...) = std(_get_local_storage(x), kw...) |> localize
 
 # TODO
 # * NNlib
diff --git a/test/tarray.jl b/test/tarray.jl
index 0ce037a6..9c20bd90 100644
--- a/test/tarray.jl
+++ b/test/tarray.jl
@@ -55,7 +55,7 @@
         @test_throws MethodError TArray{Int,2}(undef, 4)
 
         ta3 = TArray{Int, 4}(4, 3, 2, 1)
-        ta4 = Libtask._get(ta3)
+        ta4 = Libtask._get_local_storage(ta3)
         @test ta3[3] == ta4[3]
 
         ta5 = TArray{Int}(4)
@@ -116,7 +116,7 @@
 
         @test repeat(ta, 1, 2) == hcat(ta, ta)
 
-        @test ta .+ ta == Libtask._get(ta) .+ Libtask._get(ta)
+        @test ta .+ ta == Libtask._get_local_storage(ta) .+ Libtask._get_local_storage(ta)
     end
 
     @testset "task copy" begin

From bddfc26b7bba2615fb99770f40a5447255e9698c Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Wed, 27 Jan 2021 20:20:02 +0000
Subject: [PATCH 13/14] fix broadcasting

---
 src/tarray.jl | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/src/tarray.jl b/src/tarray.jl
index 3ddbb64e..5212646f 100644
--- a/src/tarray.jl
+++ b/src/tarray.jl
@@ -48,7 +48,6 @@ TArray{T}(::UndefInitializer, dim::NTuple{N,Int}) where {T,N} = TArray(T, dim)
 TArray{T,N}(d::Vararg{<:Integer,N}) where {T,N} = TArray(T, d)
 TArray{T,N}(::UndefInitializer, d::Vararg{<:Integer,N}) where {T,N} = TArray{T,N}(d)
 TArray{T,N}(dim::NTuple{N,Int}) where {T,N} = TArray(T, dim)
-TArray{T, N, Array{T, N}}(::UndefInitializer, dim::NTuple{N,Int}) where {T,N} = TArray(T, dim)
 
 function TArray(T::Type, dim)
     N_dim = length(dim)
@@ -282,10 +281,8 @@ Base.:*(x::AbstractArray, y::TArray) = x * _get_local_storage(y) |> localize
 Base.:*(x::TArray, y::AbstractArray) = _get_local_storage(x) * y |> localize
 
 # broadcast
-Base.BroadcastStyle(::Type{TArray{T, N, A}}) where {T, N, A} = Broadcast.ArrayStyle{TArray{T, N, A}}()
-# Broadcast.broadcasted(::Broadcast.ArrayStyle{TArray}, f, args...) = f.(_get_local_storage.(args)...) |> localize
-Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{TArray{T, N, A}}}, ::Type{T}) where {T, N, A} =
-    similar(TArray{T, N, Array{T, N}}, axes(bc))
+Base.BroadcastStyle(::Type{<:TArray}) = Broadcast.ArrayStyle{TArray}()
+Broadcast.broadcasted(::Broadcast.ArrayStyle{TArray}, f, args...) = f.(_get_local_storage.(args)...) |> localize
 
 import LinearAlgebra
 import LinearAlgebra:  \, /, inv, det, logdet, logabsdet, norm

From ad7ee1234ed9912f824f86ce44d579c55b8d6615 Mon Sep 17 00:00:00 2001
From: Hong Ge <hg344@cam.ac.uk>
Date: Thu, 28 Jan 2021 16:02:33 +0000
Subject: [PATCH 14/14] Update Project.toml

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 213b838d..1cabdfa6 100644
--- a/Project.toml
+++ b/Project.toml
@@ -3,7 +3,7 @@ uuid = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f"
 license = "MIT"
 desc = "C shim for task copying in Turing"
 repo = "https://github.com/TuringLang/Libtask.jl.git"
-version = "0.5.1"
+version = "0.5.2"
 
 [deps]
 Libtask_jll = "3ae2931a-708c-5973-9c38-ccf7496fb450"