diff --git a/NEWS.md b/NEWS.md index ae6880b5a8d0d..056608323df19 100644 --- a/NEWS.md +++ b/NEWS.md @@ -40,6 +40,8 @@ New library features * `@test_throws "some message" triggers_error()` can now be used to check whether the displayed error text contains "some message" regardless of the specific exception type. Regular expressions, lists of strings, and matching functions are also supported. ([#41888) +* A new `CheckIndexStyle` trait allows one to dispatch to efficient + forms of bounds-checking without overloading `Base.checkindex`. ([#42029]) Standard library changes ------------------------ diff --git a/base/abstractarray.jl b/base/abstractarray.jl index 5ded231fa2be2..2a23f55e3daf0 100644 --- a/base/abstractarray.jl +++ b/base/abstractarray.jl @@ -715,13 +715,18 @@ checkindex(::Type{Bool}, inds::AbstractUnitRange, i) = checkindex(::Type{Bool}, inds::AbstractUnitRange, i::Real) = (first(inds) <= i) & (i <= last(inds)) checkindex(::Type{Bool}, inds::AbstractUnitRange, ::Colon) = true checkindex(::Type{Bool}, inds::AbstractUnitRange, ::Slice) = true -function checkindex(::Type{Bool}, inds::AbstractUnitRange, r::AbstractRange) +checkindex(::Type{Bool}, inds::AbstractUnitRange, I::AbstractArray{Bool}) = false +checkindex(::Type{Bool}, inds::AbstractUnitRange, I::AbstractVector{Bool}) = + _checkindex(Bool, inds, CheckIndexStyle(typeof(I)), I) +checkindex(::Type{Bool}, inds::AbstractUnitRange, I::AbstractArray) = + _checkindex(Bool, inds, CheckIndexStyle(typeof(I)), I) + +function _checkindex(::Type{Bool}, inds::AbstractUnitRange, ::CheckIndexFirstLast, r) @_propagate_inbounds_meta isempty(r) | (checkindex(Bool, inds, first(r)) & checkindex(Bool, inds, last(r))) end -checkindex(::Type{Bool}, indx::AbstractUnitRange, I::AbstractVector{Bool}) = indx == axes1(I) -checkindex(::Type{Bool}, indx::AbstractUnitRange, I::AbstractArray{Bool}) = false -function checkindex(::Type{Bool}, inds::AbstractUnitRange, I::AbstractArray) +_checkindex(::Type{Bool}, inds::AbstractUnitRange, ::CheckIndexAxes, I) = inds == axes1(I) # TODO: comparison should be values-only +function _checkindex(::Type{Bool}, inds::AbstractUnitRange, ::CheckIndexAll, I) @inline b = true for i in I diff --git a/base/indices.jl b/base/indices.jl index 28028f23c72a3..b91465325f36e 100644 --- a/base/indices.jl +++ b/base/indices.jl @@ -102,6 +102,89 @@ IndexStyle(A::AbstractArray, B::AbstractArray...) = IndexStyle(IndexStyle(A), In IndexStyle(::IndexLinear, ::IndexLinear) = IndexLinear() IndexStyle(::IndexStyle, ::IndexStyle) = IndexCartesian() + +abstract type CheckIndexStyle end +""" + Base.CheckIndexFirstLast() + +A [`Base.CheckIndexStyle`](@ref) trait-value indicating that bounds-checking +only need test the first and last elements in an index vector. + +# Examples +```jldoctest +julia> r = 3:7; Base.CheckIndexStyle(r) +Base.CheckIndexFirstLast() +``` + +Ranges are declared `CheckIndexFirstLast` because `x[r]` can be tested +for out-of-bounds indexing using just the first and last elements of `r`. + +See also [`Base.CheckIndexAll`](@ref) and [`Base.CheckIndexAxes`](@ref). +""" +struct CheckIndexFirstLast <: CheckIndexStyle end + +""" + Base.CheckIndexAll() + +A [`Base.CheckIndexStyle`](@ref) trait-value indicating that bounds-checking +must test all elements in an index vector. + +# Examples +```jldoctest +julia> idx = [3,4,5,6,7]; Base.CheckIndexStyle(idx) +Base.CheckIndexAll() +``` + +Since the entries in `idx` could be arbitrary, we have to check each +entry for bounds violations. + +See also [`Base.CheckIndexFirstLast`](@ref) and [`Base.CheckIndexAxes`](@ref). +""" +struct CheckIndexAll <: CheckIndexStyle end + +""" + Base.CheckIndexAxes() + +A [`CheckIndexStyle`](@ref) trait-value indicating that bounds-checking +should consider the axes of the index rather than the values of the +index. This is used in cases where the index acts as a filter to +select elements. + +# Examples +```jldoctest +julia> idx = [true, false, true]; Base.CheckIndexStyle(idx) +Base.CheckIndexAxes() +``` + +When `idx` is used in `x[idx]`, it returns the entries in `x` +corresponding to `true` entries in `idx`. Consequently, indexing +should insist on `idx` and `x` having identical axes. + +See also [`Base.CheckIndexFirstLast`](@ref) and [`Base.CheckIndexAll`](@ref). +""" +struct CheckIndexAxes <: CheckIndexStyle end + +""" + CheckIndexStyle(idx) + CheckIndexStyle(typeof(idx)) + +`CheckIndexStyle` specifies how bounds-checking of `x[idx]` should be performed. +Certain types of `idx`, such as ranges, may have particularly efficient ways +to perform the bounds-checking. When you define a new [`AbstractArray`](@ref) type, +you can choose to define a specific value for this trait: + + Base.CheckIndexStyle(::Type{<:MyRange}) = CheckIndexFirstLast() + +The default is [`CheckIndexAll()`](@ref), except for `AbstractVector`s with `Bool` +`eltype` (which default to [`Base.CheckIndexAxes()`](@ref)) and subtypes of `AbstractRange` +(which default to [`Base.CheckIndexFirstLast()`](@ref).) +""" +CheckIndexStyle(A::AbstractArray) = CheckIndexStyle(typeof(A)) +CheckIndexStyle(::Type{Union{}}) = CheckIndexAll() +CheckIndexStyle(::Type{<:AbstractArray{T}}) where T = T === Bool ? CheckIndexAxes() : CheckIndexAll() +CheckIndexStyle(::Type{<:AbstractRange{T}}) where T = T === Bool ? CheckIndexAxes() : CheckIndexFirstLast() +CheckIndexStyle(::Type{<:FastContiguousSubArray{T,N,P}}) where {T, N, P} = CheckIndexStyle(P) + # array shape rules promote_shape(::Tuple{}, ::Tuple{}) = () diff --git a/doc/src/base/arrays.md b/doc/src/base/arrays.md index 1dc2d8ed926af..9d1df79873775 100644 --- a/doc/src/base/arrays.md +++ b/doc/src/base/arrays.md @@ -57,6 +57,10 @@ Base.eachindex Base.IndexStyle Base.IndexLinear Base.IndexCartesian +Base.CheckIndexStyle +Base.CheckIndexAll +Base.CheckIndexFirstLast +Base.CheckIndexAxes Base.conj! Base.stride Base.strides diff --git a/test/abstractarray.jl b/test/abstractarray.jl index 32c367a7a50a8..b19734c719a57 100644 --- a/test/abstractarray.jl +++ b/test/abstractarray.jl @@ -22,6 +22,10 @@ A = rand(5,4,3) @test checkbounds(Bool, A, 5, 12) == false @test checkbounds(Bool, A, 1, 13) == false @test checkbounds(Bool, A, 6, 12) == false + + x = [1.0, 2.0] + @test x[false:true] == [2.0] + @test checkindex(Bool, axes(x, 1), false:true) == true end @testset "single CartesianIndex" begin diff --git a/test/offsetarray.jl b/test/offsetarray.jl index 7621e14013627..0d0fd90785cd1 100644 --- a/test/offsetarray.jl +++ b/test/offsetarray.jl @@ -277,6 +277,13 @@ v = view(A0, i1, 1) v = view(A0, 1:1, i1) @test axes(v) === (Base.OneTo(1), OffsetArrays.IdOffsetRange(Base.OneTo(2), -5)) +@test checkindex(Bool, 1:2, OffsetArray(1:2, (0,))) +@test checkindex(Bool, 1:2, OffsetArray([1,2], (0,))) +@test checkindex(Bool, 1:2, OffsetArray(false:true, (0,))) +@test !checkindex(Bool, 1:2, OffsetArray(1:3, (0,))) +@test !checkindex(Bool, 1:2, OffsetArray([1,3], (0,))) +@test !checkindex(Bool, 1:3, OffsetArray(false:true, (0,))) + # copyto! and fill! a = OffsetArray{Int}(undef, (-3:-1,)) fill!(a, -1) diff --git a/test/testhelpers/OffsetArrays.jl b/test/testhelpers/OffsetArrays.jl index 27c666c9dacbd..eac05a490edbe 100644 --- a/test/testhelpers/OffsetArrays.jl +++ b/test/testhelpers/OffsetArrays.jl @@ -364,6 +364,8 @@ Broadcast.broadcast_unalias(dest::OffsetArray, src::OffsetArray) = parent(dest) const OffsetRange{T} = OffsetArray{T,1,<:AbstractRange{T}} const IIUR = IdentityUnitRange{S} where S<:AbstractUnitRange{T} where T<:Integer +Base.CheckIndexStyle(::Type{<:OffsetRange{T}}) where T = T === Bool ? Base.CheckIndexAxes() : Base.CheckIndexFirstLast() + Base.step(a::OffsetRange) = step(parent(a)) @propagate_inbounds Base.getindex(a::OffsetRange, r::OffsetRange) = OffsetArray(a[parent(r)], r.offsets)