Skip to content

Commit 7a43b4c

Browse files
jishnubjishnub
authored
index AbstractUnitRanges with offset ranges (#244)
1 parent e9f33f0 commit 7a43b4c

File tree

5 files changed

+129
-66
lines changed

5 files changed

+129
-66
lines changed

‎Project.toml

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
name = "OffsetArrays"
22
uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
3-
version = "1.9.2"
3+
version = "1.10.0"
44

55
[deps]
66
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

‎src/OffsetArrays.jl

+21-11
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ end
99

1010
export OffsetArray, OffsetMatrix, OffsetVector
1111

12+
const IIUR = IdentityUnitRange{<:AbstractUnitRange{<:Integer}}
13+
1214
include("axes.jl")
1315
include("utils.jl")
1416
include("origin.jl")
@@ -434,10 +436,8 @@ Base.dataids(A::OffsetArray) = Base.dataids(parent(A))
434436
Broadcast.broadcast_unalias(dest::OffsetArray, src::OffsetArray) = parent(dest) === parent(src) ? src : Broadcast.unalias(dest, src)
435437

436438
### Special handling for AbstractRange
437-
438439
const OffsetRange{T} = OffsetVector{T,<:AbstractRange{T}}
439440
const OffsetUnitRange{T} = OffsetVector{T,<:AbstractUnitRange{T}}
440-
const IIUR = IdentityUnitRange{S} where S<:AbstractUnitRange{T} where T<:Integer
441441

442442
Base.step(a::OffsetRange) = step(parent(a))
443443

@@ -501,23 +501,33 @@ end
501501
IdOffsetRange(_subtractoffset(parent(r), of), of)
502502
end
503503

504+
@inline function _boundscheck_index_retaining_axes(r, s)
505+
@boundscheck checkbounds(r, s)
506+
@inbounds pr = r[UnitRange(s)]
507+
_indexedby(pr, axes(s))
508+
end
509+
@inline _boundscheck_return(r, s) = (@boundscheck checkbounds(r, s); s)
510+
504511
for OR in [:IIUR, :IdOffsetRange]
505-
for R in [:StepRange, :StepRangeLen, :LinRange, :UnitRange]
506-
@eval @inline function Base.getindex(r::$R, s::$OR)
507-
@boundscheck checkbounds(r, s)
508-
@inbounds pr = r[UnitRange(s)]
509-
_indexedby(pr, axes(s))
510-
end
512+
for R in [:StepRange, :StepRangeLen, :LinRange, :AbstractUnitRange]
513+
@eval @inline Base.getindex(r::$R, s::$OR) = _boundscheck_index_retaining_axes(r, s)
511514
end
512515

513516
# this method is needed for ambiguity resolution
514517
@eval @inline function Base.getindex(r::StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision}, s::$OR) where T
515-
@boundscheck checkbounds(r, s)
516-
@inbounds pr = r[UnitRange(s)]
517-
_indexedby(pr, axes(s))
518+
_boundscheck_index_retaining_axes(r, s)
518519
end
519520
end
520521

522+
# These methods are added to avoid ambiguities with Base.
523+
# The ones involving Base types should be ported to Base and version-limited here
524+
@inline Base.getindex(r::IdentityUnitRange, s::IIUR) = _boundscheck_return(r, s)
525+
@inline Base.getindex(r::IdentityUnitRange, s::IdOffsetRange) = _boundscheck_return(r, s)
526+
if IdentityUnitRange !== Base.Slice
527+
@inline Base.getindex(r::Base.Slice, s::IIUR) = _boundscheck_return(r, s)
528+
@inline Base.getindex(r::Base.Slice, s::IdOffsetRange) = _boundscheck_return(r, s)
529+
end
530+
521531
# eltype conversion
522532
# This may use specialized map methods for the parent
523533
Base.map(::Type{T}, O::OffsetArray) where {T} = parent_call(x -> map(T, x), O)

‎src/axes.jl

+8
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,14 @@ for T in [:AbstractUnitRange, :StepRange]
227227
end
228228
end
229229

230+
# These methods are necessary to avoid ambiguity
231+
for R in [:IIUR, :IdOffsetRange]
232+
@eval @inline function Base.getindex(r::IdOffsetRange, s::$R)
233+
@boundscheck checkbounds(r, s)
234+
return _getindex(r, s)
235+
end
236+
end
237+
230238
# offset-preserve broadcasting
231239
Broadcast.broadcasted(::Base.Broadcast.DefaultArrayStyle{1}, ::typeof(-), r::IdOffsetRange{T}, x::Integer) where T =
232240
IdOffsetRange{T}(r.parent .- x, r.offset)

‎test/customranges.jl

+82
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
# Useful for testing indexing
2+
struct ZeroBasedRange{T,A<:AbstractRange{T}} <: AbstractRange{T}
3+
a :: A
4+
function ZeroBasedRange(a::AbstractRange{T}) where {T}
5+
@assert !Base.has_offset_axes(a)
6+
new{T, typeof(a)}(a)
7+
end
8+
end
9+
10+
struct ZeroBasedUnitRange{T,A<:AbstractUnitRange{T}} <: AbstractUnitRange{T}
11+
a :: A
12+
function ZeroBasedUnitRange(a::AbstractUnitRange{T}) where {T}
13+
@assert !Base.has_offset_axes(a)
14+
new{T, typeof(a)}(a)
15+
end
16+
end
17+
18+
for Z in [:ZeroBasedRange, :ZeroBasedUnitRange]
19+
@eval Base.parent(A::$Z) = A.a
20+
@eval Base.first(A::$Z) = first(A.a)
21+
@eval Base.length(A::$Z) = length(A.a)
22+
@eval Base.last(A::$Z) = last(A.a)
23+
@eval Base.size(A::$Z) = size(A.a)
24+
@eval Base.axes(A::$Z) = map(x -> IdentityUnitRange(0:x-1), size(A.a))
25+
@eval Base.getindex(A::$Z, i::Int) = A.a[i + 1]
26+
@eval Base.getindex(A::$Z, i::Integer) = A.a[i + 1]
27+
@eval Base.firstindex(A::$Z) = 0
28+
@eval Base.step(A::$Z) = step(A.a)
29+
@eval OffsetArrays.no_offset_view(A::$Z) = A.a
30+
@eval function Base.show(io::IO, A::$Z)
31+
show(io, A.a)
32+
print(io, " with indices $(axes(A,1))")
33+
end
34+
end
35+
36+
for Z in [:ZeroBasedRange, :ZeroBasedUnitRange]
37+
for R in [:AbstractRange, :AbstractUnitRange, :StepRange]
38+
@eval @inline function Base.getindex(A::$Z, r::$R{<:Integer})
39+
@boundscheck checkbounds(A, r)
40+
OffsetArrays._indexedby(A.a[r .+ 1], axes(r))
41+
end
42+
end
43+
44+
for R in [:ZeroBasedUnitRange, :ZeroBasedRange]
45+
@eval @inline function Base.getindex(A::$Z, r::$R{<:Integer})
46+
@boundscheck checkbounds(A, r)
47+
OffsetArrays._indexedby(A.a[r.a .+ 1], axes(r))
48+
end
49+
end
50+
51+
for R in [:IIUR, :IdOffsetRange]
52+
@eval @inline function Base.getindex(A::$Z, r::$R)
53+
@boundscheck checkbounds(A, r)
54+
OffsetArrays._indexedby(A.a[r .+ 1], axes(r))
55+
end
56+
end
57+
58+
for R in [:AbstractUnitRange, :IdOffsetRange, :IdentityUnitRange, :SliceIntUR, :StepRange, :StepRangeLen, :LinRange]
59+
@eval @inline function Base.getindex(A::$R, r::$Z)
60+
@boundscheck checkbounds(A, r)
61+
OffsetArrays._indexedby(A[r.a], axes(r))
62+
end
63+
end
64+
@eval @inline function Base.getindex(A::StepRangeLen{<:Any,<:Base.TwicePrecision,<:Base.TwicePrecision}, r::$Z)
65+
@boundscheck checkbounds(A, r)
66+
OffsetArrays._indexedby(A[r.a], axes(r))
67+
end
68+
end
69+
70+
# A basic range that does not have specialized vector indexing methods defined
71+
# In this case the best that we may do is to return an OffsetArray
72+
# Despite this, an indexing operation involving this type should preserve the axes of the indices
73+
struct CustomRange{T,A<:AbstractRange{T}} <: AbstractRange{T}
74+
a :: A
75+
end
76+
Base.parent(r::CustomRange) = r.a
77+
Base.size(r::CustomRange) = size(parent(r))
78+
Base.axes(r::CustomRange) = axes(parent(r))
79+
Base.first(r::CustomRange) = first(parent(r))
80+
Base.last(r::CustomRange) = last(parent(r))
81+
Base.step(r::CustomRange) = step(parent(r))
82+
Base.getindex(r::CustomRange, i::Int) = getindex(parent(r), i)

‎test/runtests.jl

+17-54
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using OffsetArrays
2-
using OffsetArrays: IdentityUnitRange, no_offset_view
2+
using OffsetArrays: IdentityUnitRange, no_offset_view, IIUR
3+
using Base: Slice
34
using OffsetArrays: IdOffsetRange
45
using Test, Aqua, Documenter
56
using LinearAlgebra
@@ -11,6 +12,8 @@ using StaticArrays
1112
using FillArrays
1213
using DistributedArrays
1314

15+
const SliceIntUR = Slice{<:AbstractUnitRange{<:Integer}}
16+
1417
DocMeta.setdocmeta!(OffsetArrays, :DocTestSetup, :(using OffsetArrays); recursive=true)
1518

1619
# https://github.com/JuliaLang/julia/pull/29440
@@ -26,58 +29,7 @@ struct TupleOfRanges{N}
2629
x ::NTuple{N, UnitRange{Int}}
2730
end
2831

29-
# Useful for testing indexing
30-
struct ZeroBasedRange{T,A<:AbstractRange{T}} <: AbstractRange{T}
31-
a :: A
32-
function ZeroBasedRange(a::AbstractRange{T}) where {T}
33-
@assert !Base.has_offset_axes(a)
34-
new{T, typeof(a)}(a)
35-
end
36-
end
37-
38-
struct ZeroBasedUnitRange{T,A<:AbstractUnitRange{T}} <: AbstractUnitRange{T}
39-
a :: A
40-
function ZeroBasedUnitRange(a::AbstractUnitRange{T}) where {T}
41-
@assert !Base.has_offset_axes(a)
42-
new{T, typeof(a)}(a)
43-
end
44-
end
45-
46-
for Z in [:ZeroBasedRange, :ZeroBasedUnitRange]
47-
@eval Base.parent(A::$Z) = A.a
48-
@eval Base.first(A::$Z) = first(A.a)
49-
@eval Base.length(A::$Z) = length(A.a)
50-
@eval Base.last(A::$Z) = last(A.a)
51-
@eval Base.size(A::$Z) = size(A.a)
52-
@eval Base.axes(A::$Z) = map(x -> IdentityUnitRange(0:x-1), size(A.a))
53-
@eval Base.getindex(A::$Z, i::Int) = A.a[i + 1]
54-
@eval Base.firstindex(A::$Z) = 0
55-
@eval Base.axes(A::$Z) = map(x -> IdentityUnitRange(0:x-1), size(A.a))
56-
@eval Base.getindex(A::$Z, i::Integer) = A.a[i + 1]
57-
@eval Base.step(A::$Z) = step(A.a)
58-
@eval OffsetArrays.no_offset_view(A::$Z) = A.a
59-
@eval function Base.show(io::IO, A::$Z)
60-
show(io, A.a)
61-
print(io, " with indices $(axes(A,1))")
62-
end
63-
64-
for R in [:AbstractRange, :AbstractUnitRange, :StepRange]
65-
@eval @inline function Base.getindex(A::$Z, r::$R{<:Integer})
66-
@boundscheck checkbounds(A, r)
67-
OffsetArrays._indexedby(A.a[r .+ 1], axes(r))
68-
end
69-
end
70-
for R in [:UnitRange, :StepRange, :StepRangeLen, :LinRange]
71-
@eval @inline function Base.getindex(A::$R, r::$Z)
72-
@boundscheck checkbounds(A, r)
73-
OffsetArrays._indexedby(A[r.a], axes(r))
74-
end
75-
end
76-
@eval @inline function Base.getindex(A::StepRangeLen{<:Any,<:Base.TwicePrecision,<:Base.TwicePrecision}, r::$Z)
77-
@boundscheck checkbounds(A, r)
78-
OffsetArrays._indexedby(A[r.a], axes(r))
79-
end
80-
end
32+
include("customranges.jl")
8133

8234
function same_value(r1, r2)
8335
length(r1) == length(r2) || return false
@@ -1128,6 +1080,10 @@ end
11281080
OffsetArray(IdOffsetRange(IdOffsetRange(10:1000, -1), 1), 3), # offset index
11291081

11301082
# AbstractRanges
1083+
Base.OneTo(1000),
1084+
CustomRange(Base.OneTo(1000)),
1085+
Slice(Base.OneTo(1000)),
1086+
SOneTo(1000),
11311087
1:1000,
11321088
UnitRange(1.0, 1000.0),
11331089
1:3:1000,
@@ -1144,6 +1100,7 @@ end
11441100
ZeroBasedUnitRange(1:1000), # offset range
11451101
ZeroBasedRange(1:1000), # offset range
11461102
ZeroBasedRange(1:1:1000), # offset range
1103+
CustomRange(ZeroBasedRange(1:1:1000)), # offset range
11471104
]
11481105

11491106
# AbstractArrays with 1-based indices
@@ -1158,7 +1115,7 @@ end
11581115
test_indexing_axes_and_vals(r1, r2)
11591116
test_indexing_axes_and_vals(r1, collect(r2))
11601117

1161-
if r1 isa AbstractRange && axes(r2, 1) isa Base.OneTo
1118+
if r1 isa AbstractRange && !(r1 isa CustomRange) && axes(r2, 1) isa Base.OneTo
11621119
@test r1[r2] isa AbstractRange
11631120
end
11641121
end
@@ -1269,13 +1226,18 @@ end
12691226
OffsetArray(IdOffsetRange(IdOffsetRange(10:1000, -1), 1), 3), # offset index
12701227

12711228
# AbstractRanges
1229+
Base.OneTo(1000),
1230+
Slice(Base.OneTo(1000)),
1231+
SOneTo(1000),
1232+
CustomRange(Base.OneTo(1000)),
12721233
1:1000,
12731234
UnitRange(1.0, 1000.0),
12741235
1:2:2000,
12751236
2000:-1:1,
12761237
1.0:2.0:2000.0,
12771238
StepRangeLen(Float64(1), Float64(1000), 1000),
12781239
LinRange(1.0, 2000.0, 2000),
1240+
Base.Slice(Base.OneTo(1000)), # 1-based index
12791241
IdOffsetRange(Base.OneTo(1000)), # 1-based index
12801242
IdOffsetRange(1:1000, 0), # 1-based index
12811243
IdOffsetRange(Base.OneTo(1000), 4), # offset index
@@ -1288,6 +1250,7 @@ end
12881250
ZeroBasedRange(1:1000), # offset index
12891251
ZeroBasedRange(1:1:1000), # offset index
12901252
ZeroBasedUnitRange(IdentityUnitRange(1:1000)), # offset index
1253+
CustomRange(ZeroBasedUnitRange(IdentityUnitRange(1:1000))), # offset index
12911254
]
12921255

12931256
# AbstractArrays with offset axes

0 commit comments

Comments
 (0)