first draft for blockedpermutation

Author: ogauthe

src/TensorAlgebra.jl

@@ -1,7 +1,7 @@
 module TensorAlgebra
+include("blockedtuple.jl")
 include("blockedpermutation.jl")
 include("BaseExtensions/BaseExtensions.jl")
-include("blockedtuple.jl")
 include("fusedims.jl")
 include("splitdims.jl")
 include("contract/contract.jl")

src/blockedpermutation.jl

@@ -28,30 +28,12 @@ collect_tuple(t::Tuple) = t
 
 const TupleOfTuples{N} = Tuple{Vararg{Tuple{Vararg{Int}},N}}
 
-abstract type AbstractBlockedPermutation{BlockLength,Length} end
-
-BlockArrays.blocks(blockedperm::AbstractBlockedPermutation) = error("Not implemented")
-
-function Base.Tuple(blockedperm::AbstractBlockedPermutation)
-  return flatten_tuples(blocks(blockedperm))
-end
-
-function BlockArrays.blocklengths(blockedperm::AbstractBlockedPermutation)
-  return length.(blocks(blockedperm))
-end
-
-function BlockArrays.blockfirsts(blockedperm::AbstractBlockedPermutation)
-  return _blockfirsts(blocklengths(blockedperm))
-end
-
-function BlockArrays.blocklasts(blockedperm::AbstractBlockedPermutation)
-  return _blocklasts(blocklengths(blockedperm))
-end
+#
+# ==============================  AbstractBlockedPermutation  ==============================
+#
+abstract type AbstractBlockedPermutation{BlockLength,Length} <: AbstractBlockTuple end
 
-Base.iterate(permblocks::AbstractBlockedPermutation) = iterate(Tuple(permblocks))
-function Base.iterate(permblocks::AbstractBlockedPermutation, state)
-  return iterate(Tuple(permblocks), state)
-end
+widened_constructorof(::Type{<:AbstractBlockedPermutation}) = BlockedTuple
 
 # Block a permutation based on the specified lengths.
 # blockperm((4, 3, 2, 1), (2, 2)) == blockedperm((4, 3), (2, 1))
@@ -67,24 +49,6 @@ function Base.invperm(blockedperm::AbstractBlockedPermutation)
   return blockperm(invperm(Tuple(blockedperm)), blocklengths(blockedperm))
 end
 
-Base.length(blockedperm::AbstractBlockedPermutation) = length(Tuple(blockedperm))
-function BlockArrays.blocklength(blockedperm::AbstractBlockedPermutation)
-  return length(blocks(blockedperm))
-end
-
-function Base.getindex(blockedperm::AbstractBlockedPermutation, i::Int)
-  return Tuple(blockedperm)[i]
-end
-
-function Base.getindex(blockedperm::AbstractBlockedPermutation, I::AbstractUnitRange)
-  perm = Tuple(blockedperm)
-  return [perm[i] for i in I]
-end
-
-function Base.getindex(blockedperm::AbstractBlockedPermutation, b::Block)
-  return blocks(blockedperm)[Int(b)]
-end
-
 # Like `BlockRange`.
 function blockeachindex(blockedperm::AbstractBlockedPermutation)
   return ntuple(i -> Block(i), blocklength(blockedperm))
@@ -148,6 +112,9 @@ function blockedperm_indexin(collection, subs...)
   return blockedperm(map(sub -> BaseExtensions.indexin(sub, collection), subs)...)
 end
 
+#
+# ==================================  BlockedPermutation  ==================================
+#
 struct BlockedPermutation{BlockLength,Length,Blocks<:TupleOfTuples{BlockLength}} <:
        AbstractBlockedPermutation{BlockLength,Length}
   blocks::Blocks
@@ -158,15 +125,28 @@ struct BlockedPermutation{BlockLength,Length,Blocks<:TupleOfTuples{BlockLength}}
   end
 end
 
+Base.Tuple(blockedperm::BlockedPermutation) = flatten_tuples(blocks(blockedperm))
+
+BlockedTuple(bp::BlockedPermutation) = tuplemortar(blocks(bp))
+
 BlockArrays.blocks(blockedperm::BlockedPermutation) = getfield(blockedperm, :blocks)
 
+function BlockArrays.blocklengths(
+  ::Type{<:BlockedPermutation{<:Any,<:Any,Blocks}}
+) where {Blocks}
+  return fieldcount.(fieldtypes(Blocks))
+end
+
 function blockedperm(length::Val, permblocks::Tuple{Vararg{Int}}...)
   @assert value(length) == sum(Base.length, permblocks; init=zero(Bool))
   blockedperm = _BlockedPermutation(permblocks)
   @assert isperm(blockedperm)
   return blockedperm
 end
 
+#
+# ==============================  BlockedTrivialPermutation  ===============================
+#
 trivialperm(length::Union{Integer,Val}) = ntuple(identity, length)
 
 struct BlockedTrivialPermutation{BlockLength,Length,Blocks<:TupleOfTuples{BlockLength}} <:
@@ -180,6 +160,8 @@ struct BlockedTrivialPermutation{BlockLength,Length,Blocks<:TupleOfTuples{BlockL
   end
 end
 
+Base.Tuple(blockedperm::BlockedTrivialPermutation) = flatten_tuples(blocks(blockedperm))
+
 BlockArrays.blocks(blockedperm::BlockedTrivialPermutation) = getfield(blockedperm, :blocks)
 
 function blockedtrivialperm(blocklengths::Tuple{Vararg{Int}})

test/test_basics.jl

@@ -1,71 +1,13 @@
-using BlockArrays: blockfirsts, blocklasts, blocklength, blocklengths, blocks
-using Combinatorics: permutations
 using EllipsisNotation: var".."
 using LinearAlgebra: norm, qr
-using TensorAlgebra: TensorAlgebra, blockedperm, blockedperm_indexin, fusedims, splitdims
+using TensorAlgebra, TensorAlgebra, fusedims, splitdims
 # TODO: Remove dependency on NDTensors, create a GPUTestUtils.jl package.
 using NDTensors: NDTensors
 include(joinpath(pkgdir(NDTensors), "test", "NDTensorsTestUtils", "NDTensorsTestUtils.jl"))
 using .NDTensorsTestUtils: default_rtol
 using Test: @test, @test_broken, @testset
 const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
-@testset "BlockedPermutation" begin
-  p = blockedperm((3, 4, 5), (2, 1))
-  @test Tuple(p) === (3, 4, 5, 2, 1)
-  @test isperm(p)
-  @test length(p) == 5
-  @test blocks(p) == ((3, 4, 5), (2, 1))
-  @test blocklength(p) == 2
-  @test blocklengths(p) == (3, 2)
-  @test blockfirsts(p) == (1, 4)
-  @test blocklasts(p) == (3, 5)
-  @test invperm(p) == blockedperm((5, 4, 1), (2, 3))
 
-  # Empty block.
-  p = blockedperm((3, 2), (), (1,))
-  @test Tuple(p) === (3, 2, 1)
-  @test isperm(p)
-  @test length(p) == 3
-  @test blocks(p) == ((3, 2), (), (1,))
-  @test blocklength(p) == 3
-  @test blocklengths(p) == (2, 0, 1)
-  @test blockfirsts(p) == (1, 3, 3)
-  @test blocklasts(p) == (2, 2, 3)
-  @test invperm(p) == blockedperm((3, 2), (), (1,))
-
-  # Split collection into `BlockedPermutation`.
-  p = blockedperm_indexin(("a", "b", "c", "d"), ("c", "a"), ("b", "d"))
-  @test p == blockedperm((3, 1), (2, 4))
-
-  # Singleton dimensions.
-  p = blockedperm((2, 3), 1)
-  @test p == blockedperm((2, 3), (1,))
-
-  # First dimensions are unspecified.
-  p = blockedperm(.., (4, 3))
-  @test p == blockedperm(1, 2, (4, 3))
-  # Specify length
-  p = blockedperm(.., (4, 3); length=Val(6))
-  @test p == blockedperm(1, 2, 5, 6, (4, 3))
-
-  # Last dimensions are unspecified.
-  p = blockedperm((4, 3), ..)
-  @test p == blockedperm((4, 3), 1, 2)
-  # Specify length
-  p = blockedperm((4, 3), ..; length=Val(6))
-  @test p == blockedperm((4, 3), 1, 2, 5, 6)
-
-  # Middle dimensions are unspecified.
-  p = blockedperm((4, 3), .., 1)
-  @test p == blockedperm((4, 3), 2, 1)
-  # Specify length
-  p = blockedperm((4, 3), .., 1; length=Val(6))
-  @test p == blockedperm((4, 3), 2, 5, 6, 1)
-
-  # No dimensions are unspecified.
-  p = blockedperm((3, 2), .., 1)
-  @test p == blockedperm((3, 2), 1)
-end
 @testset "TensorAlgebra" begin
   @testset "fusedims (eltype=$elt)" for elt in elts
     a = randn(elt, 2, 3, 4, 5)

test/test_blockedpermutation.jl

@@ -0,0 +1,65 @@
+using Test: @test, @test_broken, @testset
+
+using BlockArrays: blockfirsts, blocklasts, blocklength, blocklengths, blocks
+using Combinatorics: permutations
+
+using TensorAlgebra: BlockedTuple, blockedperm, blockedperm_indexin
+
+@testset "BlockedPermutation" begin
+  p = blockedperm((3, 4, 5), (2, 1))
+  @test Tuple(p) === (3, 4, 5, 2, 1)
+  @test isperm(p)
+  @test length(p) == 5
+  @test blocks(p) == ((3, 4, 5), (2, 1))
+  @test blocklength(p) == 2
+  @test blocklengths(p) == (3, 2)
+  @test blockfirsts(p) == (1, 4)
+  @test blocklasts(p) == (3, 5)
+  @test invperm(p) == blockedperm((5, 4, 1), (2, 3))
+
+  # Empty block.
+  p = blockedperm((3, 2), (), (1,))
+  @test Tuple(p) === (3, 2, 1)
+  @test isperm(p)
+  @test length(p) == 3
+  @test blocks(p) == ((3, 2), (), (1,))
+  @test blocklength(p) == 3
+  @test blocklengths(p) == (2, 0, 1)
+  @test blockfirsts(p) == (1, 3, 3)
+  @test blocklasts(p) == (2, 2, 3)
+  @test invperm(p) == blockedperm((3, 2), (), (1,))
+  @test BlockedTuple(p) == BlockedTuple{(2, 0, 1)}((3, 2, 1))
+
+  # Split collection into `BlockedPermutation`.
+  p = blockedperm_indexin(("a", "b", "c", "d"), ("c", "a"), ("b", "d"))
+  @test p == blockedperm((3, 1), (2, 4))
+
+  # Singleton dimensions.
+  p = blockedperm((2, 3), 1)
+  @test p == blockedperm((2, 3), (1,))
+
+  # First dimensions are unspecified.
+  p = blockedperm(.., (4, 3))
+  @test p == blockedperm(1, 2, (4, 3))
+  # Specify length
+  p = blockedperm(.., (4, 3); length=Val(6))
+  @test p == blockedperm(1, 2, 5, 6, (4, 3))
+
+  # Last dimensions are unspecified.
+  p = blockedperm((4, 3), ..)
+  @test p == blockedperm((4, 3), 1, 2)
+  # Specify length
+  p = blockedperm((4, 3), ..; length=Val(6))
+  @test p == blockedperm((4, 3), 1, 2, 5, 6)
+
+  # Middle dimensions are unspecified.
+  p = blockedperm((4, 3), .., 1)
+  @test p == blockedperm((4, 3), 2, 1)
+  # Specify length
+  p = blockedperm((4, 3), .., 1; length=Val(6))
+  @test p == blockedperm((4, 3), 2, 5, 6, 1)
+
+  # No dimensions are unspecified.
+  p = blockedperm((3, 2), .., 1)
+  @test p == blockedperm((3, 2), 1)
+end