Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add methods to index identityUnitRange/Slice with another IdentityUnitRange #41224

Merged
merged 1 commit into from
Mar 9, 2024
Merged

Add methods to index identityUnitRange/Slice with another IdentityUnitRange #41224

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 46 additions & 4 deletions base/indices.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -423,15 +423,57 @@ first(S::IdentityUnitRange) = first(S.indices)
last(S::IdentityUnitRange) = last(S.indices)
size(S::IdentityUnitRange) = (length(S.indices),)
length(S::IdentityUnitRange) = length(S.indices)
getindex(S::IdentityUnitRange, i::Int) = (@inline; @boundscheck checkbounds(S, i); i)
getindex(S::IdentityUnitRange, i::AbstractUnitRange{<:Integer}) = (@inline; @boundscheck checkbounds(S, i); i)
getindex(S::IdentityUnitRange, i::StepRange{<:Integer}) = (@inline; @boundscheck checkbounds(S, i); i)
unsafe_length(S::IdentityUnitRange) = unsafe_length(S.indices)
getindex(S::IdentityUnitRange, i::Integer) = (@inline; @boundscheck checkbounds(S, i); convert(eltype(S), i))
getindex(S::IdentityUnitRange, i::Bool) = throw(ArgumentError("invalid index: $i of type Bool"))
function getindex(S::IdentityUnitRange, i::AbstractUnitRange{<:Integer})
@inline
@boundscheck checkbounds(S, i)
return convert(AbstractUnitRange{eltype(S)}, i)
end
function getindex(S::IdentityUnitRange, i::AbstractUnitRange{Bool})
@inline
@boundscheck checkbounds(S, i)
range(first(i) ? first(S) : last(S), length = last(i))
end
function getindex(S::IdentityUnitRange, i::StepRange{<:Integer})
@inline
@boundscheck checkbounds(S, i)
return convert(AbstractRange{eltype(S)}, i)
end
function getindex(S::IdentityUnitRange, i::StepRange{Bool})
@inline
@boundscheck checkbounds(S, i)
range(first(i) ? first(S) : last(S), length = last(i), step = Int(step(i)))
end
# Indexing with offset ranges should preserve the axes of the indices
# however, this is only really possible in general with OffsetArrays.
# In some cases, though, we may obtain correct results using Base ranges
# the following methods are added to allow OffsetArrays to dispatch on the first argument without ambiguities
function getindex(S::IdentityUnitRange{<:AbstractUnitRange{<:Integer}},
i::IdentityUnitRange{<:AbstractUnitRange{<:Integer}})
@inline
@boundscheck checkbounds(S, i)
return i
end
function getindex(S::Slice{<:AbstractUnitRange{<:Integer}},
i::IdentityUnitRange{<:AbstractUnitRange{<:Integer}})
@inline
@boundscheck checkbounds(S, i)
return i
end
show(io::IO, r::IdentityUnitRange) = print(io, "Base.IdentityUnitRange(", r.indices, ")")
iterate(S::IdentityUnitRange, s...) = iterate(S.indices, s...)

# For OneTo, the values and indices of the values are identical, so this may be defined in Base.
# In general such an indexing operation would produce offset ranges
getindex(S::OneTo, I::IdentityUnitRange{<:AbstractUnitRange{<:Integer}}) = (@inline; @boundscheck checkbounds(S, I); I)
# This should also ideally return an AbstractUnitRange{eltype(S)}, but currently
# we're restricted to eltype(::IdentityUnitRange) == Int by definition
function getindex(S::OneTo, I::IdentityUnitRange{<:AbstractUnitRange{<:Integer}})
@inline
@boundscheck checkbounds(S, I)
return I
end

"""
LinearIndices(A::AbstractArray)
Expand Down
47 changes: 40 additions & 7 deletions test/ranges.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2376,13 +2376,46 @@ end
@test 0.2 * (-2:2:2) == [-0.4, 0, 0.4]
end

@testset "Indexing OneTo with IdentityUnitRange" begin
for endpt in Any[10, big(10), UInt(10)]
r = Base.OneTo(endpt)
inds = Base.IdentityUnitRange(3:5)
rs = r[inds]
@test rs === inds
@test_throws BoundsError r[Base.IdentityUnitRange(-1:100)]
@testset "IdentityUnitRange indexing" begin
@testset "Indexing into an IdentityUnitRange" begin
@testset for r in Any[-1:20, Base.OneTo(20)]
ri = Base.IdentityUnitRange(r)
@test_throws "invalid index" ri[true]
@testset for s in Any[Base.OneTo(6), Base.OneTo{BigInt}(6), 3:6, big(3):big(6), 3:2:7]
@test mapreduce(==, &, ri[s], ri[s[begin]]:step(s):ri[s[end]])
@test axes(ri[s]) == axes(s)
@test eltype(ri[s]) == eltype(ri)
end
end
@testset "Bool indices" begin
r = 1:1
@test Base.IdentityUnitRange(r)[true:true] == r[true:true]
@test Base.IdentityUnitRange(r)[true:true:true] == r[true:true:true]
@test_throws BoundsError Base.IdentityUnitRange(1:2)[true:true]
@test_throws BoundsError Base.IdentityUnitRange(1:2)[true:true:true]
end
end
@testset "Indexing with IdentityUnitRange" begin
@testset "OneTo" begin
@testset for endpt in Any[10, big(12), UInt(11)]
r = Base.OneTo(endpt)
inds = Base.IdentityUnitRange(3:5)
rs = r[inds]
@test rs == inds
@test axes(rs) == axes(inds)
@test_throws BoundsError r[Base.IdentityUnitRange(-1:100)]
end
end
@testset "IdentityUnitRange" begin
@testset for r in Any[Base.IdentityUnitRange(1:4), Base.IdentityUnitRange(Base.OneTo(4)), Base.Slice(1:4), Base.Slice(Base.OneTo(4))]
@testset for s in Any[Base.IdentityUnitRange(3:3), Base.IdentityUnitRange(Base.OneTo(2)), Base.Slice(3:3), Base.Slice(Base.OneTo(2))]
rs = r[s]
@test rs == s
@test axes(rs) == axes(s)
end
@test_throws BoundsError r[Base.IdentityUnitRange(first(r):last(r) + 1)]
end
end
end
end

Expand Down