Skip to content

Add CheckIndexStyle trait for bounds-checking #42029

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.

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

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
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
2 changes: 2 additions & 0 deletions NEWS.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`. ([#?????])

Standard library changes
------------------------
Expand Down
13 changes: 9 additions & 4 deletions base/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
82 changes: 82 additions & 0 deletions base/indices.jl
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,88 @@ 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()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we include things like FastContiguousSubArray whose parent types have CheckIndexFirstLast

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure I follow. This checks the indexes, not the array you are indexing. So if we checked only the first and last elements of idx, this would be a problem:

julia> a = [1, 2, 37, 4]
4-element Vector{Int64}:
  1
  2
 37
  4

julia> b = 1:4
1:4

julia> sa = view(a, :)
4-element view(::Vector{Int64}, :) with eltype Int64:
  1
  2
 37
  4

julia> isa(sa, Base.FastContiguousSubArray)
true

julia> b[sa]
ERROR: BoundsError: attempt to access 4-element UnitRange{Int64} at index [[1, 2, 37, 4]]
Stacktrace:
 [1] throw_boundserror(A::UnitRange{Int64}, I::Tuple{SubArray{Int64, 1, Vector{Int64}, Tuple{Base.Slice{Base.OneTo{Int64}}}, true}})
   @ Base ./abstractarray.jl:651
 [2] checkbounds
   @ ./abstractarray.jl:616 [inlined]
 [3] _getindex
   @ ./multidimensional.jl:831 [inlined]
 [4] getindex(A::UnitRange{Int64}, I::SubArray{Int64, 1, Vector{Int64}, Tuple{Base.Slice{Base.OneTo{Int64}}}, true})
   @ Base ./abstractarray.jl:1170
 [5] top-level scope
   @ REPL[6]:1

That said, special rules might apply to AbstractSortedVector but I am not sure we have one of those anywhere.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The parent type would have to return CheckIndexFirstLast, like this view(Wrapper(1:5), :).

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, I think I get it, you mean if the index is a FastContiguousSubArray. Sure, I'll add that.


# array shape rules

promote_shape(::Tuple{}, ::Tuple{}) = ()
Expand Down
4 changes: 4 additions & 0 deletions doc/src/base/arrays.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions test/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 7 additions & 0 deletions test/offsetarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions test/testhelpers/OffsetArrays.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down