Mercurial > repos > AlgTools / file diff

--- /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
+