From 6f8ba49e157d4b0de3726be8613a6ae607c5ccd2 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Sat, 9 Mar 2024 19:21:03 +0530 Subject: [PATCH] Add methods to index identityUnitRange/Slice with another IdentityUnitRange (#41224) Adding these methods lets `OffsetArrays` define `getindex(::AbstractUnitRange, ::IdentityUnitRange)` without ambiguities. This is in the domain of sanctioned type-piracy, as the result is an offset range in general and cannot be represented correctly using `Base` types. Re: https://github.com/JuliaArrays/OffsetArrays.jl/pull/244 cc: @johnnychen94 Edit: this also fixes an indexing bug in `IdentityUntiRange`: master ```julia julia> r = Base.IdentityUnitRange(-3:3) Base.IdentityUnitRange(-3:3) julia> r[2] 2 julia> r[big(2)] -2 ``` Co-authored-by: jishnub --- base/indices.jl | 50 +++++++++++++++++++++++++++++++++++++++++++++---- test/ranges.jl | 47 +++++++++++++++++++++++++++++++++++++++------- 2 files changed, 86 insertions(+), 11 deletions(-) diff --git a/base/indices.jl b/base/indices.jl index 0bd46f9787dab..9fadff1cdb5de 100644 --- a/base/indices.jl +++ b/base/indices.jl @@ -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) diff --git a/test/ranges.jl b/test/ranges.jl index 4660a96dfc16a..9284c898bcbfb 100644 --- a/test/ranges.jl +++ b/test/ranges.jl @@ -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