Add append!/push!/pushfirst!/pop!/popfirst! for BlockVector

Author: tkf

docs/src/lib/public.md

@@ -50,6 +50,11 @@ BlockArray
 undef_blocks
 UndefBlocksInitializer
 mortar
+Base.append!
+Base.push!
+Base.pushfirst!
+Base.pop!
+Base.popfirst!
 ```
 
 

src/BlockArrays.jl

@@ -48,5 +48,7 @@ include("blockarrayinterface.jl")
 include("blockbroadcast.jl")
 include("blocklinalg.jl")
 include("blockproduct.jl")
+include("blockreduce.jl")
+include("blockdeque.jl")
 
 end # module

src/blockdeque.jl

@@ -0,0 +1,152 @@
+"""
+    append!(dest::BlockVector, sources...; alias = false)
+
+Append items from `sources` to the last block of `dest`.  If `alias = true`,
+append the blocks in `sources` as-is if they are compatible with the internal
+block array type of `dest`.  Importantly, this means that mutating `sources`
+afterwards alters the items in `dest` and it may even break the invariance
+of `dest` if the length of `sources` are changed.  The elements may be copied
+even if `alias = true` if the corresponding implementation does not exist.
+
+The blocks in `dest` must not alias with `sources` or components of them.
+For example, the result of `append!(x, x)` is undefined.
+"""
+Base.append!(dest::BlockVector, sources...; alias::Bool = false) =
+    foldl((d, s) -> append!(d, s; alias = alias), sources; init = dest)
+
+function Base.append!(dest::BlockVector, src; alias::Bool = false)
+    if _blocktype(dest) === _blocktype(src) && alias
+        return append_nocopy!(dest, src)
+    else
+        return append_copy!(dest, src)
+    end
+end
+
+_blocktype(::Any) = nothing
+_blocktype(::T) where {T<:AbstractArray} = T
+_blocktype(::BlockArray{<:Any,<:Any,<:AbstractArray{T}}) where {T<:AbstractArray} = T
+_blocktype(::PseudoBlockArray{<:Any,<:Any,T}) where {T<:AbstractArray} = T
+
+function append_nocopy!(dest::BlockVector{<:Any,T}, src::BlockVector{<:Any,T}) where {T}
+    isempty(src) && return dest
+    append!(dest.blocks, src.blocks)
+    offset = last(dest.axes[1]) + 1 - src.axes[1].first
+    append!(dest.axes[1].lasts, (n + offset for n in src.axes[1].lasts))
+    return dest
+end
+
+append_nocopy!(
+    dest::BlockVector{<:Any,<:AbstractArray{T}},
+    src::PseudoBlockVector{<:Any,T},
+) where {T} = append_nocopy!(dest, src.blocks)
+
+function append_nocopy!(dest::BlockVector{<:Any,<:AbstractArray{T}}, src::T) where {T}
+    isempty(src) && return dest
+    push!(dest.blocks, src)
+    push!(dest.axes[1].lasts, last(dest.axes[1]) + length(src))
+    return dest
+end
+
+append_copy!(dest::BlockVector, src) = _append_copy!(dest, Base.IteratorSize(src), src)
+
+function _append_copy!(dest::BlockVector, ::Union{Base.HasShape,Base.HasLength}, src)
+    block = dest.blocks[end]
+    li = lastindex(block)
+    resize!(block, length(block) + length(src))
+    # Equivalent to `i = li; for x in src; ...; end` but (maybe) faster:
+    foldl(src, init = li) do i, x
+        Base.@_inline_meta
+        i += 1
+        @inbounds block[i] = x
+        return i
+    end
+    da, = dest.axes
+    da.lasts[end] += length(src)
+    return dest
+end
+
+function _append_copy!(dest::BlockVector, ::Base.SizeUnknown, src)
+    block = dest.blocks[end]
+    # Equivalent to `n = 0; for x in src; ...; end` but (maybe) faster:
+    n = foldl(src, init = 0) do n, x
+        push!(block, x)
+        return n + 1
+    end
+    da, = dest.axes
+    da.lasts[end] += n
+    return dest
+end
+
+# remove empty blocks at the end
+function _squash_lasts!(A::BlockVector)
+    while !isempty(A.blocks) && isempty(A.blocks[end])
+        pop!(A.blocks)
+        pop!(A.axes[1].lasts)
+    end
+end
+
+# remove empty blocks at the beginning
+function _squash_firsts!(A::BlockVector)
+    while !isempty(A.blocks) && isempty(A.blocks[1])
+        popfirst!(A.blocks)
+        popfirst!(A.axes[1].lasts)
+    end
+end
+
+"""
+    pop!(A::BlockVector)
+
+Pop the last element from the last non-empty block.  Remove all empty
+blocks at the end.
+"""
+function Base.pop!(A::BlockVector)
+    isempty(A) && throw(Argument("array must be nonempty"))
+    _squash_lasts!(A)
+    x = pop!(A.blocks[end])
+    lasts = A.axes[1].lasts
+    if isempty(A.blocks[end])
+        pop!(A.blocks)
+        pop!(lasts)
+    else
+        lasts[end] -= 1
+    end
+    return x
+end
+
+"""
+    popfirst!(A::BlockVector)
+
+Pop the first element from the first non-empty block.  Remove all empty
+blocks at the beginning.
+"""
+function Base.popfirst!(A::BlockVector)
+    isempty(A) && throw(Argument("array must be nonempty"))
+    _squash_firsts!(A)
+    x = popfirst!(A.blocks[1])
+    ax, = A.axes
+    if isempty(A.blocks[1])
+        popfirst!(A.blocks)
+        popfirst!(ax.lasts)
+    else
+        ax.lasts[1] -= 1
+    end
+    return x
+end
+
+"""
+    push!(dest::BlockVector, items...)
+
+Push items to the end of the last block.
+"""
+Base.push!(dest::BlockVector, items...) = append!(dest, items)
+
+"""
+    pushfirst!(A::BlockVector, items...)
+
+Push items to the beginning of the first block.
+"""
+function Base.pushfirst!(A::BlockVector, items...)
+    pushfirst!(A.blocks[1], items...)
+    A.axes[1].lasts .+= length(items)
+    return A
+end

src/blockreduce.jl

@@ -0,0 +1,18 @@
+# * Those seemingly no-op `where {F, OP}` for forcing specialization.
+#   See: https://github.com/JuliaLang/julia/pull/33917
+# * Per-block reduction strategy is correct only for vectors.
+
+# Let mapping transducer in `Base` compose an efficient nested loop:
+Base.mapfoldl(f::F, op::OP, B::BlockVector; kw...) where {F, OP} =
+    foldl(op, (f(x) for block in B.blocks for x in block); kw...)
+
+Base.mapreduce(f::F, op::OP, B::BlockVector; kw...) where {F, OP} =
+    mapfoldl(op, B.blocks; kw...) do block
+        mapreduce(f, op, block; kw...)
+    end
+
+Base.mapfoldl(f::F, op::OP, B::PseudoBlockArray; kw...) where {F, OP} =
+    mapfoldl(f, op, B.blocks; kw...)
+
+Base.mapreduce(f::F, op::OP, B::PseudoBlockArray; kw...) where {F, OP} =
+    mapreduce(f, op, B.blocks; kw...)

test/runtests.jl

@@ -9,4 +9,6 @@ using BlockArrays, LinearAlgebra, Test
     include("test_blockbroadcast.jl")
     include("test_blocklinalg.jl")
     include("test_blockproduct.jl")
+    include("test_blockreduce.jl")
+    include("test_blockdeque.jl")
 end

test/test_blockdeque.jl

@@ -0,0 +1,80 @@
+using BlockArrays, Test
+
+@testset "append!(::BlockVector, vector)" begin
+    @testset for alias in [false, true],
+        compatible in [false, true],
+        srctype in [:BlockVector, :PseudoBlockVector, :Vector]
+
+        dest = mortar([[1, 2, 3], [4, 5]])
+
+        # Create `src` array:
+        if compatible
+            T = Int
+        else
+            T = Float64
+        end
+        if srctype === :BlockVector
+            src = mortar([T[6, 7], T[8, 9]])
+        elseif srctype === :PseudoBlockVector
+            src = PseudoBlockVector(T[6:9;], [3, 1])
+        elseif srctype === :Vector
+            src = T[6:9;]
+        else
+            error("Unknown srctype = $srctype")
+        end
+
+        @test append!(dest, src; alias = alias) === dest
+        @test dest == 1:9
+        src[1] = 666
+        if alias && compatible
+            @test dest[6] == 666
+        else
+            @test dest[6] == 6
+        end
+    end
+end
+
+@testset "append!(::BlockVector, iterator)" begin
+    @testset "$label" for (label, itr) in [
+        "with length" => (x + 0 for x in 6:9),
+        "no length" => (x for x in 6:9 if x > 0),
+    ]
+        dest = mortar([[1, 2, 3], [4, 5]])
+        @test append!(dest, src) === dest
+        @test dest == 1:9
+    end
+end
+
+@testset "push!" begin
+    A = mortar([[1, 2, 3], [4, 5]])
+    push!(A, 6)
+    push!(A, 7, 8, 9)
+    @test A == 1:9
+end
+
+@testset "pushfirst!" begin
+    A = mortar([[1, 2, 3], [4, 5]])
+    pushfirst!(A, 0)
+    pushfirst!(A, -3, -2, -1)
+    @test A == -3:5
+end
+
+@testset "pop!" begin
+    A = mortar([[1, 2, 3], [4, 5]])
+    B = []
+    while !isempty(A)
+        push!(B, pop!(A))
+    end
+    @test A == []
+    @test B == 5:-1:1
+end
+
+@testset "popfirst!" begin
+    A = mortar([[1, 2, 3], [4, 5]])
+    B = []
+    while !isempty(A)
+        push!(B, popfirst!(A))
+    end
+    @test A == []
+    @test B == 1:5
+end

test/test_blockreduce.jl

@@ -0,0 +1,17 @@
+using BlockArrays, Test
+
+@testset "foldl" begin
+    x = mortar([rand(3), rand(2)])
+    @test foldl(push!, x; init = []) == collect(x)
+
+    x = PseudoBlockVector(rand(3), [1, 2])
+    @test foldl(push!, x; init = []) == collect(x)
+end
+
+@testset "reduce" begin
+    x = mortar([rand(Int, 3), rand(Int, 2)])
+    @test reduce(+, x) == sum(collect(x))
+
+    x = PseudoBlockVector(rand(Int, 3), [1, 2])
+    @test reduce(+, x) == sum(collect(x))
+end