Skip to content

Commit cd47c71

Browse files
TokazamaTokazama
committed
Begin implementing OptionallyStaticRange
1 parent dbd3b42 commit cd47c71

File tree

2 files changed

+141
-79
lines changed

2 files changed

+141
-79
lines changed

‎src/ArrayInterface.jl‎

Lines changed: 2 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -503,85 +503,6 @@ function restructure(x::Array,y)
503503
reshape(convert(Array,y),size(x)...)
504504
end
505505

506-
"""
507-
known_first(::Type{T})
508-
509-
If `first` of an instance of type `T` is known at compile time, return it.
510-
Otherwise, return `nothing`.
511-
512-
@test isnothing(known_first(typeof(1:4)))
513-
@test isone(known_first(typeof(Base.OneTo(4))))
514-
"""
515-
known_first(x) = known_first(typeof(x))
516-
known_first(::Type{T}) where {T} = nothing
517-
known_first(::Type{Base.OneTo{T}}) where {T} = one(T)
518-
519-
"""
520-
known_last(::Type{T})
521-
522-
If `last` of an instance of type `T` is known at compile time, return it.
523-
Otherwise, return `nothing`.
524-
525-
@test isnothing(known_last(typeof(1:4)))
526-
using StaticArrays
527-
@test known_last(typeof(SOneTo(4))) == 4
528-
"""
529-
known_last(x) = known_last(typeof(x))
530-
known_last(::Type{T}) where {T} = nothing
531-
532-
"""
533-
known_step(::Type{T})
534-
535-
If `step` of an instance of type `T` is known at compile time, return it.
536-
Otherwise, return `nothing`.
537-
538-
@test isnothing(known_step(typeof(1:0.2:4)))
539-
@test isone(known_step(typeof(1:4)))
540-
"""
541-
known_step(x) = known_step(typeof(x))
542-
known_step(::Type{T}) where {T} = nothing
543-
known_step(::Type{<:AbstractUnitRange{T}}) where {T} = one(T)
544-
545-
546-
"""
547-
indices(x[, d]) -> AbstractRange
548-
549-
Given an array `x`, this returns the indices along dimension `d`. If `x` is a tuple
550-
of arrays then the indices corresponding to dimension `d` of all arrays in `x` are
551-
returned. If any indices are not equal along dimension `d` an error is thrown. A
552-
tuple may be used to specify a different dimension for each array. If `d` is not
553-
specified then indices for visiting each index of `x` is returned.
554-
"""
555-
@inline indices(x) = eachindex(x)
556-
557-
indices(x, d) = indices(axes(x, d))
558-
559-
@inline function indices(x::NTuple{N,<:Any}, dim) where {N}
560-
inds = indices(first(x), dim)
561-
@assert _check_indices(inds, Base.tail(x), dim) "The indices along dimension $dim are not equal for all $x"
562-
return inds
563-
end
564-
565-
@inline function indices(x::NTuple{N,<:Any}, dim::NTuple{N,<:Any}) where {N}
566-
ind = indices(first(x), first(dim))
567-
@assert _check_indices(ind, Base.tail(x), Base.tail(dim)) "The indices along dimension $dim are not equal for all $x"
568-
return ind
569-
end
570-
571-
@inline function _check_indices(ind, x::Tuple, dim::Tuple)
572-
for (x_i, d_i) in zip(x, dim)
573-
ind == indices(x_i, d_i) || return false
574-
end
575-
return true
576-
end
577-
578-
@inline function _check_indices(ind, x::Tuple, d)
579-
for x_i in x
580-
ind == indices(x_i, d) || return false
581-
end
582-
return true
583-
end
584-
585506
function __init__()
586507

587508
@require SuiteSparse="4607b0f0-06f3-5cda-b6b1-a6196a1729e9" begin
@@ -737,4 +658,6 @@ function __init__()
737658
end
738659
end
739660

661+
include("ranges.jl")
662+
740663
end

‎src/ranges.jl‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
2+
"""
3+
known_first(::Type{T})
4+
5+
If `first` of an instance of type `T` is known at compile time, return it.
6+
Otherwise, return `nothing`.
7+
8+
@test isnothing(known_first(typeof(1:4)))
9+
@test isone(known_first(typeof(Base.OneTo(4))))
10+
"""
11+
known_first(x) = known_first(typeof(x))
12+
known_first(::Type{T}) where {T} = nothing
13+
known_first(::Type{Base.OneTo{T}}) where {T} = one(T)
14+
15+
"""
16+
known_last(::Type{T})
17+
18+
If `last` of an instance of type `T` is known at compile time, return it.
19+
Otherwise, return `nothing`.
20+
21+
@test isnothing(known_last(typeof(1:4)))
22+
using StaticArrays
23+
@test known_last(typeof(SOneTo(4))) == 4
24+
"""
25+
known_last(x) = known_last(typeof(x))
26+
known_last(::Type{T}) where {T} = nothing
27+
28+
"""
29+
known_step(::Type{T})
30+
31+
If `step` of an instance of type `T` is known at compile time, return it.
32+
Otherwise, return `nothing`.
33+
34+
@test isnothing(known_step(typeof(1:0.2:4)))
35+
@test isone(known_step(typeof(1:4)))
36+
"""
37+
known_step(x) = known_step(typeof(x))
38+
known_step(::Type{T}) where {T} = nothing
39+
known_step(::Type{<:AbstractUnitRange{T}}) where {T} = one(T)
40+
41+
_eltype(::Type{T}) where {T} = T
42+
_eltype(::Type{Val{V}}) where {V} = typeof(V)
43+
44+
struct OptionallyStaticRange{T<:Integer,F,S,L} <: OrdinalRange{T,T}
45+
start::F
46+
step::S
47+
stop::L
48+
49+
function OptionallyStaticRange(start::F, step::S, stop::L) where {F,S,L}
50+
T = promote_type(_eltype(F), _eltype(S), eltype(L))
51+
return new{T,F,S,L}(start, step, stop)
52+
end
53+
end
54+
55+
Base.first(r::OptionallyStaticRange{<:Any,Val{F}}) where {F} = F
56+
Base.first(r::OptionallyStaticRange{<:Any,<:Any}) = getfield(r, :start)
57+
58+
Base.step(r::OptionallyStaticRange{<:Any,<:Any,Val{S}}) where {S} = S
59+
Base.step(r::OptionallyStaticRange{<:Any,<:Any,<:Any}) = getfield(r, :step)
60+
61+
Base.last(r::OptionallyStaticRange{<:Any,<:Any,<:Any,Val{L}}) where {L} = L
62+
Base.last(r::OptionallyStaticRange{<:Any,<:Any,<:Any,<:Any}) = getfield(r, :stop)
63+
64+
ArrayInterface.known_first(::OptionallyStaticRange{<:Any,Val{F}}) where {F} = F
65+
ArrayInterface.known_step(::OptionallyStaticRange{<:Any,<:Any,Val{S}}) where {S} =S
66+
ArrayInterface.known_last(::OptionallyStaticRange{<:Any,<:Any,<:Any,Val{L}}) where {L} = L
67+
68+
function Base.isempty(r::OptionallyStaticRange)
69+
return (first(r) != last(r)) & ((step(r) > zero(step(r))) != (last(r) > first(r)))
70+
end
71+
72+
@inline function Base.length(r::OptionallyStaticRange{T}) where {T}
73+
if isempty(r)
74+
return zero(T)
75+
else
76+
if known_step(r) === oneunit(T)
77+
if known_first(r) === oneunit(T)
78+
return last(r)
79+
else
80+
return last(r) - first(r) + step(r)
81+
end
82+
else
83+
return Integer(div((last(r) - first(r)) + step(r), step(r)))
84+
end
85+
end
86+
end
87+
88+
89+
isempty(r::StepRange) =
90+
(r.start != r.stop) & ((r.step > zero(r.step)) != (r.stop > r.start))
91+
isempty(r::AbstractUnitRange) = first(r) > last(r)
92+
isempty(r::StepRangeLen) = length(r) == 0
93+
isempty(r::LinRange) = length(r) == 0
94+
95+
# add methods to support ArrayInterface
96+
97+
_try_static(x, y) = Val(x)
98+
_try_static(::Nothing, y) = Val(y)
99+
_try_static(x, ::Nothing) = Val(x)
100+
_try_static(::Nothing, ::Nothing) = nothing
101+
102+
@inline function _pick_range(x, y)
103+
fst = _try_static(known_first(x), known_first(y))
104+
fst = fst === nothing ? first(x) : fst
105+
106+
st = _try_static(known_step(x), known_step(y))
107+
st = st === nothing ? step(x) : st
108+
109+
lst = _try_static(known_last(x), known_last(y))
110+
lst = lst === nothing ? last(x) : lst
111+
return OptionallyStaticRange(fst, st, lst)
112+
end
113+
114+
"""
115+
indices(x[, d]) -> AbstractRange
116+
117+
Given an array `x`, this returns the indices along dimension `d`. If `x` is a tuple
118+
of arrays then the indices corresponding to dimension `d` of all arrays in `x` are
119+
returned. If any indices are not equal along dimension `d` an error is thrown. A
120+
tuple may be used to specify a different dimension for each array. If `d` is not
121+
specified then indices for visiting each index of `x` is returned.
122+
"""
123+
@inline indices(x) = eachindex(x)
124+
125+
indices(x, d) = indices(axes(x, d))
126+
127+
@inline function indices(x::NTuple{N,<:Any}, dim) where {N}
128+
inds = map(x_i -> indices(x_i, dim), x)
129+
@assert all(isequal(first(inds)), Base.tail(inds)) "Not all specified axes are equal: $inds"
130+
return reduce(_pick_range, inds)
131+
end
132+
133+
@inline function indices(x::NTuple{N,<:Any}, dim::NTuple{N,<:Any}) where {N}
134+
inds = map(indices, x, dim)
135+
@assert all(isequal(first(inds)), Base.tail(inds)) "Not all specified axes are equal: $inds"
136+
return reduce(_pick_range, inds)
137+
end
138+
139+

0 commit comments

Comments
 (0)