Merge pull request #98 from Tokazama/axes_methods

Make working with axes (and their types) easy

Author: chriselrod

src/ArrayInterface.jl

@@ -13,7 +13,7 @@ parameterless_type(x::Type) = __parameterless_type(x)
 """
     parent_type(::Type{T})
 
-Returns the parent array that `x` wraps.
+Returns the parent array that type `T` wraps.
 """
 parent_type(x) = parent_type(typeof(x))
 parent_type(::Type{<:SubArray{T,N,P}}) where {T,N,P} = P
@@ -704,6 +704,12 @@ end
   end
 end
 
+include("static.jl")
+include("ranges.jl")
+include("dimensions.jl")
+include("indexing.jl")
+include("stridelayout.jl")
+
 function __init__()
 
   @require SuiteSparse="4607b0f0-06f3-5cda-b6b1-a6196a1729e9" begin
@@ -746,6 +752,10 @@ function __init__()
     stride_rank(::Type{T}) where {N, T <: StaticArrays.StaticArray{<:Any,<:Any,N}} = StrideRank{ntuple(identity, Val{N}())}()
     dense_dims(::Type{<:StaticArrays.StaticArray{S,T,N}}) where {S,T,N} = DenseDims{ntuple(_ -> true, Val(N))}()
     defines_strides(::Type{<:StaticArrays.MArray}) = true
+
+    @generated function axes_types(::Type{<:StaticArrays.StaticArray{S}}) where {S}
+        return Tuple{[StaticArrays.SOneTo{s} for s in S.parameters]...}
+    end
     @generated function size(A::StaticArrays.StaticArray{S}) where {S}
         t = Expr(:tuple); Sp = S.parameters
         for n in 1:length(Sp)
@@ -896,10 +906,4 @@ function __init__()
   end
 end
 
-include("static.jl")
-include("ranges.jl")
-include("dimensions.jl")
-include("indexing.jl")
-include("stridelayout.jl")
-
 end

src/stridelayout.jl

@@ -74,7 +74,7 @@ end
 
 
 """
-contiguous_axis_indicator(::Type{T}) -> Tuple{Vararg{<:Val}}
+    contiguous_axis_indicator(::Type{T}) -> Tuple{Vararg{<:Val}}
 
 Returns a tuple boolean `Val`s indicating whether that axis is contiguous.
 """
@@ -84,14 +84,14 @@ contiguous_axis_indicator(::Nothing, ::Val) = nothing
 Base.@pure contiguous_axis_indicator(::Contiguous{N}, ::Val{D}) where {N,D} = ntuple(d -> Val{d == N}(), Val{D}())
 
 """
-If the contiguous dimension is not the dimension with `Stride_rank{1}`:
+If the contiguous dimension is not the dimension with `StrideRank{1}`:
 """
 struct ContiguousBatch{N} end
 Base.@pure ContiguousBatch(N::Int) = ContiguousBatch{N}()
 _get(::ContiguousBatch{N}) where {N} = N
 
 """
-contiguous_batch_size(::Type{T}) -> ContiguousBatch{N}
+    contiguous_batch_size(::Type{T}) -> ContiguousBatch{N}
 
 Returns the Base.size of contiguous batches if `!isone(stride_rank(T, contiguous_axis(T)))`.
 If `isone(stride_rank(T, contiguous_axis(T)))`, then it will return `ContiguousBatch{0}()`.
@@ -126,7 +126,7 @@ Base.collect(::StrideRank{R}) where {R} = collect(R)
 @inline Base.getindex(::StrideRank{R}, ::Val{I}) where {R,I} = StrideRank{permute(R, I)}()
 
 """
-rank_to_sortperm(::StrideRank) -> NTuple{N,Int}
+    rank_to_sortperm(::StrideRank) -> NTuple{N,Int}
 
 Returns the `sortperm` of the stride ranks.
 """
@@ -177,7 +177,9 @@ stride_rank(x, i) = stride_rank(x)[i]
 stride_rank(::Type{R}) where {T, N, S, A <: Array{S}, R <: Base.ReinterpretArray{T, N, S, A}} = StrideRank{ntuple(identity, Val{N}())}()
 
 """
-is_column_major(A) -> Val{true/false}()
+    is_column_major(A) -> Val{true/false}()
+
+Returns `Val{true}` if elements of `A` are stored in column major order. Otherwise returns `Val{false}`.
 """
 is_column_major(A) = is_column_major(stride_rank(A), contiguous_batch_size(A))
 is_column_major(::Nothing, ::Any) = Val{false}()
@@ -197,7 +199,7 @@ Base.@pure DenseDims(D::NTuple{<:Any,Bool}) = DenseDims{D}()
 @inline Base.getindex(::DenseDims{D}, i::Integer) where {D} = D[i]
 @inline Base.getindex(::DenseDims{D}, ::Val{I}) where {D,I} = DenseDims{permute(D, I)}()
 """
-dense_dims(::Type{T}) -> NTuple{N,Bool}
+    dense_dims(::Type{T}) -> NTuple{N,Bool}
 
 Returns a tuple of indicators for whether each axis is dense.
 An axis `i` of array `A` is dense if `stride(A, i) * Base.size(A, i) == stride(A, j)` where `stride_rank(A)[i] + 1 == stride_rank(A)[j]`.
@@ -250,7 +252,7 @@ permute(t::NTuple{N}, I::NTuple{N,Int}) where {N} = ntuple(n -> t[I[n]], Val{N}(
 end
 
 """
-  strides(A)
+    strides(A) -> Tuple
 
 Returns the strides of array `A`. If any strides are known at compile time,
 these should be returned as `Static` numbers. For example:
@@ -274,8 +276,196 @@ while still producing correct behavior when using valid cartesian indices, such
 strides(A) = Base.strides(A)
 strides(A, d) = strides(A)[to_dims(A, d)]
 
+@generated function _perm_tuple(::Type{T}, ::Val{P}) where {T,P}
+    out = Expr(:curly, :Tuple)
+    for p in P
+        push!(out.args, :(T.parameters[$p]))
+    end
+    Expr(:block, Expr(:meta, :inline), out)
+end
+
+"""
+    axes_types(::Type{T}[, d]) -> Type
+
+Returns the type of the axes for `T`
+"""
+axes_types(x) = axes_types(typeof(x))
+axes_types(x, d) = axes_types(typeof(x), d)
+@inline axes_types(::Type{T}, d) where {T} = axes_types(T).parameters[to_dims(T, d)]
+function axes_types(::Type{T}) where {T}
+    if parent_type(T) <: T
+        return Tuple{Vararg{OptionallyStaticUnitRange{One,Int},ndims(T)}}
+    else
+        return axes_types(parent_type(T))
+    end
+end
+axes_types(::Type{T}) where {T<:Adjoint} = _perm_tuple(axes_types(parent_type(T)), Val((2, 1)))
+axes_types(::Type{T}) where {T<:Transpose} = _perm_tuple(axes_types(parent_type(T)), Val((2, 1)))
+function axes_types(::Type{T}) where {I1,T<:PermutedDimsArray{<:Any,<:Any,I1}}
+    return _perm_tuple(axes_types(parent_type(T)), Val(I1))
+end
+function axes_types(::Type{T}) where {T<:AbstractRange}
+    if known_length(T) === nothing
+        return Tuple{OptionallyStaticUnitRange{One,Int}}
+    else
+        return Tuple{OptionallyStaticUnitRange{One,StaticInt{known_length(T)}}}
+    end
+end
+
+@inline function axes_types(::Type{T}) where {P,I,T<:SubArray{<:Any,<:Any,P,I}}
+    return _sub_axes_types(Val(ArrayStyle(T)), I, axes_types(P))
+end
+@generated function _sub_axes_types(::Val{S}, ::Type{I}, ::Type{PI}) where {S,I<:Tuple,PI<:Tuple}
+    out = Expr(:curly, :Tuple)
+    d = 1
+    for i in I.parameters
+        ad = argdims(S, i)
+        if ad > 0
+            push!(out.args, :(sub_axis_type($(PI.parameters[d]), $i)))
+            d += ad
+        else
+            d += 1
+        end
+    end
+    Expr(:block, Expr(:meta, :inline), out)
+end
+
+@inline function axes_types(::Type{T}) where {T<:Base.ReinterpretArray}
+    return _reinterpret_axes_types(axes_types(parent_type(T)), eltype(T), eltype(parent_type(T)))
+end
+@generated function _reinterpret_axes_types(::Type{I}, ::Type{T}, ::Type{S}) where {I<:Tuple,T,S}
+    out = Expr(:curly, :Tuple)
+    for i in 1:length(I.parameters)
+        if i === 1
+            push!(out.args, reinterpret_axis_type(I.parameters[1], T, S))
+        else
+            push!(out.args, I.parameters[i])
+        end
+    end
+    Expr(:block, Expr(:meta, :inline), out)
+end
+
+
+# These methods help handle identifying axes that dont' directly propagate from the
+# parent array axes. They may be worth making a formal part of the API, as they provide
+# a low traffic spot to change what axes_types produces.
+@inline function sub_axis_type(::Type{A}, ::Type{I}) where {A,I}
+    if known_length(I) === nothing
+        return OptionallyStaticUnitRange{One,Int}
+    else
+        return OptionallyStaticUnitRange{One,StaticInt{known_length(I)}}
+    end
+end
+
+@inline function reinterpret_axis_type(::Type{A}, ::Type{T}, ::Type{S}) where {A,T,S}
+    if known_length(A) === nothing
+        return OptionallyStaticUnitRange{One,Int}
+    else
+        return OptionallyStaticUnitRange{One,StaticInt{Int(known_length(A) / (sizeof(T) / sizeof(S)))}}
+    end
+end
+
 """
-  offsets(A)
+    known_offsets(::Type{T}[, d]) -> Tuple
+
+Returns a tuple of offset values known at compile time. If the offset of a given axis is
+not known at compile time `nothing` is returned its position.
+"""
+@inline known_offsets(x, d) = known_offsets(x)[to_dims(x, d)]
+known_offsets(x) = known_offsets(typeof(x))
+@generated function known_offsets(::Type{T}) where {T}
+    out = Expr(:tuple)
+    for p in axes_types(T).parameters
+        push!(out.args, known_first(p))
+    end
+    return out
+end
+
+"""
+    known_size(::Type{T}[, d]) -> Tuple
+
+Returns the size of each dimension for `T` known at compile time. If a dimension does not
+have a known size along a dimension then `nothing` is returned in its position.
+"""
+@inline known_size(x, d) = known_size(x)[to_dims(x, d)]
+known_size(x) = known_size(typeof(x))
+known_size(::Type{T}) where {T} = _known_size(axes_types(T))
+@generated function _known_size(::Type{Axs}) where {Axs<:Tuple}
+    out = Expr(:tuple)
+    for p in Axs.parameters
+        push!(out.args, :(known_length($p)))
+    end
+    return Expr(:block, Expr(:meta, :inline), out)
+end
+
+"""
+    known_strides(::Type{T}[, d]) -> Tuple
+
+Returns the strides of array `A` known at compile time. Any strides that are not known at
+compile time are represented by `nothing`.
+"""
+known_strides(x) = known_strides(typeof(x))
+known_strides(x, d) = known_strides(x)[to_dims(x, d)]
+known_strides(::Type{T}) where {T<:Vector} = (1,)
+@inline function known_strides(::Type{T}) where {T<:Adjoint{<:Any,<:AbstractVector}}
+    strd = first(known_strides(parent_type(T)))
+    return (strd, strd)
+end
+function known_strides(::Type{T}) where {T<:Adjoint}
+    return permute(known_strides(parent_type(T)), Val{(2,1)}())
+end
+function known_strides(::Type{T}) where {T<:Transpose}
+    return permute(known_strides(parent_type(T)), Val{(2,1)}())
+end
+@inline function known_strides(::Type{T}) where {T<:Transpose{<:Any,<:AbstractVector}}
+    strd = first(known_strides(parent_type(T)))
+    return (strd, strd)
+end
+@inline function known_strides(::Type{T}) where {I1,T<:PermutedDimsArray{<:Any,<:Any,I1}}
+    return permute(known_strides(parent_type(T)), Val{I1}())
+end
+@inline function known_strides(::Type{T}) where {I1,T<:SubArray{<:Any,<:Any,<:Any,I1}}
+    return _sub_strides(Val(ArrayStyle(T)), I1, Val(known_strides(parent_type(T))))
+end
+
+@generated function _sub_strides(::Val{S}, ::Type{I}, ::Val{P}) where {S,I<:Tuple,P}
+    out = Expr(:tuple)
+    d = 1
+    for i in I.parameters
+        ad = argdims(S, i)
+        if ad > 0
+            push!(out.args, P[d])
+            d += ad
+        else
+            d += 1
+        end
+    end
+    Expr(:block, Expr(:meta, :inline), out)
+end
+
+function known_strides(::Type{T}) where {T}
+    if ndims(T) === 1
+        return (1,)
+    else
+        return _known_strides(Val(Base.front(known_size(T))))
+    end
+end
+@generated function _known_strides(::Val{S}) where {S}
+    out = Expr(:tuple)
+    N = length(S)
+    push!(out.args, 1)
+    for s in S
+        if s === nothing || out.args[end] === nothing
+            push!(out.args, nothing)
+        else
+            push!(out.args, out.args[end] * s)
+        end
+    end
+    return Expr(:block, Expr(:meta, :inline), out)
+end
+
+"""
+    offsets(A) -> Tuple
 
 Returns offsets of indices with respect to 0. If values are known at compile time,
 it should return them as `Static` numbers.
@@ -294,7 +484,7 @@ end
     strd = stride(parent(x), One())
     (strd, strd)
 end
-                
+
 @generated function _strides(A::AbstractArray{T,N}, s::NTuple{N}, ::Contiguous{C}) where {T,N,C}
     if C ≤ 0 || C > N
         return Expr(:block, Expr(:meta,:inline), :s)
@@ -325,15 +515,11 @@ if VERSION ≥ v"1.6.0-DEV.1581"
         quote
             $(Expr(:meta,:inline))
             @inbounds $stup
-        end        
+        end
     end
 end
 
-@inline function offsets(x, i)
-    inds = indices(x, i)
-    start = known_first(inds)
-    isnothing(start) ? first(inds) : StaticInt(start)
-end
+@inline offsets(x, i) = static_first(indices(x, i))
 # @inline offsets(A::AbstractArray{<:Any,N}) where {N} = ntuple(n -> offsets(A, n), Val{N}())
 # Explicit tuple needed for inference.
 @generated function offsets(A::AbstractArray{<:Any,N}) where {N}
@@ -344,6 +530,7 @@ end
     Expr(:block, Expr(:meta, :inline), t)
 end
 
+@inline size(v::AbstractVector) = (static_length(axes_types(v, 1)),)
 @inline size(B::Union{Transpose{T,A},Adjoint{T,A}}) where {T,A<:AbstractMatrix{T}} = permute(size(parent(B)), Val{(2,1)}())
 @inline size(B::PermutedDimsArray{T,N,I1,I2,A}) where {T,N,I1,I2,A<:AbstractArray{T,N}} = permute(size(parent(B)), Val{I1}())
 @inline size(A::AbstractArray, ::StaticInt{N}) where {N} = size(A)[N]