Skip to content

Commit a0790a2

Browse files
jishnubjishnub
authored
Fix indexing with OffsetRanges (#205)
* Fix indexing with OffsetRanges * Fix indexing IdOffsetRange * Add tests * Bugfix in indexing IdOffsetRange * bugfix in indexing UnitRanges * Don't specialize getindex with colon * remove comments in tests
1 parent b190e15 commit a0790a2

File tree

4 files changed

+204
-16
lines changed

4 files changed

+204
-16
lines changed

‎src/OffsetArrays.jl

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -342,14 +342,14 @@ const IIUR = IdentityUnitRange{S} where S<:AbstractUnitRange{T} where T<:Integer
342342

343343
Base.step(a::OffsetRange) = step(parent(a))
344344

345-
@propagate_inbounds Base.getindex(a::OffsetRange, r::OffsetRange) = OffsetArray(a[parent(r)], r.offsets)
345+
@propagate_inbounds function Base.getindex(a::OffsetRange, r::OffsetRange)
346+
OffsetArray(a.parent[r.parent .- a.offsets[1]], axes(r))
347+
end
346348
@propagate_inbounds function Base.getindex(a::OffsetRange, r::IdOffsetRange)
347-
OffsetArray(a.parent[r.parent .+ (r.offset - a.offsets[1])], r.offset)
349+
OffsetArray(a.parent[r.parent .+ (r.offset - a.offsets[1])], axes(r))
348350
end
349-
@propagate_inbounds Base.getindex(r::OffsetRange, s::IIUR) =
350-
OffsetArray(r[s.indices], s)
351-
@propagate_inbounds Base.getindex(a::OffsetRange, r::AbstractRange) = a.parent[r .- a.offsets[1]]
352-
@propagate_inbounds Base.getindex(a::AbstractRange, r::OffsetRange) = OffsetArray(a[parent(r)], r.offsets)
351+
@propagate_inbounds Base.getindex(a::OffsetRange, r::AbstractRange) = _maybewrapaxes(a.parent[r .- a.offsets[1]], axes(r,1))
352+
@propagate_inbounds Base.getindex(a::AbstractRange, r::OffsetRange) = OffsetArray(a[parent(r)], axes(r))
353353

354354
for OR in [:IIUR, :IdOffsetRange]
355355
for R in [:StepRange, :StepRangeLen, :LinRange, :UnitRange]
@@ -359,12 +359,18 @@ for OR in [:IIUR, :IdOffsetRange]
359359
# this method is needed for ambiguity resolution
360360
@eval @propagate_inbounds Base.getindex(r::StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision}, s::$OR) where T =
361361
OffsetArray(r[no_offset_view(s)], axes(s))
362+
363+
#= Integer UnitRanges may return an appropriate AbstractUnitRange{<:Integer}, as the result may be used in indexing, and
364+
indexing is faster with ranges =#
365+
@eval @propagate_inbounds function Base.getindex(r::UnitRange{<:Integer}, s::$OR)
366+
rs = r[no_offset_view(s)]
367+
offset_s = first(axes(s,1)) - 1
368+
IdOffsetRange(UnitRange(rs .- offset_s), offset_s)
369+
end
362370
end
363371

364-
#= Integer UnitRanges may return an appropriate AbstractUnitRange{<:Integer}, as the result may be used in indexing, and
365-
indexing is faster with ranges =#
366-
@propagate_inbounds Base.getindex(r::UnitRange{<:Integer}, s::IdOffsetRange) = IdOffsetRange(r[no_offset_view(s)] .- s.offset, s.offset)
367-
@propagate_inbounds Base.getindex(r::UnitRange{<:Integer}, s::IIUR) = IdentityUnitRange(r[no_offset_view(s)])
372+
# This is technically breaking, so it might be incorporated in the next major release
373+
# Base.getindex(a::OffsetRange, ::Colon) = OffsetArray(a.parent[:], a.offsets)
368374

369375
function Base.show(io::IO, r::OffsetRange)
370376
show(io, r.parent)

‎src/axes.jl

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -157,13 +157,14 @@ end
157157

158158
@propagate_inbounds Base.getindex(r::IdOffsetRange, i::Integer) = r.parent[i - r.offset] + r.offset
159159
@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::AbstractUnitRange{<:Integer})
160-
return r.parent[s .- r.offset] .+ r.offset
160+
offset_s = first(axes(s,1)) - 1
161+
pr = r.parent[s .- r.offset] .+ (r.offset - offset_s)
162+
_maybewrapoffset(pr, offset_s, axes(s,1))
161163
end
162-
@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::IdentityUnitRange)
163-
return IdOffsetRange(r.parent[s .- r.offset], r.offset)
164-
end
165-
@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::IdOffsetRange)
166-
return IdOffsetRange(r.parent[s.parent .+ (s.offset - r.offset)] .+ (r.offset - s.offset), s.offset)
164+
# The following method is required to avoid falling back to getindex(::AbstractUnitRange, ::StepRange{<:Integer})
165+
@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::StepRange{<:Integer})
166+
rs = r.parent[s .- r.offset] .+ r.offset
167+
return no_offset_view(rs)
167168
end
168169

169170
# offset-preserve broadcasting

‎src/utils.jl

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,3 +71,11 @@ function _checkindices(N::Integer, indices, label)
7171
throw_argumenterror(N, indices, label) = throw(ArgumentError(label*" $indices are not compatible with a $(N)D array"))
7272
N == length(indices) || throw_argumenterror(N, indices, label)
7373
end
74+
75+
_maybewrapaxes(A::AbstractVector, ::Base.OneTo) = no_offset_view(A)
76+
_maybewrapaxes(A::AbstractVector, ax) = OffsetArray(A, ax)
77+
78+
_maybewrapoffset(r::AbstractUnitRange, of, ::Base.OneTo) = no_offset_view(r)
79+
_maybewrapoffset(r::AbstractVector, of, ::Base.OneTo) = no_offset_view(r)
80+
_maybewrapoffset(r::AbstractUnitRange, of, ::Any) = IdOffsetRange(UnitRange(r), of)
81+
_maybewrapoffset(r::AbstractVector, of, axs) = OffsetArray(r .+ of, axs)

‎test/runtests.jl

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,56 @@ struct TupleOfRanges{N}
2424
x ::NTuple{N, UnitRange{Int}}
2525
end
2626

27+
# Useful for testing indexing
28+
struct ZeroBasedRange{T,A<:AbstractRange{T}} <: AbstractRange{T}
29+
a :: A
30+
function ZeroBasedRange(a::AbstractRange{T}) where {T}
31+
@assert !Base.has_offset_axes(a)
32+
new{T, typeof(a)}(a)
33+
end
34+
end
35+
36+
struct ZeroBasedUnitRange{T,A<:AbstractUnitRange{T}} <: AbstractUnitRange{T}
37+
a :: A
38+
function ZeroBasedUnitRange(a::AbstractUnitRange{T}) where {T}
39+
@assert !Base.has_offset_axes(a)
40+
new{T, typeof(a)}(a)
41+
end
42+
end
43+
44+
for Z in [:ZeroBasedRange, :ZeroBasedUnitRange]
45+
@eval Base.parent(A::$Z) = A.a
46+
@eval Base.first(A::$Z) = first(A.a)
47+
@eval Base.length(A::$Z) = length(A.a)
48+
@eval Base.last(A::$Z) = last(A.a)
49+
@eval Base.size(A::$Z) = size(A.a)
50+
@eval Base.axes(A::$Z) = map(x -> 0:x-1, size(A.a))
51+
@eval Base.getindex(A::$Z, i::Int) = A.a[i + 1]
52+
@eval Base.step(A::$Z) = step(A.a)
53+
@eval OffsetArrays.no_offset_view(A::$Z) = A.a
54+
@eval function Base.show(io::IO, A::$Z)
55+
show(io, A.a)
56+
print(io, " with indices $(axes(A,1))")
57+
end
58+
59+
for R in [:AbstractRange, :AbstractUnitRange, :StepRange]
60+
@eval @inline function Base.getindex(A::$Z, r::$R{<:Integer})
61+
@boundscheck checkbounds(A, r)
62+
OffsetArray(A.a[r .+ 1], axes(r))
63+
end
64+
end
65+
for R in [:UnitRange, :StepRange, :StepRangeLen, :LinRange]
66+
@eval @inline function Base.getindex(A::$R, r::$Z)
67+
@boundscheck checkbounds(A, r)
68+
OffsetArray(A[r.a], axes(r))
69+
end
70+
end
71+
@eval @inline function Base.getindex(A::StepRangeLen{<:Any,<:Base.TwicePrecision,<:Base.TwicePrecision}, r::$Z)
72+
@boundscheck checkbounds(A, r)
73+
OffsetArray(A[r.a], axes(r))
74+
end
75+
end
76+
2777
function same_value(r1, r2)
2878
length(r1) == length(r2) || return false
2979
for (v1, v2) in zip(r1, r2)
@@ -722,6 +772,24 @@ end
722772
end
723773
end
724774

775+
_comp(::Type{<:Integer}) = ==
776+
_comp(::Type{<:Real}) = ≈
777+
function test_indexing_axes_and_vals(r1, r2)
778+
r12 = r1[r2]
779+
op = _comp(eltype(r1))
780+
if axes(r12, 1) != axes(r2, 1)
781+
@show r1 r2 r12 axes(r12, 1) axes(r2, 1)
782+
end
783+
@test axes(r12, 1) == axes(r2, 1)
784+
if axes(r12, 1) == axes(r2, 1)
785+
@test op(first(r12), r1[first(r2)])
786+
@test op(last(r12), r1[last(r2)])
787+
for i in eachindex(r2)
788+
@test op(r12[i], r1[r2[i]])
789+
end
790+
end
791+
end
792+
725793
@testset "Vector indexing" begin
726794
A0 = [1 3; 2 4]
727795
A = OffsetArray(A0, (-1,2))
@@ -740,6 +808,55 @@ end
740808
@test A[0, 3:4] == S[0, 3:4] == [1,3]
741809
@test A[1, [4,3]] == S[1, [4,3]] == [4,2]
742810
@test A[:, :] == S[:, :] == A
811+
812+
r1 = OffsetArray(IdentityUnitRange(100:1000), 3)
813+
r2 = r1[:]
814+
@test r2 == r1
815+
816+
for r1 in [
817+
# AbstractArrays
818+
OffsetArray(10:1000, 0), # 1-based index
819+
OffsetArray(10:3:1000, 3), # offset index
820+
OffsetArray(10.0:3:1000.0, 0), # 1-based index
821+
OffsetArray(10.0:3:1000.0, 3), # offset index
822+
OffsetArray(IdOffsetRange(10:1000, 1), -1), # 1-based index
823+
OffsetArray(IdOffsetRange(10:1000, 1), 3), # offset index
824+
OffsetArray(IdOffsetRange(IdOffsetRange(10:1000, -4), 1), 3), # 1-based index
825+
OffsetArray(IdOffsetRange(IdOffsetRange(10:1000, -1), 1), 3), # offset index
826+
827+
# AbstractRanges
828+
1:1000,
829+
1:3:1000,
830+
1.0:3.0:1000.0,
831+
IdOffsetRange(ZeroBasedUnitRange(1:1000), 1), # 1-based index
832+
IdOffsetRange(ZeroBasedUnitRange(1:1000), 2), # offset index
833+
ZeroBasedUnitRange(1:1000), # offset range
834+
ZeroBasedRange(1:1000), # offset range
835+
ZeroBasedRange(1:1:1000), # offset range
836+
]
837+
838+
# AbstractArrays with 1-based indices
839+
for r2 in [
840+
OffsetArray(5:80, 0),
841+
OffsetArray(5:2:80, 0),
842+
OffsetArray(IdentityUnitRange(5:80), -4),
843+
OffsetArray(IdOffsetRange(5:80), 0),
844+
]
845+
846+
test_indexing_axes_and_vals(r1, r2)
847+
end
848+
849+
# AbstractRanges with 1-based indices
850+
for r2 in [
851+
5:80,
852+
5:2:80,
853+
IdOffsetRange(5:80),
854+
IdOffsetRange(ZeroBasedUnitRange(4:79), 1),
855+
]
856+
857+
test_indexing_axes_and_vals(r1, r2)
858+
end
859+
end
743860
end
744861

745862
@testset "Vector indexing with offset ranges" begin
@@ -765,6 +882,62 @@ end
765882
for i in axes(ax,1)
766883
@test a[ax[i]] == a[ax][i]
767884
end
885+
886+
for r1 in [
887+
# AbstractArrays
888+
OffsetArray(10:1000, 0), # 1-based index
889+
OffsetArray(10:1000, 3), # offset index
890+
OffsetArray(10:3:1000, 0), # 1-based index
891+
OffsetArray(10:3:1000, 3), # offset index
892+
OffsetArray(10.0:3:1000.0, 0), # 1-based index
893+
OffsetArray(10.0:3:1000.0, 3), # offset index
894+
OffsetArray(IdOffsetRange(10:1000, -3), 3), # 1-based index
895+
OffsetArray(IdOffsetRange(10:1000, 1), 3), # offset index
896+
OffsetArray(IdOffsetRange(IdOffsetRange(10:1000, -4), 1), 3), # 1-based index
897+
OffsetArray(IdOffsetRange(IdOffsetRange(10:1000, -1), 1), 3), # offset index
898+
899+
# AbstractRanges
900+
1:1000,
901+
1:2:2000,
902+
1.0:2.0:2000.0,
903+
LinRange(1.0, 2000.0, 2000),
904+
IdOffsetRange(1:1000, 0), # 1-based index
905+
IdOffsetRange(ZeroBasedUnitRange(1:1000), 1), # 1-based index
906+
IdOffsetRange(ZeroBasedUnitRange(1:1000), 2), # offset index
907+
IdentityUnitRange(ZeroBasedUnitRange(1:1000)), # 1-based index
908+
ZeroBasedUnitRange(1:1000), # offset index
909+
ZeroBasedRange(1:1000), # offset index
910+
ZeroBasedRange(1:1:1000), # offset index
911+
ZeroBasedUnitRange(IdentityUnitRange(1:1000)), # offset index
912+
]
913+
914+
# AbstractArrays with offset axes
915+
for r2 in [OffsetArray(5:80, 40), OffsetArray(5:2:80, 40),
916+
OffsetArray(IdentityUnitRange(5:80), 2),
917+
OffsetArray(IdOffsetRange(5:80, 1), 3),
918+
OffsetArray(IdOffsetRange(IdOffsetRange(5:80, 4), 1), 3),
919+
OffsetArray(IdOffsetRange(IdentityUnitRange(5:80), 1), 3),
920+
OffsetArray(IdentityUnitRange(IdOffsetRange(5:80, 1)), 3),
921+
]
922+
923+
test_indexing_axes_and_vals(r1, r2)
924+
end
925+
926+
# AbstractRanges with offset axes
927+
for r2 in [IdOffsetRange(5:80, 1),
928+
IdentityUnitRange(5:80),
929+
IdOffsetRange(IdOffsetRange(5:80, 2), 1),
930+
IdOffsetRange(IdOffsetRange(IdOffsetRange(5:80, -1), 2), 1),
931+
IdentityUnitRange(IdOffsetRange(1:10, 5)),
932+
IdOffsetRange(IdentityUnitRange(15:20), -2),
933+
ZeroBasedUnitRange(5:80),
934+
ZeroBasedRange(5:80),
935+
ZeroBasedRange(5:2:80),
936+
]
937+
938+
test_indexing_axes_and_vals(r1, r2)
939+
end
940+
end
768941
end
769942

770943
@testset "CartesianIndexing" begin

0 commit comments

Comments
 (0)