layered array

layered_array.jl

using Revise
using ArgCheck

struct LayeredArray{T, N, L} <: AbstractArray{T,N}
    layers::L
    function LayeredArray(layers)
        @argcheck !isempty(layers)
        @argcheck first(layers) isa AbstractArray
        l = first(layers)
        L = typeof(layers)
        N = ndims(l)
        T = eltype(l)
        for layer in layers
            @argcheck layer isa AbstractArray{T, N}
        end
        ret = new{T,N,L}(layers)
        ret
    end
end

fallbacklayer(l::LayeredArray) = last(l.layers)
Base.size(l::LayeredArray) = size(fallbacklayer(l))
Base.axes(l::LayeredArray) = Base.axes(fallbacklayer(l))
Base.IndexStyle(l::Type{<: LayeredArray}) = Base.IndexCartesian()
function Base.getindex(o::LayeredArray, ci::CartesianIndex)
    l = get_containing_layer(o, ci)
    l[ci]
end
function get_containing_layer(o::LayeredArray, ci::CartesianIndex)
    for layer in o.layers
        if ci in CartesianIndices(layer)
            return layer
        end
    end
    err = BoundsError(o, ci)
    throw(err)
end
function Base.getindex(l::LayeredArray, inds...)
    ci = CartesianIndices(l)[inds...]
    l[ci]
end


using Test
l = LayeredArray([[10,2], [1,6,3]])
@inferred l == [10, 2, 3]