Fix zero dimensional BlockArray and BlockedArray (#410)

* add broken tests

* fix zerodim constructor and getindex

* fix blockcheckbounds

* fix tests

* fix views

* fix setindex

* fix view

* fix Array

* fix BlockedArray view

* test linalg

* fix linear algebra

* test for product

* test ^

* same behavior as Array for .* and .^

* fix reshape

* test view(::ReshapedArray)

* fix getindex(::Block{0})

* revert unneeded changes

Author: ogauthe

src/abstractblockarray.jl

@@ -101,6 +101,7 @@ end
 
 blockcheckbounds(A::AbstractArray{T, N}, i::Block{N}) where {T,N} = blockcheckbounds(A, i.n...)
 blockcheckbounds(A::AbstractArray{T, N}, i::Vararg{Block{1},N}) where {T,N} = blockcheckbounds(A, Int.(i)...)
+blockcheckbounds(::AbstractArray{T, 0}) where {T} = true
 blockcheckbounds(A::AbstractVector{T}, i::Block{1}) where {T} = blockcheckbounds(A, Int(i))
 
 """
@@ -186,6 +187,7 @@ viewblock(block_arr, block) = Base.invoke(view, Tuple{AbstractArray, Any}, block
     blkind = BlockRange(blocksize(block_arr))[Int(block)]
     view(block_arr, blkind)
 end
+@inline view(zerodim::AbstractBlockArray{<:Any,0}) = view(zerodim.blocks[])
 @inline view(block_arr::AbstractBlockVector, block::Block{1}) = viewblock(block_arr, block)
 @propagate_inbounds view(block_arr::AbstractBlockArray, block::Block{1}...) = view(block_arr, Block(block))
 

src/blockarray.jl

@@ -265,6 +265,12 @@ function BlockArray{T}(arr::AbstractArray{T, N}, baxes::Tuple{Vararg{AbstractUni
     return _BlockArray(blocks, baxes)
 end
 
+function BlockArray{T}(arr::AbstractArray{T, 0}, ::Tuple{}) where T
+    blocks = Array{Array{T, 0},0}(undef)
+    fill!(blocks, arr)
+    return _BlockArray(blocks)
+end
+
 BlockArray{T}(arr::AbstractArray{<:Any, N}, baxes::Tuple{Vararg{AbstractUnitRange{<:Integer},N}}) where {T,N} =
     BlockArray{T}(convert(AbstractArray{T, N}, arr), baxes)
 
@@ -463,12 +469,21 @@ const OffsetAxis = Union{Integer, UnitRange, Base.OneTo, Base.IdentityUnitRange}
     return v
 end
 
+@inline function getindex(block_arr::BlockArray{T, 0}) where T
+    return blocks(block_arr)[][]
+end
+
 @inline function setindex!(block_arr::BlockArray{T, N}, v, i::Vararg{Integer, N}) where {T,N}
     @boundscheck checkbounds(block_arr, i...)
     @inbounds block_arr[findblockindex.(axes(block_arr), i)...] = v
     return block_arr
 end
 
+@inline function setindex!(block_arr::BlockArray{<:Any, 0}, v)
+    blocks(block_arr)[][] = v
+end
+
+
 ############
 # Indexing #
 ############
@@ -527,6 +542,11 @@ end
 ########
 # Misc #
 ########
+function Base.Array(zerodim::BlockArray{T, 0}) where {T}
+    arr = Array{T}(undef)
+    arr[] = zerodim[]
+    return arr
+end
 
 function Base.Array(block_array::BlockArray{T, N, R}) where {T,N,R}
     arr = Array{eltype(T)}(undef, size(block_array))
@@ -547,6 +567,8 @@ end
 # Temporary work around
 Base.reshape(block_array::BlockArray, axes::Tuple{Vararg{AbstractUnitRange{<:Integer},N}}) where N =
     reshape(BlockedArray(block_array), axes)
+Base.reshape(block_array::BlockArray, ::Tuple{}) =
+    reshape(BlockedArray(block_array), ())  # zerodim
 Base.reshape(block_array::BlockArray, dims::Tuple{Int,Vararg{Int}}) =
     reshape(BlockedArray(block_array), dims)
 Base.reshape(block_array::BlockArray, axes::Tuple{Union{Integer,Base.OneTo}, Vararg{Union{Integer,Base.OneTo}}}) =

src/blockedarray.jl

@@ -99,6 +99,7 @@ const BlockedVector{T} = BlockedArray{T, 1}
 const BlockedVecOrMat{T} = Union{BlockedMatrix{T}, BlockedVector{T}}
 
 # Auxiliary outer constructors
+BlockedArray(x::Number, ::Tuple{}) = x  # zero dimensional
 @inline BlockedArray(blocks::R, baxes::BS) where {T,N,R<:AbstractArray{T,N},BS<:Tuple{Vararg{AbstractUnitRange{<:Integer},N}}} =
     BlockedArray{T, N, R,BS}(blocks, baxes)
 
@@ -240,6 +241,7 @@ end
 ############
 # Indexing #
 ############
+@inline view(block_arr::BlockedArray{<:Any, 0}) = view(block_arr.blocks)
 
 @inline function viewblock(block_arr::BlockedArray, block)
     range = getindex.(axes(block_arr), Block.(block.n))
@@ -300,6 +302,8 @@ Base.reshape(parent::BlockedArray, shp::Tuple{Union{Int,Base.OneTo}, Vararg{Unio
     reshape(parent, Base.to_shape(shp))
 Base.reshape(parent::BlockedArray, dims::Tuple{Int,Vararg{Int}}) =
     Base._reshape(parent, dims)
+Base.reshape(block_array::BlockedArray, ::Tuple{}) =
+    _blocked_reshape(block_array, ())  # zero dim
 
 """
     resize!(a::BlockedVector, N::Block) -> BlockedVector

src/blockindices.jl

@@ -202,6 +202,7 @@ BlockIndexRange(block::Block{N}, inds::Vararg{AbstractUnitRange{<:Integer},N}) w
 
 block(R::BlockIndexRange) = R.block
 
+getindex(::Block{0}) = Block()
 getindex(B::Block{N}, inds::Vararg{Integer,N}) where N = BlockIndex(B,inds)
 getindex(B::Block{N}, inds::Vararg{AbstractUnitRange{<:Integer},N}) where N = BlockIndexRange(B,inds)
 getindex(B::Block{1}, inds::Colon) = B

src/views.jl

@@ -88,9 +88,13 @@ end
     return reshape(view(A.parent, I[1:M]...), Val(N))
 end
 
-@propagate_inbounds  function Base.unsafe_view(
-        A::Array{<:Any, N},
-        I::Vararg{BlockSlice{<:BlockIndexRange{1}}, N}) where {N}
+@propagate_inbounds function Base.unsafe_view(
+    A::Array,
+    I1::BlockSlice{<:BlockIndexRange{1}},
+    Is::Vararg{BlockSlice{<:BlockIndexRange{1}}},
+)
+    I = (I1, Is...)
+    @assert ndims(A) == length(I)
     return view(A, map(x -> x.indices, I)...)
 end
 

test/test_blockarrays.jl

@@ -318,6 +318,61 @@ end
             @test all(isone, M)
         end
 
+        @testset "zero dim" begin
+            zerodim = ones()
+            @test view(zerodim) isa AbstractArray{Float64, 0}  #  check no type-piracy
+
+            ret = BlockArray{Float64}(undef)
+            @test ret isa BlockArray{Float64, 0}
+            fill!(ret, 0)
+            @test size(ret) == ()
+            @test all(iszero, ret)
+            @test ret[Block()] == zeros()
+            @test ret[Block()[]] == zeros()
+            @test ret[] == 0
+            @test view(ret, Block()) == zeros()
+            @test Array(ret) == zeros()
+            ret[] = 1
+            @test ret[] == 1
+            @test view(ret) == ones()
+            view(ret)[] = 0
+            @test ret[] == 0
+
+            ret = BlockArrays.BlockArray(zeros())
+            @test ret isa BlockArray{Float64, 0}
+            @test size(ret) == ()
+            @test all(iszero, ret)
+            @test ret[Block()] == zeros()
+
+            ret = BlockArrays.BlockArray(zeros(1,1))
+            @test reshape(ret, ()) isa AbstractBlockArray{Float64, 0}  # may be BlockedArray
+            @test size(reshape(ret, ())) == ()
+
+            ret = BlockedArray{Float64}(undef)
+            @test ret isa BlockedArray{Float64, 0}
+            fill!(ret, 0)
+            @test size(ret) == ()
+            @test all(iszero, ret)
+            @test ret[] == 0
+            @test ret[Block()] == zeros()
+            @test ret[Block()[]] == zeros()
+            @test Array(ret) == zeros()
+            ret[] = 1
+            @test ret[] == 1
+            @test view(ret) == ones()
+            view(ret)[] = 0
+            @test ret[] == 0
+
+            ret = BlockedArray(zeros())
+            @test size(ret) == ()
+            @test all(iszero, ret)
+            @test ret[Block()] == zeros()
+
+            ret = BlockArrays.BlockedArray(zeros(1,1))
+            @test reshape(ret, ()) isa BlockedArray{Float64, 0}
+            @test size(reshape(ret, ())) == ()
+        end
+
         @testset "BlockVector" begin
             a_data = [1,2,3]
             a = BlockVector(a_data,[1,2])

test/test_blocklinalg.jl

@@ -7,6 +7,34 @@ import ArrayLayouts: DenseRowMajor, ColumnMajor, StridedLayout
 bview(a, b) = Base.invoke(view, Tuple{AbstractArray,Any}, a, b)
 
 @testset "Linear Algebra" begin
+    @testset "zerodim" begin
+    a = BlockArray{Float64}(2*ones())
+    @test 2a isa BlockArray{Float64,0}
+    @test (2a)[] == 4
+    @test a + a isa BlockArray{Float64,0}
+    @test a + a == 2a
+    @test norm(a) == 2
+
+    # same behavior as Array
+    @test a .* a isa Float64
+    @test a .* a == 4
+    @test a .^ 2 isa Float64
+    @test a .^ 2 == 4
+
+    a = BlockedArray{Float64}(2*ones())
+    @test 2a isa BlockedArray{Float64,0}
+    @test (2a)[] == 4
+    @test a + a isa BlockedArray{Float64,0}
+    @test a + a == 2a
+    @test norm(a) == 2
+
+    # same behavior as Array
+    @test a .* a isa Float64
+    @test a .* a == 4
+    @test a .^ 2 isa Float64
+    @test a .^ 2 == 4
+    end
+
     @testset "BlockArray scalar * matrix" begin
         A = BlockArray{Float64}(randn(6,6), fill(2,3), 1:3)
         @test 2A == A*2 == 2Matrix(A)

test/test_blockviews.jl

@@ -21,6 +21,9 @@ bview(a, b) = Base.invoke(view, Tuple{AbstractArray,Any}, a, b)
         @test collect(b) == [2,3]
         @test b[1] == 2
         @test b[1:2] == 2:3
+
+        rba = reshape(BlockedArray(collect(1:4),[2,2]), (2,2))
+        @test view(rba, Block(1,1)[1:1,1:1]) == ones(1,1)
     end
 
     @testset "block view" begin