[WIP] ProductedArrays

Project.toml

@@ -2,6 +2,9 @@ name = "MappedArrays"
 uuid = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
 version = "0.3.0"
 
+[deps]
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+
 [compat]
 FixedPointNumbers = "0.6.1, 0.7, 0.8"
 julia = "1"

src/MappedArrays.jl

@@ -4,6 +4,10 @@ using Base: @propagate_inbounds
 
 export AbstractMappedArray, MappedArray, ReadonlyMappedArray, mappedarray, of_eltype
 
+using Reexport
+include("ProductedArrays.jl")
+@reexport using .ProductedArrays
+
 abstract type AbstractMappedArray{T,N} <: AbstractArray{T,N} end
 abstract type AbstractMultiMappedArray{T,N} <: AbstractMappedArray{T,N} end
 

src/ProductedArrays.jl

@@ -0,0 +1,24 @@
+module ProductedArrays
+    export ProductedArray
+
+    struct ProductedArray{T, N, AAs<:Tuple{Vararg{AbstractArray}}} <: AbstractArray{T, N}
+        data::AAs
+    end
+    function ProductedArray(data...)
+        ProductedArray{typeof(map(first, data)), mapreduce(ndims, +, data), typeof(data)}(data)
+    end
+
+    @inline Base.size(A::ProductedArray) = mapreduce(size, (i,j)->(i...,j...), A.data)
+
+    Base.@propagate_inbounds function Base.getindex(A::ProductedArray{T, N}, inds::Vararg{Int, N}) where {T, N}
+        map((x, i)->x[i...], A.data, _split_indices(A, inds))
+    end
+
+    # TODO: this fails to inline and thus gives about 1.5ns overhead to getindex
+    @inline function _split_indices(A::ProductedArray{T, N}, inds::NTuple{N, Int}) where {T, N}
+        # TODO: this line is repeatedly computed
+        pos = (firstindex(A.data)-1, accumulate(+, map(ndims, A.data))...)
+
+        return ntuple(i->inds[pos[i]+1:pos[i+1]], length(pos)-1)
+    end
+end