diff --git a/src/CloseOpenIntervals.jl b/src/CloseOpenIntervals.jl index 809b291..6f4cf6f 100644 --- a/src/CloseOpenIntervals.jl +++ b/src/CloseOpenIntervals.jl @@ -3,23 +3,23 @@ module CloseOpenIntervals using Static: StaticInt, Zero, One export CloseOpen, SafeCloseOpen -abstract type AbstractCloseOpen{L <: Integer, U <: Integer} <: AbstractUnitRange{Int} end +abstract type AbstractCloseOpen{L <: Union{Int,StaticInt}, U <: Union{Int,StaticInt}} <: AbstractUnitRange{Int} end for T ∈ (:CloseOpen,:SafeCloseOpen) @eval begin - struct $T{L <: Integer, U <: Integer} <: AbstractCloseOpen{L,U} + struct $T{L <: Union{Int,StaticInt}, U <: Union{Int,StaticInt}} <: AbstractCloseOpen{L,U} start::L upper::U - @inline $T{L,U}(l::L,u::U) where {L <: Integer, U <: Integer} = new{L,U}(l,u) + @inline $T{L,U}(l::L,u::U) where {L <: Union{Int,StaticInt}, U <: Union{Int,StaticInt}} = new{L,U}(l,u) end @inline $T(s::S, u::U) where {S,U} = $T{S,U}(s, u) - @inline $T(len::T) where {T<:Integer} = $T{Zero,T}(Zero(), len) + @inline $T(len::T) where {T<:Union{Int,StaticInt}} = $T{Zero,T}(Zero(), len) end end """ CloseOpen([start=0], U) <: AbstractUnitRange{Int} SafeCloseOpen([start=0], U) <: AbstractUnitRange{Int} - + Define close-open unit ranges, i.e. `CloseOpen(0,10)` iterates from from `0:9`. Close-open ranges can be more convenient in some circumstances, e.g. when partitioning a larger array. @@ -29,7 +29,7 @@ partitioning a larger array. function foo(x) nt = Threads.nthreads() d, r = divrem(length(x), nt) - i = firstindex(x) + i = firstindex(x) Threads.@sync for j in 1:nt stop = i + d + (r >= j) Threads.@spawn bar!(\$(@view(x[CloseOpen(i, stop)]))) @@ -60,16 +60,16 @@ See `?AbstractCloseOpen` for more information. SafeCloseOpen @inline Base.first(r::AbstractCloseOpen) = getfield(r,:start) -@inline Base.first(r::AbstractCloseOpen{StaticInt{F}}) where {F} = StaticInt{F}() -@inline Base.step(::AbstractCloseOpen) = One() -@inline Base.last(r::AbstractCloseOpen{<:Integer,I}) where {I} = getfield(r,:upper) - oneunit(I) -@inline Base.last(r::AbstractCloseOpen{<:Integer,StaticInt{L}}) where {L} = StaticInt{L}() - One() +@inline Base.first(r::AbstractCloseOpen{StaticInt{F}}) where {F} = F +@inline Base.step(::AbstractCloseOpen) = 1 +@inline Base.last(r::AbstractCloseOpen{<:Union{Int,StaticInt},I}) where {I} = getfield(r,:upper) - 1 +@inline Base.last(r::AbstractCloseOpen{<:Union{Int,StaticInt},StaticInt{L}}) where {L} = L - 1 @inline Base.length(r::AbstractCloseOpen) = getfield(r,:upper) - getfield(r,:start) @inline Base.length(r::AbstractCloseOpen{Zero}) = getfield(r,:upper) @inline Base.iterate(r::CloseOpen) = (i = Int(first(r)); (i, i)) @inline Base.iterate(r::SafeCloseOpen) = (i = Int(first(r)); i ≥ getfield(r, :upper) ? nothing : (i, i)) -@inline Base.iterate(r::AbstractCloseOpen, i::Integer) = (i += one(i)) ≥ getfield(r, :upper) ? nothing : (i, i) +@inline Base.iterate(r::AbstractCloseOpen, i::Union{Int,StaticInt}) = (i += one(i)) ≥ getfield(r, :upper) ? nothing : (i, i) import ArrayInterface ArrayInterface.known_first(::Type{<:AbstractCloseOpen{StaticInt{F}}}) where {F} = F