From 358fa5637233aab396644b0945252e65630b9a07 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo@hexoasis.com>
Date: Fri, 14 Feb 2020 14:00:24 +0800
Subject: [PATCH] Minor code refactor and more test cases for TArray

---
 src/tarray.jl    | 22 ++++++++++++++--------
 src/tref.jl      |  2 +-
 test/runtests.jl |  1 +
 test/tarray3.jl  | 32 ++++++++++++++++++++++++++++++++
 4 files changed, 48 insertions(+), 9 deletions(-)
 create mode 100644 test/tarray3.jl

diff --git a/src/tarray.jl b/src/tarray.jl
index 6f304d32..73b6a3f0 100644
--- a/src/tarray.jl
+++ b/src/tarray.jl
@@ -53,7 +53,7 @@ function Base.getindex(S::TArray{T, N}, I::Vararg{Int,N}) where {T, N}
     return d[I...]
 end
 
-function Base.setindex!(S::TArray{T, N}, x, I::Vararg{Int,N}) where {T, N}
+function Base.setindex!(S::TArray{T, N}, x::T, I::Vararg{Int,N}) where {T, N}
     n, d = task_local_storage(S.ref)
     cn   = n_copies()
     newd = d
@@ -65,7 +65,7 @@ function Base.setindex!(S::TArray{T, N}, x, I::Vararg{Int,N}) where {T, N}
     newd[I...] = x
 end
 
-function Base.push!(S::TArray, x)
+function Base.push!(S::TArray{T}, x::T) where T
     n, d = task_local_storage(S.ref)
     cn   = n_copies()
     newd = d
@@ -102,21 +102,27 @@ end
 
 function Base.display(S::TArray)
     arr = S.orig_task.storage[S.ref][2]
-    @warn "display(::TArray) prints the originating task's storage, not the current task's storage. Please use show(::TArray) to display the current task's version of a TArray."
+    @warn "display(::TArray) prints the originating task's storage, " *
+        "not the current task's storage. " *
+        "Please use show(::TArray) to display the current task's version of a TArray."
     display(arr)
 end
 
 Base.show(io::IO, S::TArray) = Base.show(io::IO, task_local_storage(S.ref)[2])
-Base.size(S::TArray) = Base.size(task_local_storage(S.ref)[2])
-Base.ndims(S::TArray) = Base.ndims(task_local_storage(S.ref)[2])
 
 # Base.get(t::Task, S) = S
 # Base.get(t::Task, S::TArray) = (t.storage[S.ref][2])
 Base.get(S::TArray) = (current_task().storage[S.ref][2])
 
-# Implements eltype, firstindex, lastindex, and iterate
-# functions.
-for F in (:eltype, :firstindex, :lastindex, :iterate)
+##
+# Iterator Interface
+IteratorSize(::Type{TArray{T, N}}) where {T, N} = HasShape{N}()
+IteratorEltype(::Type{TArray}) = HasEltype()
+
+# Implements iterate, eltype, length, and size functions,
+# as well as firstindex, lastindex, ndims, and axes
+for F in (:iterate, :eltype, :length, :size,
+          :firstindex, :lastindex, :ndims, :axes)
     @eval Base.$F(a::TArray, args...) = $F(get(a), args...)
 end
 
diff --git a/src/tref.jl b/src/tref.jl
index 1c0edfed..0b0af90c 100644
--- a/src/tref.jl
+++ b/src/tref.jl
@@ -1,5 +1,5 @@
 ##########
-# TRef #
+#  TRef  #
 ##########
 
 """
diff --git a/test/runtests.jl b/test/runtests.jl
index a8c02f5b..1e381554 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -2,3 +2,4 @@ include("clonetask.jl")
 include("brokentask.jl")
 include("tarray.jl")
 include("tarray2.jl")
+include("tarray3.jl")
diff --git a/test/tarray3.jl b/test/tarray3.jl
new file mode 100644
index 00000000..9fa47bcd
--- /dev/null
+++ b/test/tarray3.jl
@@ -0,0 +1,32 @@
+using Libtask
+using Test
+
+@testset "TArray Copying Test" begin
+    DATA = Dict{Task, Array}()
+
+    function f_ct()
+        ta = tzeros(UInt64, 4);
+        for i in 1:4
+            ta[i] = hash(current_task())
+            DATA[current_task()] = convert(Array, ta)
+            produce(ta[i])
+        end
+    end
+
+    t = CTask(f_ct)
+
+    @test consume(t) == hash(t) # index = 1
+    @test consume(t) == hash(t) # index = 2
+
+    a = copy(t);
+
+    @test consume(a) == hash(a) # index = 3
+    @test consume(a) == hash(a) # index = 4
+
+    @test consume(t) == hash(t) # index = 3
+
+
+    @test DATA[t] == [hash(t), hash(t), hash(t), 0]
+    @test DATA[a] == [hash(t), hash(t), hash(a), hash(a)]
+
+end