support blocks, better precompilation / less compilation

Jutho · Feb 1, 2019 · 8e4db13 · 8e4db13
1 parent 1d32e9e
commit 8e4db13
Show file tree

Hide file tree

Showing 4 changed files with 116 additions and 50 deletions.
diff --git a/src/TensorOperations.jl b/src/TensorOperations.jl
@@ -109,45 +109,108 @@ cachesize() = cache.currentsize
 # Some precompile statements
 #----------------------------
 function _precompile_()
+    AVector = Vector{Any}
     ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
-    precompile(Tuple{typeof(_findfirst), typeof(identity), Int})
-    precompile(Tuple{typeof(_findnext), typeof(identity), Int})
-    precompile(Tuple{typeof(_ncontree!), Int, Int})
-    precompile(Tuple{typeof(conjexpr), Expr})
-    precompile(Tuple{typeof(deindexify_contraction), Int, Int, Expr, Int, Vector, Vector, Int})
-    precompile(Tuple{typeof(deindexify_generaltensor), Int, Int, Expr, Int, Vector, Vector, Int})
-    precompile(Tuple{typeof(deindexify_linearcombination), Int, Int, Expr, Int, Vector, Vector, Int})
-    precompile(Tuple{typeof(deindexify), Int, Int, Expr, Int, Vector, Vector, Int})
-    precompile(Tuple{typeof(deindexify), Int, Int, Expr, Int, Vector, Vector})
-    precompile(Tuple{typeof(expandconj), Expr})
-    precompile(Tuple{typeof(expandconj), Int})
-    precompile(Tuple{typeof(getallindices), Expr})
-    precompile(Tuple{typeof(getallindices), Int})
-    precompile(Tuple{typeof(geteltype), Expr})
-    precompile(Tuple{typeof(geteltype), Int})
-    precompile(Tuple{typeof(getindices), Expr})
-    precompile(Tuple{typeof(hastraceindices), Int})
-    precompile(Tuple{typeof(isgeneraltensor), Expr})
-    precompile(Tuple{typeof(isindex), Int})
-    precompile(Tuple{typeof(isnconstyle), Vector})
-    precompile(Tuple{typeof(isscalarexpr), Expr})
-    precompile(Tuple{typeof(istensor), Expr})
-    precompile(Tuple{typeof(istensorexpr), Int})
-    precompile(Tuple{typeof(makegeneraltensor), Int})
-    precompile(Tuple{typeof(makeindex), Int})
-    precompile(Tuple{typeof(makescalar), Expr})
-    precompile(Tuple{typeof(maketensor), Int})
-    precompile(Tuple{typeof(ncontree), Vector})
-    precompile(Tuple{typeof(processcontractorder), Expr, Int})
-    precompile(Tuple{typeof(processcontractorder), Int, Int})
-    precompile(Tuple{typeof(tensorify), Expr, Int})
-    precompile(Tuple{typeof(tensorify), Expr})
-    precompile(Tuple{typeof(tensorify), Int, Int})
-    precompile(Tuple{typeof(tree2expr), Int, Int})
-    precompile(Tuple{typeof(unique2), Array{Any, 1}})
-    precompile(Tuple{typeof(unique2), Array{Int64, 1}})
-    precompile(Tuple{typeof(use_blas)})
-    precompile(Tuple{typeof(use_cache)})
+    @assert precompile(Tuple{typeof(_findfirst), Base.Fix2{typeof(Base.isequal), Symbol}, AVector})
+    @assert precompile(Tuple{typeof(_findnext), Base.Fix2{typeof(Base.isequal), Int64}, AVector, Int64})
+    @assert precompile(Tuple{typeof(_intersect), Base.BitArray{1}, Base.BitArray{1}})
+    @assert precompile(Tuple{typeof(_intersect), Base.BitSet, Base.BitSet})
+    @assert precompile(Tuple{typeof(_intersect), UInt128, UInt128})
+    @assert precompile(Tuple{typeof(_intersect), UInt32, UInt32})
+    @assert precompile(Tuple{typeof(_intersect), UInt64, UInt64})
+    @assert precompile(Tuple{typeof(_isemptyset), Base.BitArray{1}})
+    @assert precompile(Tuple{typeof(_isemptyset), Base.BitSet})
+    @assert precompile(Tuple{typeof(_isemptyset), UInt128})
+    @assert precompile(Tuple{typeof(_isemptyset), UInt32})
+    @assert precompile(Tuple{typeof(_isemptyset), UInt64})
+    @assert precompile(Tuple{typeof(_ncontree!), AVector, Vector{Vector{Int64}}})
+    @assert precompile(Tuple{typeof(_setdiff), Base.BitArray{1}, Base.BitArray{1}})
+    @assert precompile(Tuple{typeof(_setdiff), Base.BitSet, Base.BitSet})
+    @assert precompile(Tuple{typeof(_setdiff), UInt128, UInt128})
+    @assert precompile(Tuple{typeof(_setdiff), UInt32, UInt32})
+    @assert precompile(Tuple{typeof(_setdiff), UInt64, UInt64})
+    @assert precompile(Tuple{typeof(_union), Base.BitArray{1}, Base.BitArray{1}})
+    @assert precompile(Tuple{typeof(_union), Base.BitSet, Base.BitSet})
+    @assert precompile(Tuple{typeof(_union), UInt128, UInt128})
+    @assert precompile(Tuple{typeof(_union), UInt32, UInt32})
+    @assert precompile(Tuple{typeof(_union), UInt64, UInt64})
+    @assert precompile(Tuple{typeof(addcost), Power{:χ, Int64}, Power{:χ, Int64}})
+    @assert precompile(Tuple{typeof(degree), Power{:x, Int64}})
+    @assert precompile(Tuple{typeof(deindexify_contraction), Int, Int, Expr, Int, AVector, AVector, Int})
+    @assert precompile(Tuple{typeof(deindexify_generaltensor), Int, Int, Expr, Int, AVector, AVector, Int})
+    @assert precompile(Tuple{typeof(deindexify_linearcombination), Int, Int, Expr, Int, AVector, AVector, Int})
+    @assert precompile(Tuple{typeof(deindexify), Int, Int, Expr, Int, AVector, AVector, Int})
+    @assert precompile(Tuple{typeof(deindexify), Int, Int, Expr, Int, AVector, AVector})
+    @assert precompile(Tuple{typeof(deindexify), Expr, Bool, Expr, Int64, AVector, AVector})
+    @assert precompile(Tuple{typeof(deindexify), Nothing, Bool, Expr, Bool, AVector, AVector, Bool})
+    @assert precompile(Tuple{typeof(deindexify), Symbol, Bool, Expr, Bool, AVector, AVector})
+    @assert precompile(Tuple{typeof(disable_blas)})
+    @assert precompile(Tuple{typeof(disable_cache)})
+    @assert precompile(Tuple{typeof(enable_blas)})
+    @assert precompile(Tuple{typeof(enable_cache)})
+    @assert precompile(Tuple{typeof(expandconj), Expr})
+    @assert precompile(Tuple{typeof(expandconj), Symbol})
+    @assert precompile(Tuple{typeof(getallindices), Expr})
+    @assert precompile(Tuple{typeof(getallindices), Int})
+    @assert precompile(Tuple{typeof(getallindices), Symbol})
+    @assert precompile(Tuple{typeof(geteltype), Expr})
+    @assert precompile(Tuple{typeof(getindices), Symbol})
+    @assert precompile(Tuple{typeof(getindices), Expr})
+    @assert precompile(Tuple{typeof(getlhsrhs), Expr})
+    @assert precompile(Tuple{typeof(hastraceindices), Expr})
+    @assert precompile(Tuple{typeof(isassignment), Expr})
+    @assert precompile(Tuple{typeof(isdefinition), Expr})
+    @assert precompile(Tuple{typeof(isgeneraltensor), Expr})
+    @assert precompile(Tuple{typeof(isindex), Symbol})
+    @assert precompile(Tuple{typeof(isindex), Int})
+    @assert precompile(Tuple{typeof(isnconstyle), Array{AVector, 1}})
+    @assert precompile(Tuple{typeof(isscalarexpr), Expr})
+    @assert precompile(Tuple{typeof(isscalarexpr), Float64})
+    @assert precompile(Tuple{typeof(isscalarexpr), LineNumberNode})
+    @assert precompile(Tuple{typeof(isscalarexpr), Symbol})
+    @assert precompile(Tuple{typeof(istensor), Expr})
+    @assert precompile(Tuple{typeof(istensorexpr), Expr})
+    @assert precompile(Tuple{typeof(makegeneraltensor), Expr})
+    @assert precompile(Tuple{typeof(makeindex), Int})
+    @assert precompile(Tuple{typeof(makeindex), Symbol})
+    @assert precompile(Tuple{typeof(makescalar), Expr})
+    @assert precompile(Tuple{typeof(makescalar), Float64})
+    @assert precompile(Tuple{typeof(maketensor), Expr})
+    @assert precompile(Tuple{typeof(mulcost), Power{:χ, Int64}, Power{:χ, Int64}})
+    @assert precompile(Tuple{typeof(ncontree), Vector{AVector}})
+    @assert precompile(Tuple{typeof(optdata), Expr})
+    @assert precompile(Tuple{typeof(optimaltree), Vector{AVector}, Base.Dict{Any, Power{:χ, Int64}}})
+    @assert precompile(Tuple{typeof(parsecost), Expr})
+    @assert precompile(Tuple{typeof(parsecost), Int64})
+    @assert precompile(Tuple{typeof(parsecost), Symbol})
+    @assert precompile(Tuple{typeof(processcontractorder), Expr, Nothing})
+    @assert precompile(Tuple{typeof(processcontractorder), Expr, Int})
+    @assert precompile(Tuple{typeof(processcontractorder), Int, Int})
+    @assert precompile(Tuple{typeof(processcontractorder), Symbol, Nothing})
+    @assert precompile(Tuple{typeof(storeset), Type{Base.BitArray{1}}, AVector, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{Base.BitArray{1}}, Array{Int64, 1}, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{Base.BitArray{1}}, Base.Set{Int64}, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{Base.BitSet}, AVector, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{Base.BitSet}, Array{Int64, 1}, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{Base.BitSet}, Base.Set{Int64}, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{UInt128}, AVector, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{UInt128}, Array{Int64, 1}, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{UInt128}, Base.Set{Int64}, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{UInt32}, AVector, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{UInt32}, Array{Int64, 1}, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{UInt32}, Base.Set{Int64}, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{UInt64}, AVector, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{UInt64}, Array{Int64, 1}, Int64})
+    @assert precompile(Tuple{typeof(storeset), Type{UInt64}, Base.Set{Int64}, Int64})
+    @assert precompile(Tuple{typeof(tensorify), Expr, Nothing})
+    @assert precompile(Tuple{typeof(tensorify), Expr, Int})
+    @assert precompile(Tuple{typeof(tensorify), Expr})
+    @assert precompile(Tuple{typeof(tensorify), Int, Int})
+    @assert precompile(Tuple{typeof(tree2expr), Int, Int})
+    @assert precompile(Tuple{typeof(unique2), AVector})
+    @assert precompile(Tuple{typeof(unique2), Array{Int64, 1}})
+    @assert precompile(Tuple{typeof(use_blas)})
+    @assert precompile(Tuple{typeof(use_cache)})
 end
 _precompile_()
 

diff --git a/src/indexnotation/tensorexpressions.jl b/src/indexnotation/tensorexpressions.jl
@@ -148,21 +148,21 @@ end
 function maketensor(ex)
     if isa(ex, Expr) && (ex.head == :ref || ex.head == :typed_hcat)
         object = esc(ex.args[1])
-        leftind = map(makeindex, ex.args[2:end])
+        leftind = Any[makeindex(x) for x in ex.args[2:end]]
         rightind = Any[]
         return (object, leftind, rightind)
     elseif isa(ex, Expr) && ex.head == :typed_vcat
         length(ex.args) <= 3 || throw(ArgumentError("invalid tensor index expression: $ex"))
         object = esc(ex.args[1])
         if isa(ex.args[2], Expr) && ex.args[2].head == :row
-            leftind = map(makeindex, ex.args[2].args)
+            leftind = Any[makeindex(x) for x in ex.args[2].args]
         elseif ex.args[2] == :_
             leftind = Any[]
         else
             leftind = Any[makeindex(ex.args[2])]
         end
         if length(ex.args) > 2 && isa(ex.args[3], Expr) && ex.args[3].head == :row
-            rightind = map(makeindex, ex.args[3].args)
+            rightind = Any[makeindex(x) for x in ex.args[3].args]
         elseif length(ex.args) == 2 || ex.args[3] == :_
             rightind = Any[]
         else

diff --git a/src/indexnotation/tensormacro.jl b/src/indexnotation/tensormacro.jl
@@ -203,6 +203,9 @@ function tensorify(ex::Expr, optdata = nothing)
     if ex.head == :block
         return Expr(ex.head, map(x->tensorify(x, optdata), ex.args)...)
     end
+    if ex.head == :for
+        return Expr(ex.head, esc(ex.args[1]), tensorify(ex.args[2], optdata))
+    end
     # constructions of the form: a = @tensor ...
     if isscalarexpr(ex)
         return makescalar(ex)
@@ -288,7 +291,7 @@ function tree2expr(args, tree)
 end
 
 # deindexify: parse tensor operations
-function deindexify(dst, β, ex::Expr, α, leftind::Vector, rightind::Vector, istemporary = false)
+function deindexify(dst, β, ex::Expr, α, leftind::Vector{Any}, rightind::Vector{Any}, istemporary = false)
     if isgeneraltensor(ex)
         return deindexify_generaltensor(dst, β, ex, α, leftind, rightind, istemporary)
     elseif ex.head == :call && (ex.args[1] == :+ || ex.args[1] == :-) # linear combination
@@ -309,7 +312,7 @@ function deindexify(dst, β, ex::Expr, α, leftind::Vector, rightind::Vector, is
     throw(ArgumentError("problem with parsing $ex"))
 end
 
-function deindexify_generaltensor(dst, β, ex::Expr, α, leftind::Vector, rightind::Vector, istemporary = false)
+function deindexify_generaltensor(dst, β, ex::Expr, α, leftind::Vector{Any}, rightind::Vector{Any}, istemporary = false)
     src, srcleftind, srcrightind, α2, conj = makegeneraltensor(ex)
     srcind = vcat(srcleftind, srcrightind)
     conjarg = conj ? :(:C) : :(:N)
@@ -362,7 +365,7 @@ function deindexify_generaltensor(dst, β, ex::Expr, α, leftind::Vector, righti
         end
     end
 end
-function deindexify_linearcombination(dst, β, ex::Expr, α, leftind::Vector, rightind::Vector, istemporary = false)
+function deindexify_linearcombination(dst, β, ex::Expr, α, leftind::Vector{Any}, rightind::Vector{Any}, istemporary = false)
     if ex.head == :call && (ex.args[1] == :+ || ex.args[1] == :-) # addition: add one by one
         if dst === nothing
             αnew = Expr(:call, :*, α, Expr(:call, :one, geteltype(ex)))
@@ -388,7 +391,7 @@ function deindexify_linearcombination(dst, β, ex::Expr, α, leftind::Vector, ri
         throw(ArgumentError("unable to deindexify linear combination: $ex"))
     end
 end
-function deindexify_contraction(dst, β, ex::Expr, α, leftind::Vector, rightind::Vector, istemporary = false)
+function deindexify_contraction(dst, β, ex::Expr, α, leftind::Vector{Any}, rightind::Vector{Any}, istemporary = false)
     @assert ex.head == :call && ex.args[1] == :* && length(ex.args) == 3 &&
         istensorexpr(ex.args[2]) && istensorexpr(ex.args[3])
     exA = ex.args[2]

diff --git a/test/auxiliary.jl b/test/auxiliary.jl
@@ -46,13 +46,13 @@
         @test !isindex(:(5.1))
 
         @test istensor(:(a[1,2,3]))
-        @test maketensor(:(a[1,2,3])) == (esc(:a), [1,2,3], [])
+        @test maketensor(:(a[1,2,3])) == (esc(:a), Any[1,2,3], Any[])
         @test istensor(:(a[5][a b c]))
-        @test maketensor(:(a[5][a b c])) == (esc(:(a[5])), [:a,:b,:c], [])
+        @test maketensor(:(a[5][a b c])) == (esc(:(a[5])), Any[:a,:b,:c], Any[])
         @test istensor(:(cos(y)[a b c; 1 2 3]))
-        @test maketensor(:(cos(y)[a b c; 1 2 3])) == (esc(:(cos(y))), [:a,:b,:c], [1,2,3])
+        @test maketensor(:(cos(y)[a b c; 1 2 3])) == (esc(:(cos(y))), Any[:a,:b,:c], Any[1,2,3])
         @test istensor(:(x[_; 1 2 3]))
-        @test maketensor(:(x[_; 1 2 3])) == (esc(:x), [], [1,2,3])
+        @test maketensor(:(x[_; 1 2 3])) == (esc(:x), Any[], Any[1,2,3])
         @test !istensor(:(2*a[1,2,3]))
         @test !istensor(:(a[1 2 3; 4 5 6; 7 8 9]))
         @test !istensor(:(conj(a[5][a b c])))
@@ -64,7 +64,7 @@
         @test isgeneraltensor(:(x*a[5][a b c]))
         @test makegeneraltensor(:(x*a[5][a b c])) == (esc(:(a[5])), [:a,:b,:c], [], makescalar(:(x*1)), false)
         @test isgeneraltensor(:(3*conj(a*cos(y)[a b c; 1 2 3])))
-        @test makegeneraltensor(:(3*conj(a*cos(y)[a b c; 1 2 3]))) == (esc(:(cos(y))), [:a,:b,:c], [1,2,3], makescalar(:(3*conj(a*1))), true)
+        @test makegeneraltensor(:(3*conj(a*cos(y)[a b c; 1 2 3]))) == (esc(:(cos(y))),  Any[:a,:b,:c],  Any[1,2,3], makescalar(:(3*conj(a*1))), true)
         @test !isgeneraltensor(:(1/a[1,2,3]))
         @test !isgeneraltensor(:(a[1 2 3; 4 5 6]\x))
         @test !isgeneraltensor(:(cos(y)[a b c; 1 2 3]*b[4,5]))