Mercurial > repos > AlgTools / file revision

#######################################
# Simultaneous view to multiple arrays
#######################################

"""
Implements a type for zipped arrays; see `ZipArray`, `AbstractZipArray`,
`ZipVector`, and `AbstractZipVector`.
"""
module ZipArrays

export ZipArray,
       ZipVector,
       AbstractZipArray,
       AbstractZipVector

"""
`ZipArray{S,T,N,A,B}`

`N`-dimensional array of element type `Tuple{S,T}`, the component of the tuple stored in separate
arrays of type `A <: AbstractArray{S,N}` and `B <: AbstractArray{T,N}`.
"""
struct ZipArray{S,T,N,A <: AbstractArray{S,N}, B <: AbstractArray{T,N}} <: AbstractArray{Tuple{S,T}, N}
    a :: A
    b :: B

    function ZipArray(a :: A, b :: B) where {S,T,N,A <: AbstractArray{S,N}, B <: AbstractArray{T,N}}
        @assert(IndexStyle(A)==IndexStyle(B))
        @assert(size(a)==size(b))
        return new{S,T,N,A,B}(a, b)
    end

end

"""
`AbstractZipArray{S,T,N}`

A way to refer to a `ZipArray` without specifying the container types.
"""
AbstractZipArray{S,T,N} = ZipArray{S,T,N,A,B} where {A <: AbstractArray{S,N}, B <: AbstractArray{T,N}}

"""
`ZipVector{S,T,A,B}`

One-dimensional `ZipArray`.
"""
ZipVector{S,T,A,B} = ZipArray{S,T,1,A,B}

"""
`AbstractZipVector{S,T}`

One-dimensional `AbstractZipArray`.
"""
AbstractZipVector{S,T} = AbstractZipArray{S,T,1}


function Base.IndexStyle( :: Type{<:ZipArray{S,T,N,A,B}}) where {S,T,N,A,B}
    return IndexStyle(A)
end

@inline function Base.size(z :: AbstractZipArray{S,T,N}) where {S,T,N}
    return size(z.a)
end

@inline function Base.getindex(z :: AbstractZipArray{S,T,N}, i) where {S,T,N}
    return (getindex(z.a, i), getindex(z.b, i))
end

@inline function Base.setindex!(z :: AbstractZipArray{S,T,N}, (s, t) :: Tuple{S,T}, i) where {S,T,N}
    return (setindex!(z.a, s, i), setindex!(z.b, t, i))
end

@inline function Base.pop!(z :: AbstractZipVector{S,T}) where {S,T}
    s = pop!(z.a)
    t = pop!(z.b)
    return (s, t)
end

@inline function Base.push!(z :: AbstractZipVector{S,T}, items :: Vararg{Tuple{S,T},M}) where {S,T,M}
    push!(z.a, s for (s, _) ∈ items)
    push!(z.b, t for (_, t) ∈ items)
end

end # module