Skip to content

Commit 2883ce0

Browse files
jishnubjishnub
authored
Logical indexing for IdOffsetRange (#224)
* Logical indexing for IdOffsetRange * Fix tests on v1.0 * simplify logical indexing * remove branches in logical indexing * Add comment * Helper functions for type-check and getindex * move type-checking functions out * preserve axes if all indices are true * don't preserve indices
1 parent 39427dd commit 2883ce0

File tree

2 files changed

+179
-16
lines changed

2 files changed

+179
-16
lines changed

‎src/axes.jl‎

Lines changed: 55 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -77,17 +77,30 @@ struct IdOffsetRange{T<:Integer,I<:AbstractUnitRange{T}} <: AbstractUnitRange{T}
7777
parent::I
7878
offset::T
7979

80-
IdOffsetRange{T,I}(r::I, offset::T) where {T<:Integer,I<:AbstractUnitRange{T}} = new{T,I}(r, offset)
80+
function IdOffsetRange{T,I}(r::I, offset::T) where {T<:Integer,I<:AbstractUnitRange{T}}
81+
_bool_check(T, r, offset)
82+
new{T,I}(r, offset)
83+
end
8184

8285
#= This method is necessary to avoid a StackOverflowError in IdOffsetRange{T,I}(r::IdOffsetRange, offset::Integer).
8386
The type signature in that method is more specific than IdOffsetRange{T,I}(r::I, offset::T),
8487
so it ends up calling itself if I <: IdOffsetRange.
8588
=#
8689
function IdOffsetRange{T,IdOffsetRange{T,I}}(r::IdOffsetRange{T,I}, offset::T) where {T<:Integer,I<:AbstractUnitRange{T}}
90+
_bool_check(T, r, offset)
8791
new{T,IdOffsetRange{T,I}}(r, offset)
8892
end
8993
end
9094

95+
function _bool_check(::Type{Bool}, r, offset)
96+
# disallow the construction of IdOffsetRange{Bool, UnitRange{Bool}}(true:true, true)
97+
if offset && (first(r) || last(r))
98+
throw(ArgumentError("values = $r and offset = $offset can not produce a boolean range"))
99+
end
100+
return nothing
101+
end
102+
_bool_check(::Type, r, offset) = nothing
103+
91104
# Construction/coercion from arbitrary AbstractUnitRanges
92105
function IdOffsetRange{T,I}(r::AbstractUnitRange, offset::Integer = 0) where {T<:Integer,I<:AbstractUnitRange{T}}
93106
rc, o = offset_coerce(I, r)
@@ -157,35 +170,61 @@ offset_coerce(::Type{I}, r::AbstractUnitRange) where I<:AbstractUnitRange =
157170
Base.reduced_index(i::IdOffsetRange) = typeof(i)(first(i):first(i))
158171
# Workaround for #92 on Julia < 1.4
159172
Base.reduced_index(i::IdentityUnitRange{<:IdOffsetRange}) = typeof(i)(first(i):first(i))
160-
for f in [:firstindex, :lastindex, :first, :last]
173+
for f in [:firstindex, :lastindex]
161174
@eval @inline Base.$f(r::IdOffsetRange) = $f(r.parent) + r.offset
162175
end
176+
for f in [:first, :last]
177+
# coerce the type to deal with values that get promoted on addition (eg. Bool)
178+
@eval @inline Base.$f(r::IdOffsetRange) = eltype(r)($f(r.parent) + r.offset)
179+
end
163180

164181
@inline function Base.iterate(r::IdOffsetRange)
165182
ret = iterate(r.parent)
166183
ret === nothing && return nothing
167-
return (ret[1] + r.offset, ret[2])
184+
return (eltype(r)(ret[1] + r.offset), ret[2])
168185
end
169186
@inline function Base.iterate(r::IdOffsetRange, i)
170187
ret = iterate(r.parent, i)
171188
ret === nothing && return nothing
172-
return (ret[1] + r.offset, ret[2])
189+
return (eltype(r)(ret[1] + r.offset), ret[2])
173190
end
174191

175192
@inline function Base.getindex(r::IdOffsetRange, i::Integer)
193+
i isa Bool && throw(ArgumentError("invalid index: $i of type Bool"))
176194
@boundscheck checkbounds(r, i)
177-
@inbounds r.parent[i - r.offset] + r.offset
178-
end
179-
@inline function Base.getindex(r::IdOffsetRange, s::AbstractUnitRange{<:Integer})
180-
@boundscheck checkbounds(r, s)
181-
@inbounds pr = r.parent[_subtractoffset(s, r.offset)] .+ r.offset
182-
_indexedby(pr, axes(s))
183-
end
184-
# The following method is required to avoid falling back to getindex(::AbstractUnitRange, ::StepRange{<:Integer})
185-
@inline function Base.getindex(r::IdOffsetRange, s::StepRange{<:Integer})
186-
@boundscheck checkbounds(r, s)
187-
@inbounds rs = r.parent[s .- r.offset] .+ r.offset
188-
return no_offset_view(rs)
195+
@inbounds eltype(r)(r.parent[i - r.offset] + r.offset)
196+
end
197+
198+
# Logical indexing following https://github.com/JuliaLang/julia/pull/31829
199+
#= Helper function to perform logical indxeing for boolean ranges
200+
The code implemented is a branch-free version of the following:
201+
202+
range(first(s) ? first(r) : last(r), length=Int(last(s)))
203+
204+
See https://github.com/JuliaArrays/OffsetArrays.jl/pull/224#discussion_r595635143
205+
206+
Logical indexing does not preserve indices, unlike other forms of vector indexing
207+
=#
208+
@inline function _getindex(r, s::AbstractUnitRange{Bool})
209+
range(first(r) * first(s) + last(r) * !first(s), length=Int(last(s)))
210+
end
211+
@inline function _getindex(r, s::StepRange{Bool})
212+
range(first(r) * first(s) + last(r) * !first(s), step = oneunit(step(s)), length=Int(last(s)))
213+
end
214+
@inline function _getindex(r, s::AbstractUnitRange)
215+
@inbounds rs = r.parent[_subtractoffset(s, r.offset)] .+ r.offset
216+
_indexedby(rs, axes(s))
217+
end
218+
@inline function _getindex(r, s::StepRange)
219+
rs = @inbounds r.parent[s .- r.offset] .+ r.offset
220+
_indexedby(rs, axes(s))
221+
end
222+
223+
for T in [:AbstractUnitRange, :StepRange]
224+
@eval @inline function Base.getindex(r::IdOffsetRange, s::$T{<:Integer})
225+
@boundscheck checkbounds(r, s)
226+
return _getindex(r, s)
227+
end
189228
end
190229

191230
# offset-preserve broadcasting

‎test/runtests.jl‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,130 @@ end
268268
@test OrdinalRange{BigInt,BigInt}(r2) === r2
269269
end
270270
end
271+
272+
@testset "Bool IdOffsetRange (issue #223)" begin
273+
for b1 in [false, true], b2 in [false, true]
274+
r = IdOffsetRange(b1:b2)
275+
@test first(r) === b1
276+
@test last(r) === b2
277+
end
278+
@test_throws ArgumentError IdOffsetRange(true:true, true)
279+
@test_throws ArgumentError IdOffsetRange{Bool,UnitRange{Bool}}(true:true, true)
280+
@test_throws ArgumentError IdOffsetRange{Bool,IdOffsetRange{Bool,UnitRange{Bool}}}(IdOffsetRange(true:true), true)
281+
end
282+
283+
@testset "Logical indexing" begin
284+
@testset "indexing with a single bool" begin
285+
r = IdOffsetRange(1:2)
286+
@test_throws ArgumentError r[true]
287+
@test_throws ArgumentError r[false]
288+
end
289+
@testset "indexing wtih a Bool UnitRange" begin
290+
r = IdOffsetRange(1:0)
291+
292+
@test r[true:false] == 1:0
293+
@test r[true:false] == collect(r)[true:false]
294+
@test_throws BoundsError r[true:true]
295+
@test_throws BoundsError r[false:false]
296+
@test_throws BoundsError r[false:true]
297+
298+
r = IdOffsetRange(1:1)
299+
300+
@test r[true:true] == 1:1
301+
@test r[true:true] == collect(r)[true:true]
302+
303+
@test r[false:false] == 1:0
304+
@test r[false:false] == collect(r)[false:false]
305+
306+
@test_throws BoundsError r[true:false]
307+
@test_throws BoundsError r[false:true]
308+
309+
r = IdOffsetRange(1:2)
310+
311+
@test r[false:true] == 2:2
312+
@test r[false:true] == collect(r)[false:true]
313+
314+
@test_throws BoundsError r[true:true]
315+
@test_throws BoundsError r[true:false]
316+
@test_throws BoundsError r[false:false]
317+
end
318+
@testset "indexing with a Bool IdOffsetRange" begin
319+
# bounds-checking requires the axes of the indices to match that of the array
320+
function testlogicalindexing(r, r2)
321+
r3 = r[r2];
322+
@test no_offset_view(r3) == collect(r)[collect(r2)]
323+
end
324+
325+
r = IdOffsetRange(10:9)
326+
r2 = IdOffsetRange(true:false)
327+
testlogicalindexing(r, r2)
328+
329+
r = IdOffsetRange(10:10)
330+
r2 = IdOffsetRange(false:false)
331+
testlogicalindexing(r, r2)
332+
r2 = IdOffsetRange(true:true)
333+
testlogicalindexing(r, r2)
334+
335+
r = IdOffsetRange(10:10, 1)
336+
r2 = IdOffsetRange(false:false, 1) # effectively true:true with indices 2:2
337+
testlogicalindexing(r, r2)
338+
339+
r = IdOffsetRange(10:11)
340+
r2 = IdOffsetRange(false:true)
341+
testlogicalindexing(r, r2)
342+
end
343+
@testset "indexing wtih a Bool StepRange" begin
344+
r = IdOffsetRange(1:0)
345+
346+
@test r[true:true:false] == 1:1:0
347+
@test_throws BoundsError r[true:true:true]
348+
@test_throws BoundsError r[false:true:false]
349+
@test_throws BoundsError r[false:true:true]
350+
351+
r = IdOffsetRange(1:1)
352+
353+
@test r[true:true:true] == 1:1:1
354+
@test r[true:true:true] == collect(r)[true:true:true]
355+
@test axes(r[true:true:true], 1) == 1:1
356+
357+
@test r[false:true:false] == 1:1:0
358+
@test r[false:true:false] == collect(r)[false:true:false]
359+
360+
# StepRange{Bool,Int}
361+
s = StepRange(true, 1, true)
362+
@test r[s] == 1:1:1
363+
@test r[s] == collect(r)[s]
364+
365+
s = StepRange(true, 2, true)
366+
@test r[s] == 1:1:1
367+
@test r[s] == collect(r)[s]
368+
369+
s = StepRange(false, 1, false)
370+
@test r[s] == 1:1:0
371+
@test r[s] == collect(r)[s]
372+
373+
s = StepRange(false, 2, false)
374+
@test r[s] == 1:1:0
375+
@test r[s] == collect(r)[s]
376+
377+
@test_throws BoundsError r[true:true:false]
378+
@test_throws BoundsError r[false:true:true]
379+
380+
r = IdOffsetRange(1:2)
381+
382+
@test r[false:true:true] == 2:1:2
383+
@test r[false:true:true] == collect(r)[false:true:true]
384+
385+
# StepRange{Bool,Int}
386+
s = StepRange(false, 1, true)
387+
@test r[s] == 2:1:2
388+
@test r[s] == collect(r)[s]
389+
390+
@test_throws BoundsError r[true:true:true]
391+
@test_throws BoundsError r[true:true:false]
392+
@test_throws BoundsError r[false:true:false]
393+
end
394+
end
271395
end
272396

273397
# used in testing the constructor

0 commit comments

Comments
 (0)