OptionallyStaticRange -> OptionallyStaticUnitRange

Also indices now produces a `Base.Slice` whenever possible.

Author: Tokazama

src/ArrayInterface.jl

@@ -4,6 +4,8 @@ using Requires
 using LinearAlgebra
 using SparseArrays
 
+using Base: OneTo
+
 Base.@pure __parameterless_type(T) = Base.typename(T).wrapper
 parameterless_type(x) = parameterless_type(typeof(x))
 parameterless_type(x::Type) = __parameterless_type(x)
@@ -20,8 +22,21 @@ parent_type(::Type{Adjoint{T,S}}) where {T,S} = S
 parent_type(::Type{Transpose{T,S}}) where {T,S} = S
 parent_type(::Type{Symmetric{T,S}}) where {T,S} = S
 parent_type(::Type{<:LinearAlgebra.AbstractTriangular{T,S}}) where {T,S} = S
+parent_type(::Type{<:PermutedDimsArray{T,N,I1,I2,A}}) where {T,N,I1,I2,A} = A
+parent_type(::Type{Base.Slice{T}}) where {T} = T
 parent_type(::Type{T}) where {T} = T
 
+"""
+    known_length(::Type{T})
+
+If `length` of an instance of type `T` is known at compile time, return it.
+Otherwise, return `nothing`.
+"""
+known_length(x) = known_length(typeof(x))
+known_length(::Type{<:NTuple{N,<:Any}}) where {N} = N
+known_length(::Type{<:NamedTuple{L}}) where {L} = length(L)
+known_length(::Type{T}) where {T<:Base.Slice} = known_length(parent_type(T))
+
 """
     can_change_size(::Type{T}) -> Bool
 
@@ -536,6 +551,7 @@ function __init__()
 
     known_first(::Type{<:StaticArrays.SOneTo}) = 1
     known_last(::Type{StaticArrays.SOneTo{N}}) where {N} = N
+    known_length(::Type{StaticArrays.SOneTo{N}}) where {N} = N
 
     @require Adapt="79e6a3ab-5dfb-504d-930d-738a2a938a0e" begin
       function Adapt.adapt_storage(::Type{<:StaticArrays.SArray{S}},xs::Array) where S

src/ranges.jl

@@ -1,6 +1,6 @@
 
 """
-known_first(::Type{T})
+    known_first(::Type{T})
 
 If `first` of an instance of type `T` is known at compile time, return it.
 Otherwise, return `nothing`.
@@ -11,9 +11,11 @@ Otherwise, return `nothing`.
 known_first(x) = known_first(typeof(x))
 known_first(::Type{T}) where {T} = nothing
 known_first(::Type{Base.OneTo{T}}) where {T} = one(T)
+known_first(::Type{T}) where {T<:Base.Slice} = known_first(parent_type(T))
+
 
 """
-known_last(::Type{T})
+    known_last(::Type{T})
 
 If `last` of an instance of type `T` is known at compile time, return it.
 Otherwise, return `nothing`.
@@ -24,9 +26,10 @@ using StaticArrays
 """
 known_last(x) = known_last(typeof(x))
 known_last(::Type{T}) where {T} = nothing
+known_last(::Type{T}) where {T<:Base.Slice} = known_last(parent_type(T))
 
 """
-known_step(::Type{T})
+    known_step(::Type{T})
 
 If `step` of an instance of type `T` is known at compile time, return it.
 Otherwise, return `nothing`.
@@ -38,77 +41,134 @@ known_step(x) = known_step(typeof(x))
 known_step(::Type{T}) where {T} = nothing
 known_step(::Type{<:AbstractUnitRange{T}}) where {T} = one(T)
 
-_eltype(::Type{T}) where {T} = T
-_eltype(::Type{Val{V}}) where {V} = typeof(V)
+# add methods to support ArrayInterface
 
-struct OptionallyStaticRange{T<:Integer,F,S,L} <: OrdinalRange{T,T}
-  start::F
-  step::S
-  stop::L
+_get(x) = x
+_get(::Val{V}) where {V} = V
+_convert(::Type{T}, x) where {T} = convert(T, x)
+_convert(::Type{T}, ::Val{V}) where {T,V} = Val(convert(T, V))
 
-  function OptionallyStaticRange(start::F, step::S, stop::L) where {F,S,L}
-      T = promote_type(_eltype(F), _eltype(S), eltype(L))
-      return new{T,F,S,L}(start, step, stop)
-  end
+"""
+    OptionallyStaticUnitRange{T<:Integer}(start, stop) <: OrdinalRange{T,T}
+
+This range permits diverse representations of arrays to comunicate common information 
+about their indices. Each field may be an integer or `Val(<:Integer)` if it is known
+at compile time. An `OptionallyStaticUnitRange` is intended to be constructed internally
+from other valid indices. Therefore, users should not expect the same checks are used
+to ensure construction of a valid `OptionallyStaticUnitRange`.
+"""
+struct OptionallyStaticUnitRange{T,F,L} <: AbstractUnitRange{T}
+    start::F
+    stop::L
+
+    function OptionallyStaticUnitRange{T}(start, stop) where {T<:Real}
+        if _get(start) isa T
+            if _get(stop) isa T
+                return new{T,typeof(start),typeof(stop)}(start, stop)
+            else
+                return OptionallyStaticUnitRange{T}(start, _convert(T, stop))
+            end
+        else
+            return OptionallyStaticUnitRange{T}(_convert(T, start), stop)
+        end
+    end
+
+    function OptionallyStaticUnitRange(start, stop)
+        T = promote_type(typeof(_get(start)), typeof(_get(stop)))
+        return OptionallyStaticUnitRange{T}(start, stop)
+    end
 end
 
-Base.first(r::OptionallyStaticRange{<:Any,Val{F}}) where {F} = F
-Base.first(r::OptionallyStaticRange{<:Any,<:Any}) = getfield(r, :start)
+Base.first(r::OptionallyStaticUnitRange{<:Any,Val{F}}) where {F} = F
+Base.first(r::OptionallyStaticUnitRange{<:Any,<:Any}) = r.start
 
-Base.step(r::OptionallyStaticRange{<:Any,<:Any,Val{S}}) where {S} = S
-Base.step(r::OptionallyStaticRange{<:Any,<:Any,<:Any}) = getfield(r, :step)
+Base.step(r::OptionallyStaticUnitRange{T}) where {T} = oneunit(T)
 
-Base.last(r::OptionallyStaticRange{<:Any,<:Any,<:Any,Val{L}}) where {L} = L
-Base.last(r::OptionallyStaticRange{<:Any,<:Any,<:Any,<:Any}) = getfield(r, :stop)
+Base.last(r::OptionallyStaticUnitRange{<:Any,<:Any,Val{L}}) where {L} = L
+Base.last(r::OptionallyStaticUnitRange{<:Any,<:Any,<:Any}) = r.stop
 
-ArrayInterface.known_first(::OptionallyStaticRange{<:Any,Val{F}}) where {F} = F
-ArrayInterface.known_step(::OptionallyStaticRange{<:Any,<:Any,Val{S}}) where {S} =S
-ArrayInterface.known_last(::OptionallyStaticRange{<:Any,<:Any,<:Any,Val{L}}) where {L} = L
+known_first(::Type{<:OptionallyStaticUnitRange{<:Any,Val{F}}}) where {F} = F
+known_step(::Type{<:OptionallyStaticUnitRange{T}}) where {T} = one(T)
+known_last(::Type{<:OptionallyStaticUnitRange{<:Any,<:Any,Val{L}}}) where {L} = L
 
-function Base.isempty(r::OptionallyStaticRange)
-    return (first(r) != last(r)) & ((step(r) > zero(step(r))) != (last(r) > first(r)))
+function Base.isempty(r::OptionallyStaticUnitRange)
+    if known_first(r) === oneunit(eltype(r))
+        return unsafe_isempty_one_to(last(r))
+    else
+        return unsafe_isempty_unit_range(first(r), last(r))
+    end
 end
 
-@inline function Base.length(r::OptionallyStaticRange{T}) where {T}
-    if isempty(r)
-        return zero(T)
+unsafe_isempty_one_to(lst) = lst <= zero(lst)
+unsafe_isempty_unit_range(fst, lst) = fst > lst
+
+unsafe_isempty_unit_range(fst::T, lst::T) where {T} = Integer(lst - fst + one(T))
+
+unsafe_length_one_to(lst::T) where {T<:Int} = T(lst)
+unsafe_length_one_to(lst::T) where {T} = Integer(lst - zero(lst))
+
+Base.@propagate_inbounds function Base.getindex(r::OptionallyStaticUnitRange, i::Integer)
+    if known_first(r) === oneunit(r)
+        return get_index_one_to(r, i)
     else
-        if known_step(r) === oneunit(T)
-            if known_first(r) === oneunit(T)
-                return last(r)
-            else
-                return last(r) - first(r) + step(r)
-            end
-        else
-            return Integer(div((last(r) - first(r)) + step(r), step(r)))
-        end
+        return get_index_unit_range(r, i)
     end
 end
 
+@inline function get_index_one_to(r, i)
+    @boundscheck if ((i > 0) & (i <= last(r)))
+        throw(BoundsError(r, i))
+    end
+    return convert(eltype(r), i)
+end
 
-isempty(r::StepRange) =
-    (r.start != r.stop) & ((r.step > zero(r.step)) != (r.stop > r.start))
-isempty(r::AbstractUnitRange) = first(r) > last(r)
-isempty(r::StepRangeLen) = length(r) == 0
-isempty(r::LinRange) = length(r) == 0
-
-# add methods to support ArrayInterface
+@inline function get_index_unit_range(r, i)
+    val = first(r) + (i - 1)
+    @boundscheck if i > 0 && val <= last(r) && val >= first(r)
+        throw(BoundsError(r, i))
+    end
+    return convert(eltype(r), val)
+end
 
 _try_static(x, y) = Val(x)
 _try_static(::Nothing, y) = Val(y)
 _try_static(x, ::Nothing) = Val(x)
 _try_static(::Nothing, ::Nothing) = nothing
 
-@inline function _pick_range(x, y)
-    fst = _try_static(known_first(x), known_first(y))
-    fst = fst === nothing ? first(x) : fst
+###
+### length
+###
+@inline function known_length(::Type{T}) where {T<:AbstractUnitRange}
+  fst = known_first(T)
+  lst = known_last(T)
+  if stp === nothing || fst === nothing || lst === nothing
+    return nothing
+  else
+    if fst === oneunit(eltype(T))
+      return unsafe_length_one_to(lst)
+    else
+      return unsafe_length_unit_range(fst, lst)
+    end
+  end
+end
 
-    st = _try_static(known_step(x), known_step(y))
-    st = st === nothing ? step(x) : st
+function Base.length(r::OptionallyStaticUnitRange{T}) where {T}
+    if isempty(r)
+        return zero(T)
+    else
+        if known_one(r) === one(T)
+            return unsafe_length_one_to(last(r))
+        else
+            return unsafe_length_unit_range(first(r), last(r))
+        end
+    end
+end
 
-    lst = _try_static(known_last(x), known_last(y))
-    lst = lst === nothing ? last(x) : lst
-    return OptionallyStaticRange(fst, st, lst)
+function unsafe_length_unit_range(fst::T, lst::T) where {T<:Union{Int,Int64,Int128}}
+    return Base.checked_add(Base.checked_sub(lst, fst), one(T))
+end
+function unsafe_length_unit_range(fst::T, lst::T) where {T<:Union{UInt,UInt64,UInt128}}
+    return Base.checked_add(lst - fst, one(T))
 end
 
 """
@@ -120,7 +180,14 @@ returned. If any indices are not equal along dimension `d` an error is thrown. A
 tuple may be used to specify a different dimension for each array. If `d` is not
 specified then indices for visiting each index of `x` is returned.
 """
-@inline indices(x) = eachindex(x)
+@inline function indices(x)
+    inds = eachindex(x)
+    if inds isa AbstractUnitRange{<:Integer}
+        return Base.Slice(inds)
+    else
+        return inds
+    end
+end
 
 indices(x, d) = indices(axes(x, d))
 
@@ -136,4 +203,12 @@ end
   return reduce(_pick_range, inds)
 end
 
+@inline function _pick_range(x, y)
+    fst = _try_static(known_first(x), known_first(y))
+    fst = fst === nothing ? first(x) : fst
+
+    lst = _try_static(known_last(x), known_last(y))
+    lst = lst === nothing ? last(x) : lst
+    return Base.Slice(OptionallyStaticUnitRange(fst, lst))
+end