Merge branch 'main' into stack_as_abstractarray

ext/DimensionalDataDiskArraysExt.jl

@@ -12,7 +12,7 @@ import DiskArrays
 # cache was only introduced in DiskArrays v0.4, so
 # we lock out the method definition if the method does
 # not exist.
-@static if :cache in names(DiskArrays)
+@static if isdefined(DiskArrays, :cache)
     DiskArrays.cache(x::Union{AbstractDimStack,AbstractDimArray}; kw...) = 
         modify(A -> DiskArrays.cache(A; kw...), x)
 end

src/stack/indexing.jl

@@ -86,11 +86,12 @@ for f in (:getindex, :view, :dotview)
                 I = length(layerdims) > 0 ? layerdims : map(_ -> :, size(A))
                 Base.$f(A, I...)
             end
-            newlayers = map(f, values(s))
+            newlayers = unrolled_map(f, values(s))
             # Decide to rewrap as an AbstractDimStack, or return a scalar
-            if any(map(v -> v isa AbstractDimArray, newlayers))
+            if _any_dimarray(newlayers)
+                # TODO rethink this for many-layered stacks
                 # Some scalars, re-wrap them as zero dimensional arrays
-                non_scalar_layers = map(values(s), newlayers) do l, nl
+                non_scalar_layers = unrolled_map(values(s), newlayers) do l, nl
                     nl isa AbstractDimArray ? nl : rebuild(l, fill(nl), ())
                 end
                 rebuildsliced(Base.$f, s, NamedTuple{K}(non_scalar_layers), (dims2indices(dims(s), D)))
@@ -102,6 +103,11 @@ for f in (:getindex, :view, :dotview)
     end
 end
 
+@generated function _any_dimarray(v::Union{NamedTuple,Tuple})
+    any(T -> T <: AbstractDimArray, v.types)
+end
+
+
 
 #### setindex ####
 @propagate_inbounds Base.setindex!(s::AbstractDimStack, xs, I...; kw...) =

src/stack/stack.jl

@@ -60,7 +60,7 @@ end
 @assume_effects :foldable DD.layers(s::AbstractDimStack, k::Symbol) = s[k]
 
 @assume_effects :foldable data_eltype(nt::NamedTuple{K}) where K =
-    NamedTuple{K,Tuple{map(eltype, Tuple(nt))...}}
+    NamedTuple{K,Tuple{unrolled_map(eltype, Tuple(nt))...}}
 stacktype(s, data, dims, layerdims::NamedTuple{K}) where K =
     basetypeof(s){K,data_eltype(data),length(dims)}
 
@@ -89,10 +89,15 @@ function rebuild(s::AbstractDimStack;
     return T(data, dims, refdims, layerdims, metadata, layermetadata)
 end
 
-function rebuildsliced(f::Function, s::AbstractDimStack, layers, I)
-    layerdims = map(basedims, layers)
+function rebuildsliced(f::Function, s::AbstractDimStack, layers::NamedTuple, I)
+    layerdims = unrolled_map(basedims, layers)
     dims, refdims = slicedims(f, s, I)
-    return rebuild(s; data=map(parent, layers), dims, refdims, layerdims)
+    return rebuild(s; data=unrolled_map(parent, layers), dims, refdims, layerdims)
+end
+function rebuildsliced(f::Function, s::AbstractDimStack{K}, layers::Tuple, I) where K
+    layerdims = NamedTuple{K}(unrolled_map(basedims, layers))
+    dims, refdims = slicedims(f, s, I)
+    return rebuild(s; data=unrolled_map(parent, layers), dims, refdims, layerdims)
 end
 
 """
@@ -157,7 +162,12 @@ Base.similar(s::AbstractDimStack, ::Type{T}, dt::DimTuple; kw...) where T = mapl
 # NamedTuple-like
 @assume_effects :foldable Base.getproperty(s::AbstractDimStack, x::Symbol) = s[x]
 Base.haskey(s::AbstractDimStack{K}, k) where K = k in K
-Base.values(s::AbstractDimStack) = values(layers(s))
+Base.values(s::AbstractDimStack) = _values_gen(s)
+@generated function _values_gen(s::AbstractDimStack{K}) where K
+    Expr(:tuple, map(k -> :(s[$(QuoteNode(k))]), K)...)
+end
+Base.checkbounds(s::AbstractDimStack, I...) = checkbounds(CartesianIndices(s), I...)
+Base.checkbounds(T::Type, s::AbstractDimStack, I...) = checkbounds(T, CartesianIndices(s), I...)
 
 @inline Base.keys(s::AbstractDimStack{K}) where K = K
 @inline Base.propertynames(s::AbstractDimStack{K}) where K = K
@@ -197,7 +207,8 @@ Base.map(f, s::AbstractDimStack) = error("Use maplayers(f, stack)) instad of map
 Base.map(f, ::Union{AbstractDimStack,NamedTuple}, xs::Union{AbstractDimStack,NamedTuple}...) =
     error("Use maplayers(f, stack, args...)) instad of map(f, stack, args...)")
 
-maplayers(f, s::AbstractDimStack) = _maybestack(s, map(f, values(s)))
+maplayers(f, s::AbstractDimStack) =
+    _maybestack(s, unrolled_map(f, values(s)))
 function maplayers(
     f, x1::Union{AbstractDimStack,NamedTuple}, xs::Union{AbstractDimStack,NamedTuple}...
 )

src/utils.jl

@@ -214,3 +214,24 @@ function _as_extended_nts(nt::NamedTuple{K1}, st::AbstractDimStack{K2}, As...) w
     return (extended_layers, _as_extended_nts(nt, As...)...)
 end
 _as_extended_nts(::NamedTuple) = ()
+
+
+# Tuple map that is always unrolled
+# mostly for stack indexing performance
+_unrolled_map_inner(f, v::Type{T}) where T = 
+    Expr(:tuple, (:(f(v[$i])) for i in eachindex(T.types))...)
+_unrolled_map_inner(f, v1::Type{T}, v2::Type) where T = 
+    Expr(:tuple, (:(f(v1[$i], v2[$i])) for i in eachindex(T.types))...)
+
+@generated function unrolled_map(f, v::NamedTuple{K}) where K
+    exp = _unrolled_map_inner(f, v)
+    :(NamedTuple{K}($exp))
+end
+@generated function unrolled_map(f, v1::NamedTuple{K}, v2::NamedTuple{K}) where K
+    exp = _unrolled_map_inner(f, v1, v2)
+    :(NamedTuple{K}($exp))
+end
+@generated unrolled_map(f, v::Tuple) =
+    _unrolled_map_inner(f, v)
+@generated unrolled_map(f, v1::Tuple, v2::Tuple) = 
+    _unrolled_map_inner(f, v1, v2)