Improve single element indexing (and other quality of life improvements) (#149)

* Use LazyAxis to avoid cost of materializing axes

* to_indices doesn't iterate through indexers that don't need conversion

* Patch bump and simplify some lazy_axes returns

* add change size trait

* Document canonicalize

* Fix type stability of ranges

* Add more axes tests

Author: Tokazama

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "3.1.9"
+version = "3.1.10"
 
 [deps]
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"

src/ArrayInterface.jl

@@ -16,8 +16,6 @@ using Base: @propagate_inbounds, tail, OneTo, LogicalIndex, Slice, ReinterpretAr
 ## utilites for internal use only ##
 _int_or_static_int(::Nothing) = Int
 _int_or_static_int(x::Int) = StaticInt{x}
-_int(i::Integer) = Int(i)
-_int(i::StaticInt) = i
 
 @static if VERSION >= v"1.7.0-DEV.421"
     using Base: @aggressive_constprop
@@ -842,10 +840,10 @@ end
 end
 
 include("ranges.jl")
-include("indexing.jl")
 include("dimensions.jl")
 include("axes.jl")
 include("size.jl")
+include("indexing.jl")
 include("stridelayout.jl")
 include("broadcast.jl")
 

src/axes.jl

@@ -37,11 +37,9 @@ function axes_types(::Type{T}) where {T}
 end
 axes_types(::Type{LinearIndices{N,R}}) where {N,R} = R
 axes_types(::Type{CartesianIndices{N,R}}) where {N,R} = R
-function axes_types(::Type{T}) where {T<:VecAdjTrans}
-    return Tuple{OptionallyStaticUnitRange{One,One},axes_types(parent_type(T), One())}
-end
-function axes_types(::Type{T}) where {T<:MatAdjTrans}
-    return eachop_tuple(_get_tuple, to_parent_dims(T), axes_types(parent_type(T)))
+function axes_types(::Type{T}) where {T<:Union{Adjoint,Transpose}}
+    P = parent_type(T)
+    return Tuple{axes_types(P, static(2)), axes_types(P, static(1))}
 end
 function axes_types(::Type{T}) where {T<:PermutedDimsArray}
     return eachop_tuple(_get_tuple, to_parent_dims(T), axes_types(parent_type(T)))
@@ -133,6 +131,21 @@ function axes(a::A, dim::Integer) where {A}
         return axes(parent(a), to_parent_dims(A, dim))
     end
 end
+function axes(A::CartesianIndices{N}, dim::Integer) where {N}
+    if dim > N
+        return static(1):static(1)
+    else
+        return getfield(axes(A), Int(dim))
+    end
+end
+function axes(A::LinearIndices{N}, dim::Integer) where {N}
+    if dim > N
+        return static(1):static(1)
+    else
+        return getfield(axes(A), Int(dim))
+    end
+end
+
 axes(A::SubArray, dim::Integer) = Base.axes(A, Int(dim))  # TODO implement ArrayInterface version
 axes(A::ReinterpretArray, dim::Integer) = Base.axes(A, Int(dim))  # TODO implement ArrayInterface version
 axes(A::Base.ReshapedArray, dim::Integer) = Base.axes(A, Int(dim))  # TODO implement ArrayInterface version
@@ -160,3 +173,137 @@ axes(A::Base.ReshapedArray) = Base.axes(A)  # TODO implement ArrayInterface vers
 axes(A::CartesianIndices) = A.indices
 axes(A::LinearIndices) = A.indices
 
+"""
+    LazyAxis{N}(parent::AbstractArray)
+
+A lazy representation of `axes(parent, N)`.
+"""
+struct LazyAxis{N,P} <: AbstractUnitRange{Int}
+    parent::P
+
+    LazyAxis{N}(parent::P) where {N,P} = new{N::Int,P}(parent)
+    @inline function LazyAxis{:}(parent::P) where {P}
+        if ndims(P) === 1
+            return new{1,P}(parent)
+        else
+            return new{:,P}(parent)
+        end
+    end
+end
+
+@inline Base.parent(x::LazyAxis{N,P}) where {N,P} = axes(getfield(x, :parent), static(N))
+@inline function Base.parent(x::LazyAxis{:,P}) where {P}
+    return eachindex(IndexLinear(), getfield(x, :parent))
+end
+
+@inline parent_type(::Type{LazyAxis{N,P}}) where {N,P} = axes_types(P, static(N))
+# TODO this approach to parent_type(::Type{LazyAxis{:}}) is a bit hacky. Something like
+# LabelledArrays has a linear set of symbolic keys, which could be propagated through
+# `to_indices` for key based indexing. However, there currently isn't a good way of handling
+# that when the linear indices aren't linearly accessible through a child array (e.g, adjoint)
+# For now we just make sure the linear elements are accurate.
+parent_type(::Type{LazyAxis{:,P}}) where {P<:Array} = OneTo{Int}
+@inline function parent_type(::Type{LazyAxis{:,P}}) where {P}
+    if known_length(P) === nothing
+        return OptionallyStaticUnitRange{StaticInt{1},Int}
+    else
+        return OptionallyStaticUnitRange{StaticInt{1},StaticInt{known_length(P)}}
+    end
+end
+
+Base.keys(x::LazyAxis) = keys(parent(x))
+
+Base.IndexStyle(::Type{T}) where {T<:LazyAxis} = IndexStyle(parent_type(T))
+
+can_change_size(::Type{LazyAxis{N,P}}) where {N,P} = can_change_size(P)
+
+known_first(::Type{T}) where {T<:LazyAxis} = known_first(parent_type(T))
+
+known_length(::Type{LazyAxis{N,P}}) where {N,P} = known_size(P, N)
+known_length(::Type{LazyAxis{:,P}}) where {P} = known_length(P)
+
+@inline function known_last(::Type{T}) where {T<:LazyAxis}
+    return _lazy_axis_known_last(known_first(T), known_length(T))
+end
+_lazy_axis_known_last(start::Int, length::Int) = (length + start) - 1
+_lazy_axis_known_last(::Any, ::Any) = nothing
+
+@inline function Base.first(x::LazyAxis{N})::Int where {N}
+    if known_first(x) === nothing
+        return offsets(getfield(x, :parent), static(N))
+    else
+        return known_first(x)
+    end
+end
+@inline function Base.first(x::LazyAxis{:})::Int
+    if known_first(x) === nothing
+        return firstindex(getfield(x, :parent))
+    else
+        return known_first(x)
+    end
+end
+
+@inline function Base.length(x::LazyAxis{N})::Int where {N}
+    if known_length(x) === nothing
+        return size(getfield(x, :parent), static(N))
+    else
+        return known_length(x)
+    end
+end
+@inline function Base.length(x::LazyAxis{:})::Int
+    if known_length(x) === nothing
+        return lastindex(getfield(x, :parent))
+    else
+        return known_length(x)
+    end
+end
+
+@inline function Base.last(x::LazyAxis)::Int
+    if known_last(x) === nothing
+        if known_first(x) === 1
+            return length(x)
+        else
+            return (static_length(x) + static_first(x)) - 1
+        end
+    else
+        return known_last(x)
+    end
+end
+
+Base.to_shape(x::LazyAxis) = length(x)
+
+@inline function Base.checkindex(::Type{Bool}, x::LazyAxis, i::Integer)
+    if known_first(x) === nothing || known_last(x) === nothing
+        return checkindex(Bool, parent(x), i)
+    else  # everything is static so we don't have to retrieve the axis
+        return (!(known_first(x) > i) || !(known_last(x) < i))
+    end
+end
+
+@propagate_inbounds function Base.getindex(x::LazyAxis, i::Integer)
+    @boundscheck checkindex(Bool, x, i) || throw(BoundsError(x, i))
+    return Int(i)
+end
+@propagate_inbounds Base.getindex(x::LazyAxis, i::StepRange{T}) where {T<:Integer} = parent(x)[i]
+@propagate_inbounds Base.getindex(x::LazyAxis, i::AbstractUnitRange{<:Integer}) = parent(x)[i]
+
+Base.show(io::IO, x::LazyAxis{N}) where {N} = print(io, "LazyAxis{$N}($(parent(x))))")
+
+"""
+    lazy_axes(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
+lazy_axes(x::LinearIndices) = axes(x)
+lazy_axes(x::CartesianIndices) = axes(x)
+@inline lazy_axes(x::MatAdjTrans) = reverse(lazy_axes(parent(x)))
+@inline lazy_axes(x::VecAdjTrans) = (LazyAxis{1}(x), first(lazy_axes(parent(x))))
+@inline lazy_axes(x::PermutedDimsArray) = permute(lazy_axes(parent(x)), to_parent_dims(A))
+

src/dimensions.jl

@@ -190,16 +190,13 @@ function dimnames(::Type{T}) where {T<:SubArray}
     return eachop(dimnames, to_parent_dims(T), parent_type(T))
 end
 
-_to_int(x::Integer) = Int(x)
-_to_int(x::StaticInt) = x
-
 """
     to_dims(::Type{T}, dim) -> Integer
 
 This returns the dimension(s) of `x` corresponding to `d`.
 """
 to_dims(x, dim) = to_dims(typeof(x), dim)
-to_dims(::Type{T}, dim::Integer) where {T} = _to_int(dim)
+to_dims(::Type{T}, dim::Integer) where {T} = canonicalize(dim)
 to_dims(::Type{T}, dim::Colon) where {T} = dim
 function to_dims(::Type{T}, dim::StaticSymbol) where {T}
     i = find_first_eq(dim, dimnames(T))