Support tuples in blockedrange (#165)

* Support tuples in blockedrange

* simplify length for BlockedUnitRange

* Update blockaxis.jl

* Update test_blockindices.jl

* v0.15.1

Author: dlfivefifty

Project.toml

@@ -1,6 +1,6 @@
 name = "BlockArrays"
 uuid = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
-version = "0.15"
+version = "0.15.1"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/blockarray.jl

@@ -67,13 +67,13 @@ struct BlockArray{T, N, R <: AbstractArray{<:AbstractArray{T,N},N}, BS<:NTuple{N
 end
 
 # Auxilary outer constructors
-@inline _BlockArray(blocks::R, block_sizes::Vararg{AbstractVector{Int}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}} =
+@inline _BlockArray(blocks::R, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}} =
     _BlockArray(blocks, map(blockedrange, block_sizes))
 
 # support non-concrete eltypes in blocks
 _BlockArray(blocks::R, block_axes::BS) where {T, N, R<:AbstractArray{<:AbstractArray{V,N} where V,N}, BS<:NTuple{N,AbstractUnitRange{Int}}} =
     _BlockArray(convert(AbstractArray{AbstractArray{mapreduce(eltype,promote_type,blocks),N},N}, blocks), block_axes)
-_BlockArray(blocks::R, block_sizes::Vararg{AbstractVector{Int}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{V,N} where V,N}} =
+_BlockArray(blocks::R, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{V,N} where V,N}} =
     _BlockArray(convert(AbstractArray{AbstractArray{mapreduce(eltype,promote_type,blocks),N},N}, blocks), block_sizes...)
 
 const BlockMatrix{T, R <: AbstractMatrix{<:AbstractMatrix{T}}} = BlockArray{T, 2, R}
@@ -84,7 +84,7 @@ const BlockVecOrMat{T, R} = Union{BlockMatrix{T, R}, BlockVector{T, R}}
 # Constructors #
 ################
 
-@inline _BlockArray(::Type{R}, block_sizes::Vararg{AbstractVector{Int}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}} =
+@inline _BlockArray(::Type{R}, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}} =
     _BlockArray(R, map(blockedrange,block_sizes))
 
 function _BlockArray(::Type{R}, baxes::NTuple{N,AbstractUnitRange{Int}}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}}
@@ -93,7 +93,7 @@ function _BlockArray(::Type{R}, baxes::NTuple{N,AbstractUnitRange{Int}}) where {
     _BlockArray(blocks, baxes)
 end
 
-@inline undef_blocks_BlockArray(::Type{R}, block_sizes::Vararg{AbstractVector{Int}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}} =
+@inline undef_blocks_BlockArray(::Type{R}, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}} =
     _BlockArray(R, block_sizes...)
 
 """
@@ -109,16 +109,16 @@ julia> BlockArray(undef_blocks, Matrix{Float64}, [1,3], [2,2])
  #undef  #undef  │  #undef  #undef
 ```
 """
-@inline BlockArray(::UndefBlocksInitializer, ::Type{R}, block_sizes::Vararg{AbstractVector{Int}, N}) where {T, N, R<:AbstractArray{T,N}} =
+@inline BlockArray(::UndefBlocksInitializer, ::Type{R}, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T, N, R<:AbstractArray{T,N}} =
     undef_blocks_BlockArray(Array{R,N}, block_sizes...)
 
-@inline BlockArray{T}(::UndefBlocksInitializer, block_sizes::Vararg{AbstractVector{Int}, N}) where {T, N} =
+@inline BlockArray{T}(::UndefBlocksInitializer, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T, N} =
     BlockArray(undef_blocks, Array{T,N}, block_sizes...)
 
-@inline BlockArray{T,N}(::UndefBlocksInitializer, block_sizes::Vararg{AbstractVector{Int}, N}) where {T, N} =
+@inline BlockArray{T,N}(::UndefBlocksInitializer, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T, N} =
     BlockArray(undef_blocks, Array{T,N}, block_sizes...)
 
-@inline BlockArray{T,N,R}(::UndefBlocksInitializer, block_sizes::Vararg{AbstractVector{Int}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}} =
+@inline BlockArray{T,N,R}(::UndefBlocksInitializer, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}} =
     undef_blocks_BlockArray(R, block_sizes...)
 
 @generated function initialized_blocks_BlockArray(::Type{R}, baxes::NTuple{N,AbstractUnitRange{Int}}) where R<:AbstractArray{V,N} where {T,N,V<:AbstractArray{T,N}}
@@ -134,7 +134,7 @@ julia> BlockArray(undef_blocks, Matrix{Float64}, [1,3], [2,2])
 end
 
 
-initialized_blocks_BlockArray(::Type{R}, block_sizes::Vararg{AbstractVector{Int}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}} =
+initialized_blocks_BlockArray(::Type{R}, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}} =
     initialized_blocks_BlockArray(R, map(blockedrange,block_sizes))
 
 @inline BlockArray{T}(::UndefInitializer, baxes::NTuple{N,AbstractUnitRange{Int}}) where {T, N} =
@@ -149,20 +149,20 @@ initialized_blocks_BlockArray(::Type{R}, block_sizes::Vararg{AbstractVector{Int}
 @inline BlockArray{T,N,R,BS}(::UndefInitializer, baxes::BS) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}, BS<:NTuple{N,AbstractUnitRange{Int}}} =
     initialized_blocks_BlockArray(R, baxes)
 
-@inline BlockArray{T}(::UndefInitializer, block_sizes::Vararg{AbstractVector{Int}, N}) where {T, N} =
+@inline BlockArray{T}(::UndefInitializer, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T, N} =
     initialized_blocks_BlockArray(Array{Array{T,N},N}, block_sizes...)
 
-@inline BlockArray{T, N}(::UndefInitializer, block_sizes::Vararg{AbstractVector{Int}, N}) where {T, N} =
+@inline BlockArray{T, N}(::UndefInitializer, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T, N} =
     initialized_blocks_BlockArray(Array{Array{T,N},N}, block_sizes...)
 
-@inline BlockArray{T, N, R}(::UndefInitializer, block_sizes::Vararg{AbstractVector{Int}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}} =
+@inline BlockArray{T, N, R}(::UndefInitializer, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}} =
     initialized_blocks_BlockArray(R, block_sizes...)
 
 
 @inline BlockArray{T,N,R,BS}(::UndefInitializer, sizes::NTuple{N,Int}) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}, BS<:NTuple{N,AbstractUnitRange{Int}}} =
     BlockArray{T,N,R,BS}(undef, convert(BS, map(Base.OneTo, sizes)))
 
-function BlockArray{T}(arr::AbstractArray{V, N}, block_sizes::Vararg{AbstractVector{Int}, N}) where {T,V,N}
+function BlockArray{T}(arr::AbstractArray{V, N}, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T,V,N}
     for i in 1:N
         if sum(block_sizes[i]) != size(arr, i)
             throw(DimensionMismatch("block size for dimension $i: $(block_sizes[i]) does not sum to the array size: $(size(arr, i))"))
@@ -171,7 +171,7 @@ function BlockArray{T}(arr::AbstractArray{V, N}, block_sizes::Vararg{AbstractVec
     BlockArray{T}(arr, map(blockedrange,block_sizes))
 end
 
-BlockArray(arr::AbstractArray{T, N}, block_sizes::Vararg{AbstractVector{Int}, N}) where {T,N} =
+BlockArray(arr::AbstractArray{T, N}, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {T,N} =
     BlockArray{T}(arr, block_sizes...)
 
 @generated function BlockArray{T}(arr::AbstractArray{T, N}, baxes::NTuple{N,AbstractUnitRange{Int}}) where {T,N}
@@ -194,18 +194,18 @@ BlockArray(arr::AbstractArray{T, N}, baxes::NTuple{N,AbstractUnitRange{Int}}) wh
     BlockArray{T}(arr, baxes)
 
 BlockVector(blocks::AbstractVector, baxes::Tuple{AbstractUnitRange{Int}}) = BlockArray(blocks, baxes)
-BlockVector(blocks::AbstractVector, block_sizes::AbstractVector{Int}) = BlockArray(blocks, block_sizes)
+BlockVector(blocks::AbstractVector, block_sizes::AbstractVector{<:Integer}) = BlockArray(blocks, block_sizes)
 BlockMatrix(blocks::AbstractMatrix, baxes::NTuple{2,AbstractUnitRange{Int}}) = BlockArray(blocks, baxes)
-BlockMatrix(blocks::AbstractMatrix, block_sizes::Vararg{AbstractVector{Int},2}) = BlockArray(blocks, block_sizes...)
+BlockMatrix(blocks::AbstractMatrix, block_sizes::Vararg{AbstractVector{<:Integer},2}) = BlockArray(blocks, block_sizes...)
 
 BlockArray{T}(λ::UniformScaling, baxes::NTuple{2,AbstractUnitRange{Int}}) where T = BlockArray{T}(Matrix(λ, map(length,baxes)...), baxes)
-BlockArray{T}(λ::UniformScaling, block_sizes::Vararg{AbstractVector{Int}, 2}) where T = BlockArray{T}(λ, map(blockedrange,block_sizes))
-BlockArray(λ::UniformScaling{T}, block_sizes::Vararg{AbstractVector{Int}, 2}) where T = BlockArray{T}(λ, block_sizes...)
+BlockArray{T}(λ::UniformScaling, block_sizes::Vararg{AbstractVector{<:Integer}, 2}) where T = BlockArray{T}(λ, map(blockedrange,block_sizes))
+BlockArray(λ::UniformScaling{T}, block_sizes::Vararg{AbstractVector{<:Integer}, 2}) where T = BlockArray{T}(λ, block_sizes...)
 BlockArray(λ::UniformScaling{T}, baxes::NTuple{2,AbstractUnitRange{Int}}) where T = BlockArray{T}(λ, baxes)
 BlockMatrix(λ::UniformScaling, baxes::NTuple{2,AbstractUnitRange{Int}}) = BlockArray(λ, baxes)
-BlockMatrix(λ::UniformScaling, block_sizes::Vararg{AbstractVector{Int},2}) = BlockArray(λ, block_sizes...)
+BlockMatrix(λ::UniformScaling, block_sizes::Vararg{AbstractVector{<:Integer},2}) = BlockArray(λ, block_sizes...)
 BlockMatrix{T}(λ::UniformScaling, baxes::NTuple{2,AbstractUnitRange{Int}}) where T = BlockArray{T}(λ, baxes)
-BlockMatrix{T}(λ::UniformScaling, block_sizes::Vararg{AbstractVector{Int},2}) where T = BlockArray{T}(λ, block_sizes...)
+BlockMatrix{T}(λ::UniformScaling, block_sizes::Vararg{AbstractVector{<:Integer},2}) where T = BlockArray{T}(λ, block_sizes...)
 
 """
     mortar(blocks::AbstractArray)
@@ -244,7 +244,7 @@ true
 mortar(blocks::AbstractArray{R, N}, baxes::NTuple{N,AbstractUnitRange{Int}}) where {R, N} =
     _BlockArray(blocks, baxes)
 
-mortar(blocks::AbstractArray{R, N}, block_sizes::Vararg{AbstractVector{Int}, N}) where {R, N} =
+mortar(blocks::AbstractArray{R, N}, block_sizes::Vararg{AbstractVector{<:Integer}, N}) where {R, N} =
     _BlockArray(blocks, block_sizes...)
 
 mortar(blocks::AbstractArray) = mortar(blocks, sizes_from_blocks(blocks)...)

src/blockaxis.jl

@@ -53,11 +53,16 @@ const DefaultBlockAxis = BlockedUnitRange{Vector{Int}}
 
 BlockedUnitRange(::BlockedUnitRange) = throw(ArgumentError("Forbidden due to ambiguity"))
 _blocklengths2blocklasts(blocks) = cumsum(blocks) # extra level to allow changing default cumsum behaviour
-@inline blockedrange(blocks::AbstractVector{Int}) = _BlockedUnitRange(_blocklengths2blocklasts(blocks))
+@inline blockedrange(blocks::Union{Tuple,AbstractVector}) = _BlockedUnitRange(_blocklengths2blocklasts(blocks))
 
-@inline blockfirsts(a::BlockedUnitRange) = [a.first; @view(a.lasts[1:end-1]) .+ 1]
+@inline blockfirsts(a::BlockedUnitRange) = [a.first; @views(a.lasts[1:end-1]) .+ 1]
 @inline blocklasts(a::BlockedUnitRange) = a.lasts
-@inline blocklengths(a::BlockedUnitRange) = [first(a.lasts)-a.first+1; diff(a.lasts)]
+
+_diff(a::AbstractVector) = diff(a)
+_diff(a::Tuple) = diff(collect(a))
+@inline blocklengths(a::BlockedUnitRange) = [first(a.lasts)-a.first+1; _diff(a.lasts)]
+
+length(a::BlockedUnitRange) = isempty(a.lasts) ? 0 : Integer(last(a.lasts))
 
 """
    blockisequal(a::AbstractUnitRange{Int}, b::AbstractUnitRange{Int})
@@ -105,7 +110,9 @@ julia> blockaxes(A)[1]
  Block(2)
 ```
 """
-blockaxes(b::BlockedUnitRange) = (Block.(axes(b.lasts,1)),)
+blockaxes(b::BlockedUnitRange) = _blockaxes(b.lasts)
+_blockaxes(b::AbstractVector) = (Block.(axes(b,1)),)
+_blockaxes(b::Tuple) = (Block.(Base.OneTo(length(b))),)
 blockaxes(b) = blockaxes.(axes(b), 1)
 
 """
@@ -186,10 +193,7 @@ blocksizes(A,i) = blocklengths(axes(A,i))
 axes(b::BlockedUnitRange) = (_BlockedUnitRange(blocklasts(b) .- (first(b)-1)),)
 unsafe_indices(b::BlockedUnitRange) = axes(b)
 first(b::BlockedUnitRange) = b.first
-_last(b::BlockedUnitRange, _) = isempty(blocklasts(b)) ? first(b)-1 : last(blocklasts(b))
-last(b::BlockedUnitRange) = _last(b, axes(blocklasts(b),1))
-_length(b::BlockedUnitRange, _) = Base.invoke(length, Tuple{AbstractUnitRange{Int}}, b)
-length(b::BlockedUnitRange) = _length(b, axes(blocklasts(b),1))
+last(b::BlockedUnitRange) = isempty(blocklasts(b)) ? first(b)-1 : last(blocklasts(b))
 
 function getindex(b::BlockedUnitRange, K::Block{1})
     k = Integer(K)
@@ -223,9 +227,16 @@ function getindex(b::BlockedUnitRange, KR::BlockRange{1,Tuple{Base.OneTo{Int}}})
     _BlockedUnitRange(first(b),cs[Base.OneTo(j)])
 end
 
+_searchsortedfirst(a::AbstractVector, k) = searchsortedfirst(a, k)
+function _searchsortedfirst(a::Tuple, k)
+    k ≤ first(a) && return 1
+    1+_searchsortedfirst(tail(a), k)
+end
+_searchsortedfirst(a::Tuple{}, k) = 1
+
 function findblock(b::BlockedUnitRange, k::Integer)
     @boundscheck k in b || throw(BoundsError(b,k))
-    Block(searchsortedfirst(blocklasts(b), k))
+    Block(_searchsortedfirst(blocklasts(b), k))
 end
 
 Base.dataids(b::BlockedUnitRange) = Base.dataids(blocklasts(b))
@@ -320,3 +331,9 @@ function blocklengths(a::BlockedUnitRange{<:AbstractRange})
     @assert first(a.lasts)-a.first+1 == st
     Fill(st,length(a.lasts))
 end
+
+
+# TODO: Remove
+
+function _last end
+function _length end

test/test_blockindices.jl

@@ -301,6 +301,16 @@ end
         @test blockisequal(blockedrange(SVector(1,2,3)), blockedrange([1,2,3]))
         @test @allocated(blockedrange(SVector(1,2,3))) == 0
     end
+
+    @testset "Tuples" begin
+        # we support Tuples in addition to SVectors for InfiniteArrays.jl, which has
+        # infinite block sizes
+        s = blockedrange((5,big(100_000_000)^2))
+        @test blocklengths(s) == [5,big(100_000_000)^2]
+        @test blockaxes(s) == (Block.(1:2),)
+        @test findblock(s,3) == Block(1)
+        @test findblock(s,big(100_000_000)) == Block(2)
+    end
 end
 
 @testset "BlockSlice" begin