Merge pull request #74 from Tokazama/steprange

OptionallyStaticStepRange

Author: chriselrod

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "2.13.3"
+version = "2.13.4"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/ranges.jl

@@ -43,7 +43,7 @@ known_step(::Type{<:AbstractUnitRange{T}}) where {T} = one(T)
 # add methods to support ArrayInterface
 
 """
-    OptionallyStaticUnitRange{T<:Integer}(start, stop) <: OrdinalRange{T,T}
+    OptionallyStaticUnitRange(start, stop) <: AbstractUnitRange{Int}
 
 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
@@ -67,21 +67,15 @@ struct OptionallyStaticUnitRange{F <: Integer, L <: Integer} <: AbstractUnitRang
     end
   end
 
- function OptionallyStaticUnitRange(x::AbstractRange)
+  function OptionallyStaticUnitRange(x::AbstractRange)
     if step(x) == 1
-      fst = static_first(x)
-      lst = static_last(x)
-      return OptionallyStaticUnitRange(fst, lst)
+      return OptionallyStaticUnitRange(static_first(x), static_last(x))
     else
         throw(ArgumentError("step must be 1, got $(step(r))"))
     end
   end
 end
 
-Base.:(:)(L::Integer, ::StaticInt{U}) where {U} = OptionallyStaticUnitRange(L, StaticInt(U))
-Base.:(:)(::StaticInt{L}, U::Integer) where {L} = OptionallyStaticUnitRange(StaticInt(L), U)
-Base.:(:)(::StaticInt{L}, ::StaticInt{U}) where {L,U} = OptionallyStaticUnitRange(StaticInt(L), StaticInt(U))
-
 Base.first(r::OptionallyStaticUnitRange) = r.start
 Base.step(::OptionallyStaticUnitRange) = StaticInt(1)
 Base.last(r::OptionallyStaticUnitRange) = r.stop
@@ -90,6 +84,110 @@ known_first(::Type{<:OptionallyStaticUnitRange{StaticInt{F}}}) where {F} = F
 known_step(::Type{<:OptionallyStaticUnitRange}) = 1
 known_last(::Type{<:OptionallyStaticUnitRange{<:Any,StaticInt{L}}}) where {L} = L
 
+"""
+    OptionallyStaticStepRange(start, step, stop) <: OrdinalRange{Int,Int}
+
+Similar to [`OptionallyStaticUnitRange`](@ref), `OptionallyStaticStepRange` permits
+a combination of static and standard primitive `Int`s to construct a range. It
+specifically enables the use of ranges without a step size of 1. It may be constructed
+through the use of `OptionallyStaticStepRange` directly or using static integers with
+the range operatore (i.e. `:`).
+
+```julia
+julia> using ArrayInterface
+
+julia> x = ArrayInterface.StaticInt(2);
+
+julia> x:x:10
+ArrayInterface.StaticInt{2}():ArrayInterface.StaticInt{2}():10
+
+julia> ArrayInterface.OptionallyStaticStepRange(x, x, 10)
+ArrayInterface.StaticInt{2}():ArrayInterface.StaticInt{2}():10
+```
+"""
+struct OptionallyStaticStepRange{F <: Integer, S <: Integer, L <: Integer} <: OrdinalRange{Int,Int}
+  start::F
+  step::S
+  stop::L
+
+  function OptionallyStaticStepRange(start, step, stop)
+    if eltype(start) <: Int
+      if eltype(stop) <: Int
+        lst = _steprange_last(start, step, stop)
+        return new{typeof(start),typeof(step),typeof(lst)}(start, step, lst)
+      else
+        return OptionallyStaticStepRange(start, step, Int(stop))
+      end
+    else
+      return OptionallyStaticStepRange(Int(start), step, stop)
+    end
+  end
+
+  function OptionallyStaticStepRange(x::AbstractRange)
+    return OptionallyStaticStepRange(static_first(x), static_step(x), static_last(x))
+  end
+end
+
+# to make StepRange constructor inlineable, so optimizer can see `step` value
+@inline function _steprange_last(start::StaticInt, step::StaticInt, stop::StaticInt)
+    return StaticInt(_steprange_last(Int(start), Int(step), Int(stop)))
+end
+@inline function _steprange_last(start::Integer, step::StaticInt, stop::StaticInt)
+    if step === one(step)
+        # we don't need to check the `stop` if we know it acts like a unit range
+        return stop
+    else
+        return _steprange_last(start, Int(step), Int(stop))
+    end
+end
+@inline function _steprange_last(start::Integer, step::Integer, stop::Integer)
+  z = zero(step)
+  if step === z
+    throw(ArgumentError("step cannot be zero"))
+  else
+    if stop == start
+        return Int(stop)
+    else
+      if step > z
+        if stop > start
+          return stop - Int(unsigned(stop - start) % step)
+        else
+          return Int(start - one(start))
+        end
+      else
+        if stop > start
+          return Int(start + one(start))
+        else
+          return stop + Int(unsigned(start - stop) % -step)
+        end
+      end
+    end
+  end
+end
+Base.first(r::OptionallyStaticStepRange) = r.start
+Base.step(r::OptionallyStaticStepRange) = r.step
+Base.last(r::OptionallyStaticStepRange) = r.stop
+
+known_first(::Type{<:OptionallyStaticStepRange{StaticInt{F}}}) where {F} = F
+known_step(::Type{<:OptionallyStaticStepRange{<:Any,StaticInt{S}}}) where {S} = S
+known_last(::Type{<:OptionallyStaticStepRange{<:Any,<:Any,StaticInt{L}}}) where {L} = L
+
+Base.:(:)(L::Integer, ::StaticInt{U}) where {U} = OptionallyStaticUnitRange(L, StaticInt(U))
+Base.:(:)(::StaticInt{L}, U::Integer) where {L} = OptionallyStaticUnitRange(StaticInt(L), U)
+Base.:(:)(::StaticInt{L}, ::StaticInt{U}) where {L,U} = OptionallyStaticUnitRange(StaticInt(L), StaticInt(U))
+Base.:(:)(::StaticInt{F}, ::StaticInt{S}, ::StaticInt{L}) where {F,S,L} = OptionallyStaticStepRange(StaticInt(F), StaticInt(S), StaticInt(L))
+Base.:(:)(start::Integer, ::StaticInt{S}, ::StaticInt{L}) where {S,L} = OptionallyStaticStepRange(start, StaticInt(S), StaticInt(L))
+Base.:(:)(::StaticInt{F}, ::StaticInt{S}, stop::Integer) where {F,S} = OptionallyStaticStepRange(StaticInt(F), StaticInt(S), stop)
+Base.:(:)(::StaticInt{F}, step::Integer, ::StaticInt{L}) where {F,L} = OptionallyStaticStepRange(StaticInt(F), step, StaticInt(L))
+Base.:(:)(start::Integer, step::Integer, ::StaticInt{L}) where {L} = OptionallyStaticStepRange(start, step, StaticInt(L))
+Base.:(:)(start::Integer, ::StaticInt{S}, stop::Integer) where {S} = OptionallyStaticStepRange(start, StaticInt(S), stop)
+Base.:(:)(::StaticInt{F}, step::Integer, stop::Integer) where {F} = OptionallyStaticStepRange(StaticInt(F), step, stop)
+Base.:(:)(::StaticInt{F}, ::StaticInt{1}, ::StaticInt{L}) where {F,L} = OptionallyStaticUnitRange(StaticInt(F), StaticInt(L))
+Base.:(:)(start::Integer, ::StaticInt{1}, ::StaticInt{L}) where {L} = OptionallyStaticUnitRange(start, StaticInt(L))
+Base.:(:)(::StaticInt{F}, ::StaticInt{1}, stop::Integer) where {F} = OptionallyStaticUnitRange(StaticInt(F), stop)
+Base.:(:)(start::Integer, ::StaticInt{1}, stop::Integer) = OptionallyStaticUnitRange(start, stop)
+
+
 function Base.isempty(r::OptionallyStaticUnitRange)
   if known_first(r) === oneunit(eltype(r))
     return unsafe_isempty_one_to(last(r))
@@ -98,13 +196,29 @@ function Base.isempty(r::OptionallyStaticUnitRange)
   end
 end
 
+function Base.isempty(r::OptionallyStaticStepRange)
+  return (r.start != r.stop) & ((r.step > zero(r.step)) != (r.stop > r.start))
+end
+
 unsafe_isempty_one_to(lst) = lst <= zero(lst)
 unsafe_isempty_unit_range(fst, lst) = fst > lst
 
 unsafe_length_one_to(lst::Int) = lst
-unsafe_length_one_to(::StaticInt{L}) where {L} = lst
+unsafe_length_one_to(::StaticInt{L}) where {L} = L
+
+@inline function unsafe_length_step_range(start::Int, step::Int, stop::Int)
+    if step > 1
+        return Base.checked_add(Int(div(unsigned(stop - start), step)), 1)
+    elseif step < -1
+        return Base.checked_add(Int(div(unsigned(start - stop), -step)), 1)
+    elseif step > 0
+        return Base.checked_add(Int(div(Base.checked_sub(stop, start), step)), 1)
+    else
+        return Base.checked_add(Int(div(Base.checked_sub(rtart, stop), -step)), 1)
+    end
+end
 
-Base.@propagate_inbounds function Base.getindex(r::OptionallyStaticUnitRange, i::Integer)
+@propagate_inbounds function Base.getindex(r::OptionallyStaticUnitRange, i::Integer)
   if known_first(r) === oneunit(eltype(r))
     return get_index_one_to(r, i)
   else
@@ -121,7 +235,7 @@ end
 
 @inline function get_index_unit_range(r, i)
   val = first(r) + (i - 1)
-  @boundscheck if (i < 1) || (val > last(r) && val < first(r))
+  @boundscheck if (i < 1) || val > last(r)
     throw(BoundsError(r, i))
   end
   return convert(eltype(r), val)
@@ -130,28 +244,28 @@ end
 @inline _try_static(::StaticInt{N}, ::StaticInt{N}) where {N} = StaticInt{N}()
 @inline _try_static(::StaticInt{M}, ::StaticInt{N}) where {M, N} = @assert false "Unequal Indices: StaticInt{$M}() != StaticInt{$N}()"
 @propagate_inbounds function _try_static(::StaticInt{N}, x) where {N}
-    @boundscheck begin
-        @assert N == x "Unequal Indices: StaticInt{$N}() != x == $x"
-    end
-    return StaticInt{N}()
+  @boundscheck begin
+     @assert N == x "Unequal Indices: StaticInt{$N}() != x == $x"
+  end
+  return StaticInt{N}()
 end
 @propagate_inbounds function _try_static(x, ::StaticInt{N}) where {N}
-    @boundscheck begin
-        @assert N == x "Unequal Indices: x == $x != StaticInt{$N}()"
-    end
-    return StaticInt{N}()
+  @boundscheck begin
+     @assert N == x "Unequal Indices: x == $x != StaticInt{$N}()"
+  end
+  return StaticInt{N}()
 end
 @propagate_inbounds function _try_static(x, y)
-    @boundscheck begin
-        @assert x == y "Unequal Indicess: x == $x != $y == y"
-    end
-    return x
+  @boundscheck begin
+    @assert x == y "Unequal Indicess: x == $x != $y == y"
+  end
+  return x
 end
 
 ###
 ### length
 ###
-@inline function known_length(::Type{T}) where {T<:AbstractUnitRange}
+@inline function known_length(::Type{T}) where {T<:OptionallyStaticUnitRange}
   fst = known_first(T)
   lst = known_last(T)
   if fst === nothing || lst === nothing
@@ -165,6 +279,25 @@ end
   end
 end
 
+@inline function known_length(::Type{T}) where {T<:OptionallyStaticStepRange}
+  fst = known_first(T)
+  stp = known_step(T)
+  lst = known_last(T)
+  if fst === nothing || stp === nothing || lst === nothing
+    return nothing
+  else
+    if stp === 1
+      if fst === oneunit(eltype(T))
+        return unsafe_length_one_to(lst)
+      else
+        return unsafe_length_unit_range(fst, lst)
+      end
+    else
+      return unsafe_length_step_range(fst, stp, lst)
+    end
+  end
+end
+
 function Base.length(r::OptionallyStaticUnitRange)
   if isempty(r)
     return 0
@@ -177,6 +310,23 @@ function Base.length(r::OptionallyStaticUnitRange)
   end
 end
 
+function Base.length(r::OptionallyStaticStepRange)
+  if isempty(r)
+    return 0
+  else
+    if known_step(r) === 1
+      if known_first(r) === 1
+        return unsafe_length_one_to(last(r))
+      else
+        return unsafe_length_unit_range(first(r), last(r))
+      end
+    else
+        return unsafe_length_step_range(Int(first(r)), Int(step(r)), Int(last(r)))
+    end
+  end
+end
+
+
 unsafe_length_unit_range(start::Integer, stop::Integer) = Int((stop - start) + 1)
 
 """
@@ -219,3 +369,4 @@ end
   lst = _try_static(static_last(x), static_last(y))
   return Base.Slice(OptionallyStaticUnitRange(fst, lst))
 end
+

src/static.jl

@@ -16,6 +16,10 @@ StaticInt(::StaticInt{N}) where {N} = StaticInt{N}()
 StaticInt(::Val{N}) where {N} = StaticInt{N}()
 # Base.Val(::StaticInt{N}) where {N} = Val{N}()
 Base.convert(::Type{T}, ::StaticInt{N}) where {T<:Number,N} = convert(T, N)
+Base.Bool(x::StaticInt{N}) where {N} = Bool(N)
+Base.BigInt(x::StaticInt{N}) where {N} = BigInt(N)
+Base.Integer(x::StaticInt{N}) where {N} = x
+(::Type{T})(x::StaticInt{N}) where {T<:Integer,N} = T(N)
 (::Type{T})(x::Int) where {T<:StaticInt} = StaticInt(x)
 Base.convert(::Type{StaticInt{N}}, ::StaticInt{N}) where {N} = StaticInt{N}()
 

test/runtests.jl

@@ -190,6 +190,21 @@ using ArrayInterface: parent_type
 end
 
 @testset "Range Interface" begin
+    @testset "Range Constructors" begin
+        @test @inferred(StaticInt(1):StaticInt(10)) == 1:10
+        @test @inferred(StaticInt(1):StaticInt(2):StaticInt(10)) == 1:2:10 
+        @test @inferred(1:StaticInt(2):StaticInt(10)) == 1:2:10
+        @test @inferred(StaticInt(1):StaticInt(2):10) == 1:2:10
+        @test @inferred(StaticInt(1):2:StaticInt(10)) == 1:2:10 
+        @test @inferred(1:2:StaticInt(10)) == 1:2:10
+        @test @inferred(1:StaticInt(2):10) == 1:2:10
+        @test @inferred(StaticInt(1):2:10) == 1:2:10 
+
+        @test @inferred(StaticInt(1):StaticInt(1):StaticInt(10)) === ArrayInterface.OptionallyStaticUnitRange(StaticInt(1), StaticInt(10))
+        @test @inferred(StaticInt(1):StaticInt(1):10) === ArrayInterface.OptionallyStaticUnitRange(StaticInt(1), 10)
+        @test @inferred(1:StaticInt(1):10) === ArrayInterface.OptionallyStaticUnitRange(1, 10)
+    end
+
     @test isnothing(@inferred(ArrayInterface.known_first(typeof(1:4))))
     @test isone(@inferred(ArrayInterface.known_first(Base.OneTo(4))))
     @test isone(@inferred(ArrayInterface.known_first(typeof(Base.OneTo(4)))))
@@ -201,14 +216,24 @@ end
     @test isone(@inferred(ArrayInterface.known_step(1:4)))
     @test isone(@inferred(ArrayInterface.known_step(typeof(1:4))))
 
-    @test @inferred(length(ArrayInterface.OptionallyStaticUnitRange(1, 0))) == 0
-    @test @inferred(length(ArrayInterface.OptionallyStaticUnitRange(1, 10))) == 10
-    @test @inferred(length(ArrayInterface.OptionallyStaticUnitRange(StaticInt(1), 10))) == 10
-    @test @inferred(length(ArrayInterface.OptionallyStaticUnitRange(StaticInt(0), 10))) == 11
+    @testset "length" begin
+        @test @inferred(length(ArrayInterface.OptionallyStaticUnitRange(1, 0))) == 0
+        @test @inferred(length(ArrayInterface.OptionallyStaticUnitRange(1, 10))) == 10
+        @test @inferred(length(ArrayInterface.OptionallyStaticUnitRange(StaticInt(1), 10))) == 10
+        @test @inferred(length(ArrayInterface.OptionallyStaticUnitRange(StaticInt(0), 10))) == 11
+        @test @inferred(length(ArrayInterface.OptionallyStaticUnitRange(StaticInt(1), StaticInt(10)))) == 10
+        @test @inferred(length(ArrayInterface.OptionallyStaticUnitRange(StaticInt(0), StaticInt(10)))) == 11
+
+        @test @inferred(length(StaticInt(1):StaticInt(2):StaticInt(0))) == 0
+        @test @inferred(length(StaticInt(0):StaticInt(-2):StaticInt(1))) == 0
+    end
     @test @inferred(getindex(ArrayInterface.OptionallyStaticUnitRange(StaticInt(1), 10), 1)) == 1
     @test @inferred(getindex(ArrayInterface.OptionallyStaticUnitRange(StaticInt(0), 10), 1)) == 0
     @test_throws BoundsError getindex(ArrayInterface.OptionallyStaticUnitRange(StaticInt(1), 10), 0)
-    @test_throws BoundsError getindex(ArrayInterface.OptionallyStaticUnitRange(StaticInt(1), 10), 0)
+    @test_throws BoundsError getindex(ArrayInterface.OptionallyStaticStepRange(StaticInt(1), 2, 10), 0)
+    @test_throws BoundsError getindex(ArrayInterface.OptionallyStaticUnitRange(StaticInt(1), 10), 11)
+    @test_throws BoundsError getindex(ArrayInterface.OptionallyStaticStepRange(StaticInt(1), 2, 10), 11)
+
 end
 
 @testset "Memory Layout" begin
@@ -287,12 +312,12 @@ using OffsetArrays
     M = @MArray zeros(2,3,4); Mp = @view(PermutedDimsArray(M,(3,1,2))[:,2,:])';
     Sp2 = @view(PermutedDimsArray(S,(3,2,1))[2:3,:,:]);
     Mp2 = @view(PermutedDimsArray(M,(3,1,2))[2:3,:,2])';
-    
+
     @test @inferred(ArrayInterface.size(A)) === (3,4,5)
     @test @inferred(ArrayInterface.size(Ap)) === (2,5)
     @test @inferred(ArrayInterface.size(A)) === size(A)
     @test @inferred(ArrayInterface.size(Ap)) === size(Ap)
-    
+
     @test @inferred(ArrayInterface.size(S)) === (StaticInt(2), StaticInt(3), StaticInt(4))
     @test @inferred(ArrayInterface.size(Sp)) === (2, 2, StaticInt(3))
     @test @inferred(ArrayInterface.size(Sp2)) === (2, StaticInt(3), StaticInt(2))
@@ -405,6 +430,11 @@ end
     @test @inferred(one(StaticInt)) === StaticInt(1)
     @test @inferred(zero(StaticInt)) === StaticInt(0)
     @test eltype(one(StaticInt)) <: Int
+
+    x = StaticInt(1)
+    @test @inferred(Bool(x)) isa Bool
+    @test @inferred(BigInt(x)) isa BigInt
+    @test @inferred(Integer(x)) === x
     # test for ambiguities and correctness
     for i ∈ [StaticInt(0), StaticInt(1), StaticInt(2), 3]
         for j ∈ [StaticInt(0), StaticInt(1), StaticInt(2), 3]