Simplify lazy_axes (#333)

Previously we used a generated function for part of the implementation. This PR resolves that.

Author: Tokazama

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "6.0.19"
+version = "6.0.20"
 
 [deps]
 ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2"

src/axes.jl

@@ -237,11 +237,14 @@ Base.show(io::IO, x::LazyAxis{N}) where {N} = print(io, "LazyAxis{$N}($(parent(x
 Produces a tuple of axes where each axis is constructed lazily. If an axis of `x` is already
 constructed or it is simply retrieved.
 """
-@generated function lazy_axes(x::X) where {X}
-    Expr(:block, Expr(:meta, :inline), Expr(:tuple, [:(LazyAxis{$dim}(x)) for dim in 1:ndims(X)]...))
-end
+@inline lazy_axes(x) = lazy_axes(x, ntuple(static, StaticInt(ndims(x))))
 lazy_axes(x::Union{LinearIndices,CartesianIndices,AbstractRange}) = axes(x)
-@inline lazy_axes(x::VecAdjTrans) = (SOneTo{1}(), first(lazy_axes(parent(x))))
-@inline function lazy_axes(x::Union{PermutedDimsArray,MatAdjTrans})
-    map(GetIndex{false}(lazy_axes(parent(x))), to_parent_dims(x))
-end
+@inline function lazy_axes(x::PermutedDimsArray, ::StaticInt{N}) where {N}
+    N <= ndims(x) ? lazy_axes(parent(x), getfield(to_parent_dims(x), N)) : SOneTo{1}()
+end
+lazy_axes(x::Union{Adjoint,Transpose}, ::StaticInt{1}) = lazy_axes(parent(x), StaticInt(2))
+lazy_axes(x::Union{Adjoint,Transpose}, ::StaticInt{2}) = lazy_axes(parent(x), StaticInt(1))
+lazy_axes(x::AbstractRange, ::StaticInt{1}) = Base.axes1(x)
+lazy_axes(x, ::Colon) = LazyAxis{:}(x)
+lazy_axes(x, ::StaticInt{dim}) where {dim} = ndims(x) < dim ? SOneTo{1}() : LazyAxis{dim}(x)
+@inline lazy_axes(x, dims::Tuple) = map(Base.Fix1(lazy_axes, x), dims)

src/indexing.jl

@@ -88,8 +88,8 @@ struct IndexedMappedArray{A}
     a::A
 end
 @inline (ima::IndexedMappedArray{A})(idx::I, ::StaticInt{0}) where {A,I} = to_index(StaticInt(1):StaticInt(1), idx)
-@inline (ima::IndexedMappedArray{A})(idx::I, ::Colon) where {A,I} = to_index(LazyAxis{:}(ima.a), idx)
-@inline (ima::IndexedMappedArray{A})(idx::I, ::StaticInt{dim}) where {A,I,dim} = to_index(LazyAxis{dim}(ima.a), idx)
+@inline (ima::IndexedMappedArray{A})(idx::I, ::Colon) where {A,I} = to_index(lazy_axes(ima.a, :), idx)
+@inline (ima::IndexedMappedArray{A})(idx::I, d::StaticInt{D}) where {A,I,D} = to_index(lazy_axes(ima.a, d), idx)
 @inline function (ima::IndexedMappedArray{A})(idx::AbstractArray{Bool}, dims::Tuple) where {A}
     if (last(dims) == ndims(A)) && (IndexStyle(A) isa IndexLinear)
         return LogicalIndex{Int}(idx)
@@ -99,9 +99,8 @@ end
 end
 @inline (ima::IndexedMappedArray{A})(idx::CartesianIndex, ::Tuple) where {A} = getfield(idx, 1)
 @inline function (ima::IndexedMappedArray{A})(idx::I, dims::Tuple) where {A,I}
-    to_index(CartesianIndices(map(Base.Fix1(_to_lazy_axes, ima.a), dims)), idx)
+    to_index(CartesianIndices(lazy_axes(ima.a, dims)), idx)
 end
-@inline _to_lazy_axes(a::A, ::StaticInt{dim}) where {A,dim} = LazyAxis{dim}(a)
 
 """
     ArrayInterface.to_index([::IndexStyle, ]axis, arg) -> index