diff -r d881275c6564 -r 6dfa8001eed2 src/ZipArrays.jl --- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/ZipArrays.jl Wed Dec 15 01:09:09 2021 +0200 @@ -0,0 +1,84 @@ +####################################### +# 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 +