Fix inconsistant logical index behavior (JuliaLang#45869)

1. If we use `BitArray`/`Array{Bool}` to index, `to_indices` has an
optimiztion for linear-iteratable case.
But the corresponding check is not correct. IIUC, this optimization is
legal only when the Boolen array is the only index provided.
The first commit fix it and widen this optimization to all Boolen array.
Before this PR
```julia
julia> A = rand(2,3,4); I = rand(Bool,3,4);

julia> A[1,I] == A[1,view(I,:,:)]
ERROR: BoundsError: attempt to access 2×3×4 Array{Float64, 3} at index [1, 3×4 Matrix{Bool}]
```
After
```julia
julia> A = rand(2,3,4); I = rand(Bool,3,4);

julia> A[1,I] == A[1,view(I,:,:)]
true
```

2. On master, if the index/array has singleton trailing dimension,
boundcheck of logical index show different behavior depending on the
number of indexes provided.
If there's only one index variable, singleton dimension wil not be
ignored. The second commit fix it. (close JuliaLang#45867)

---------

Co-authored-by: Matt Bauman <[email protected]>

base/abstractarray.jl

@@ -678,15 +678,12 @@ function checkbounds(::Type{Bool}, A::AbstractArray, I...)
  checkbounds_indices(Bool, axes(A), I)
 end
 
-# Linear indexing is explicitly allowed when there is only one (non-cartesian) index
+# Linear indexing is explicitly allowed when there is only one (non-cartesian) index;
+# indices that do not allow linear indexing (e.g., logical arrays, cartesian indices, etc)
+# must add specialized methods to implement their restrictions
 function checkbounds(::Type{Bool}, A::AbstractArray, i)
  @inline
- checkindex(Bool, eachindex(IndexLinear(), A), i)
-end
-# As a special extension, allow using logical arrays that match the source array exactly
-function checkbounds(::Type{Bool}, A::AbstractArray{<:Any,N}, I::AbstractArray{Bool,N}) where N
- @inline
- axes(A) == axes(I)
+ return checkindex(Bool, eachindex(IndexLinear(), A), i)
 end
 
 """
@@ -720,16 +717,13 @@ of `IA`.
 
 See also [`checkbounds`](@ref).
 """
-function checkbounds_indices(::Type{Bool}, IA::Tuple, I::Tuple)
- @inline
- checkindex(Bool, IA[1], I[1])::Bool & checkbounds_indices(Bool, tail(IA), tail(I))
-end
-function checkbounds_indices(::Type{Bool}, ::Tuple{}, I::Tuple)
+function checkbounds_indices(::Type{Bool}, inds::Tuple, I::Tuple{Any, Vararg})
  @inline
- checkindex(Bool, OneTo(1), I[1])::Bool & checkbounds_indices(Bool, (), tail(I))
+ return checkindex(Bool, get(inds, 1, OneTo(1)), I[1])::Bool &
+ checkbounds_indices(Bool, safe_tail(inds), tail(I))
 end
-checkbounds_indices(::Type{Bool}, IA::Tuple, ::Tuple{}) = (@inline; all(x->length(x)==1, IA))
-checkbounds_indices(::Type{Bool}, ::Tuple{}, ::Tuple{}) = true
+
+checkbounds_indices(::Type{Bool}, inds::Tuple, ::Tuple{}) = (@inline; all(x->length(x)==1, inds))
 
 # check along a single dimension
 """
@@ -751,20 +745,19 @@ julia> checkindex(Bool, 1:20, 21)
 false
 ```
 """
-checkindex(::Type{Bool}, inds::AbstractUnitRange, i) =
- throw(ArgumentError("unable to check bounds for indices of type $(typeof(i))"))
+checkindex(::Type{Bool}, inds, i) = throw(ArgumentError("unable to check bounds for indices of type $(typeof(i))"))
 checkindex(::Type{Bool}, inds::AbstractUnitRange, i::Real) = (first(inds) <= i) & (i <= last(inds))
 checkindex(::Type{Bool}, inds::IdentityUnitRange, i::Real) = checkindex(Bool, inds.indices, i)
 checkindex(::Type{Bool}, inds::OneTo{T}, i::T) where {T<:BitInteger} = unsigned(i - one(i)) < unsigned(last(inds))
 checkindex(::Type{Bool}, inds::AbstractUnitRange, ::Colon) = true
 checkindex(::Type{Bool}, inds::AbstractUnitRange, ::Slice) = true
-function checkindex(::Type{Bool}, inds::AbstractUnitRange, r::AbstractRange)
- @_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, i::AbstractRange) =
+ isempty(i) | (checkindex(Bool, inds, first(i)) & checkindex(Bool, inds, last(i)))
+# range like indices with cheap `extrema`
+checkindex(::Type{Bool}, inds::AbstractUnitRange, i::LinearIndices) =
+ isempty(i) | (checkindex(Bool, inds, first(i)) & checkindex(Bool, inds, last(i)))
+
+function checkindex(::Type{Bool}, inds, I::AbstractArray)
  @inline
  b = true
  for i in I

base/indices.jl

@@ -349,15 +349,8 @@ to_indices(A, I::Tuple{}) = ()
 to_indices(A, I::Tuple{Vararg{Int}}) = I
 to_indices(A, I::Tuple{Vararg{Integer}}) = (@inline; to_indices(A, (), I))
 to_indices(A, inds, ::Tuple{}) = ()
-function to_indices(A, inds, I::Tuple{Any, Vararg{Any}})
- @inline
- head = _to_indices1(A, inds, I[1])
- rest = to_indices(A, _cutdim(inds, I[1]), tail(I))
- (head..., rest...)
-end
-
-_to_indices1(A, inds, I1) = (to_index(A, I1),)
-_cutdim(inds, I1) = safe_tail(inds)
+to_indices(A, inds, I::Tuple{Any, Vararg}) =
+ (@inline; (to_index(A, I[1]), to_indices(A, safe_tail(inds), tail(I))...))
 
 """
  Slice(indices)

base/multidimensional.jl

@@ -4,12 +4,12 @@
 module IteratorsMD
  import .Base: eltype, length, size, first, last, in, getindex, setindex!,
  min, max, zero, oneunit, isless, eachindex,
- convert, show, iterate, promote_rule
+ convert, show, iterate, promote_rule, to_indices
 
  import .Base: +, -, *, (:)
  import .Base: simd_outer_range, simd_inner_length, simd_index, setindex
- import .Base: to_indices, to_index, _to_indices1, _cutdim
- using .Base: IndexLinear, IndexCartesian, AbstractCartesianIndex, fill_to_length, tail,
+ using .Base: to_index, fill_to_length, tail, safe_tail
+ using .Base: IndexLinear, IndexCartesian, AbstractCartesianIndex,
  ReshapedArray, ReshapedArrayLF, OneTo, Fix1
  using .Base.Iterators: Reverse, PartitionIterator
  using .Base: @propagate_inbounds
@@ -465,11 +465,13 @@ module IteratorsMD
  last(iter::CartesianIndices) = CartesianIndex(map(last, iter.indices))
 
  # When used as indices themselves, CartesianIndices can simply become its tuple of ranges
- _to_indices1(A, inds, I1::CartesianIndices) = map(Fix1(to_index, A), I1.indices)
- _cutdim(inds::Tuple, I1::CartesianIndices) = split(inds, Val(ndims(I1)))[2]
-
+ @inline function to_indices(A, inds, I::Tuple{CartesianIndices{N}, Vararg}) where N
+ _, indstail = split(inds, Val(N))
+ (map(Fix1(to_index, A), I[1].indices)..., to_indices(A, indstail, tail(I))...)
+ end
  # but preserve CartesianIndices{0} as they consume a dimension.
- _to_indices1(A, inds, I1::CartesianIndices{0}) = (I1,)
+ @inline to_indices(A, inds, I::Tuple{CartesianIndices{0}, Vararg}) =
+ (first(I), to_indices(A, inds, tail(I))...)
 
  @inline in(i::CartesianIndex, r::CartesianIndices) = false
  @inline in(i::CartesianIndex{N}, r::CartesianIndices{N}) where {N} = all(map(in, i.I, r.indices))
@@ -682,17 +684,15 @@ using .IteratorsMD
 
 ## Bounds-checking with CartesianIndex
 # Disallow linear indexing with CartesianIndex
-function checkbounds(::Type{Bool}, A::AbstractArray, i::Union{CartesianIndex, AbstractArray{<:CartesianIndex}})
- @inline
+@inline checkbounds(::Type{Bool}, A::AbstractArray, i::CartesianIndex) =
  checkbounds_indices(Bool, axes(A), (i,))
+# Here we try to consume N of the indices (if there are that many available)
+@inline function checkbounds_indices(::Type{Bool}, inds::Tuple, I::Tuple{CartesianIndex,Vararg})
+ inds1, rest = IteratorsMD.split(inds, Val(length(I[1])))
+ checkindex(Bool, inds1, I[1]) & checkbounds_indices(Bool, rest, tail(I))
 end
-
-@inline checkbounds_indices(::Type{Bool}, ::Tuple{}, I::Tuple{CartesianIndex,Vararg{Any}}) =
- checkbounds_indices(Bool, (), (I[1].I..., tail(I)...))
-@inline checkbounds_indices(::Type{Bool}, IA::Tuple{Any}, I::Tuple{CartesianIndex,Vararg{Any}}) =
- checkbounds_indices(Bool, IA, (I[1].I..., tail(I)...))
-@inline checkbounds_indices(::Type{Bool}, IA::Tuple, I::Tuple{CartesianIndex,Vararg{Any}}) =
- checkbounds_indices(Bool, IA, (I[1].I..., tail(I)...))
+@inline checkindex(::Type{Bool}, inds::Tuple, I::CartesianIndex) =
+ checkbounds_indices(Bool, inds, I.I)
 
 # Indexing into Array with mixtures of Integers and CartesianIndices is
 # extremely performance-sensitive. While the abstract fallbacks support this,
@@ -702,45 +702,17 @@ end
 @propagate_inbounds setindex!(A::Array, v, i1::Union{Integer, CartesianIndex}, I::Union{Integer, CartesianIndex}...) =
  (A[to_indices(A, (i1, I...))...] = v; A)
 
-# Support indexing with an array of CartesianIndex{N}s
+## Bounds-checking with arrays of CartesianIndex{N}
+# Disallow linear indexing with an array of CartesianIndex{N}
+@inline checkbounds(::Type{Bool}, A::AbstractArray, i::AbstractArray{CartesianIndex{N}}) where {N} =
+ checkbounds_indices(Bool, axes(A), (i,))
 # Here we try to consume N of the indices (if there are that many available)
-# The first two simply handle ambiguities
-@inline function checkbounds_indices(::Type{Bool}, ::Tuple{},
- I::Tuple{AbstractArray{CartesianIndex{N}},Vararg{Any}}) where N
- checkindex(Bool, (), I[1]) & checkbounds_indices(Bool, (), tail(I))
-end
-@inline function checkbounds_indices(::Type{Bool}, IA::Tuple{Any},
- I::Tuple{AbstractArray{CartesianIndex{0}},Vararg{Any}})
- checkbounds_indices(Bool, IA, tail(I))
-end
-@inline function checkbounds_indices(::Type{Bool}, IA::Tuple{Any},
- I::Tuple{AbstractArray{CartesianIndex{N}},Vararg{Any}}) where N
- checkindex(Bool, IA, I[1]) & checkbounds_indices(Bool, (), tail(I))
-end
-@inline function checkbounds_indices(::Type{Bool}, IA::Tuple,
- I::Tuple{AbstractArray{CartesianIndex{N}},Vararg{Any}}) where N
- IA1, IArest = IteratorsMD.split(IA, Val(N))
- checkindex(Bool, IA1, I[1]) & checkbounds_indices(Bool, IArest, tail(I))
+@inline function checkbounds_indices(::Type{Bool}, inds::Tuple, I::Tuple{AbstractArray{CartesianIndex{N}},Vararg}) where N
+ inds1, rest = IteratorsMD.split(inds, Val(N))
+ checkindex(Bool, inds1, I[1]) & checkbounds_indices(Bool, rest, tail(I))
 end
-
-
-@inline function checkbounds_indices(::Type{Bool}, IA::Tuple{},
- I::Tuple{AbstractArray{Bool,N},Vararg{Any}}) where N
- return checkbounds_indices(Bool, IA, (LogicalIndex(I[1]), tail(I)...))
-end
-@inline function checkbounds_indices(::Type{Bool}, IA::Tuple,
- I::Tuple{AbstractArray{Bool,N},Vararg{Any}}) where N
- return checkbounds_indices(Bool, IA, (LogicalIndex(I[1]), tail(I)...))
-end
-
-function checkindex(::Type{Bool}, inds::Tuple, I::AbstractArray{<:CartesianIndex})
- b = true
- for i in I
- b &= checkbounds_indices(Bool, inds, (i,))
- end
- b
-end
-checkindex(::Type{Bool}, inds::Tuple, I::CartesianIndices) = all(checkindex.(Bool, inds, I.indices))
+@inline checkindex(::Type{Bool}, inds::Tuple, I::CartesianIndices) =
+ checkbounds_indices(Bool, inds, I.indices)
 
 # combined count of all indices, including CartesianIndex and
 # AbstractArray{CartesianIndex}
@@ -848,11 +820,29 @@ end
  return eltype(L)(i1, irest...), (i1 - tz, Bi, irest, c)
 end
 
-@inline checkbounds(::Type{Bool}, A::AbstractArray, I::LogicalIndex{<:Any,<:AbstractArray{Bool,1}}) =
- eachindex(IndexLinear(), A) == eachindex(IndexLinear(), I.mask)
-@inline checkbounds(::Type{Bool}, A::AbstractArray, I::LogicalIndex) = axes(A) == axes(I.mask)
-@inline checkindex(::Type{Bool}, indx::AbstractUnitRange, I::LogicalIndex) = (indx,) == axes(I.mask)
-checkindex(::Type{Bool}, inds::Tuple, I::LogicalIndex) = checkbounds_indices(Bool, inds, axes(I.mask))
+## Boundscheck for Logicalindex
+# LogicalIndex: map all calls to mask
+checkbounds(::Type{Bool}, A::AbstractArray, i::LogicalIndex) = checkbounds(Bool, A, i.mask)
+# `checkbounds_indices` has been handled via `I::AbstractArray` fallback
+checkindex(::Type{Bool}, inds::AbstractUnitRange, i::LogicalIndex) = checkindex(Bool, inds, i.mask)
+checkindex(::Type{Bool}, inds::Tuple, i::LogicalIndex) = checkindex(Bool, inds, i.mask)
+
+## Boundscheck for AbstractArray{Bool}
+# Disallow linear indexing with AbstractArray{Bool}
+checkbounds(::Type{Bool}, A::AbstractArray, i::AbstractArray{Bool}) =
+ checkbounds_indices(Bool, axes(A), (i,))
+# But allow linear indexing with AbstractVector{Bool}
+checkbounds(::Type{Bool}, A::AbstractArray, i::AbstractVector{Bool}) =
+ checkindex(Bool, eachindex(IndexLinear(), A), i)
+@inline function checkbounds_indices(::Type{Bool}, inds::Tuple, I::Tuple{AbstractArray{Bool},Vararg})
+ inds1, rest = IteratorsMD.split(inds, Val(ndims(I[1])))
+ checkindex(Bool, inds1, I[1]) & checkbounds_indices(Bool, rest, tail(I))
+end
+checkindex(::Type{Bool}, inds::AbstractUnitRange, I::AbstractVector{Bool}) = axes1(I) == inds
+checkindex(::Type{Bool}, inds::AbstractUnitRange, I::AbstractRange{Bool}) = axes1(I) == inds
+checkindex(::Type{Bool}, inds::Tuple, I::AbstractArray{Bool}) = _check_boolean_axes(inds, axes(I))
+_check_boolean_axes(inds::Tuple, axes::Tuple) = (inds[1] == axes[1]) & _check_boolean_axes(tail(inds), tail(axes))
+_check_boolean_axes(::Tuple{}, axes::Tuple) = all(==(OneTo(1)), axes)
 
 ensure_indexable(I::Tuple{}) = ()
 @inline ensure_indexable(I::Tuple{Any, Vararg{Any}}) = (I[1], ensure_indexable(tail(I))...)
@@ -863,20 +853,28 @@ ensure_indexable(I::Tuple{}) = ()
 @inline to_indices(A, I::Tuple{Vararg{Union{Integer, CartesianIndex}}}) = to_indices(A, (), I)
 # But some index types require more context spanning multiple indices
 # CartesianIndex is unfolded outside the inner to_indices for better inference
-_to_indices1(A, inds, I1::CartesianIndex) = map(Fix1(to_index, A), I1.I)
-_cutdim(inds, I1::CartesianIndex) = IteratorsMD.split(inds, Val(length(I1)))[2]
+@inline function to_indices(A, inds, I::Tuple{CartesianIndex{N}, Vararg}) where N
+ _, indstail = IteratorsMD.split(inds, Val(N))
+ (map(Fix1(to_index, A), I[1].I)..., to_indices(A, indstail, tail(I))...)
+end
 # For arrays of CartesianIndex, we just skip the appropriate number of inds
-_cutdim(inds, I1::AbstractArray{CartesianIndex{N}}) where {N} = IteratorsMD.split(inds, Val(N))[2]
+@inline function to_indices(A, inds, I::Tuple{AbstractArray{CartesianIndex{N}}, Vararg}) where N
+ _, indstail = IteratorsMD.split(inds, Val(N))
+ (to_index(A, I[1]), to_indices(A, indstail, tail(I))...)
+end
 # And boolean arrays behave similarly; they also skip their number of dimensions
-_cutdim(inds::Tuple, I1::AbstractArray{Bool}) = IteratorsMD.split(inds, Val(ndims(I1)))[2]
-# As an optimization, we allow trailing Array{Bool} and BitArray to be linear over trailing dimensions
-@inline to_indices(A, inds, I::Tuple{Union{Array{Bool,N}, BitArray{N}}}) where {N} =
- (_maybe_linear_logical_index(IndexStyle(A), A, I[1]),)
+@inline function to_indices(A, inds, I::Tuple{AbstractArray{Bool, N}, Vararg}) where N
+ _, indstail = IteratorsMD.split(inds, Val(N))
+ (to_index(A, I[1]), to_indices(A, indstail, tail(I))...)
+end
+# As an optimization, we allow the only `AbstractArray{Bool}` to be linear-iterated
+@inline to_indices(A, I::Tuple{AbstractArray{Bool}}) = (_maybe_linear_logical_index(IndexStyle(A), A, I[1]),)
 _maybe_linear_logical_index(::IndexStyle, A, i) = to_index(A, i)
 _maybe_linear_logical_index(::IndexLinear, A, i) = LogicalIndex{Int}(i)
 
 # Colons get converted to slices by `uncolon`
-_to_indices1(A, inds, I1::Colon) = (uncolon(inds),)
+@inline to_indices(A, inds, I::Tuple{Colon, Vararg}) =
+ (uncolon(inds), to_indices(A, Base.safe_tail(inds), tail(I))...)
 
 uncolon(::Tuple{}) = Slice(OneTo(1))
 uncolon(inds::Tuple) = Slice(inds[1])

doc/src/manual/arrays.md

@@ -793,38 +793,46 @@ Indexing by a boolean vector `B` is effectively the same as indexing by the
 vector of integers that is returned by [`findall(B)`](@ref). Similarly, indexing
 by a `N`-dimensional boolean array is effectively the same as indexing by the
 vector of `CartesianIndex{N}`s where its values are `true`. A logical index
-must be a vector of the same length as the dimension it indexes into, or it
-must be the only index provided and match the size and dimensionality of the
-array it indexes into. It is generally more efficient to use boolean arrays as
-indices directly instead of first calling [`findall`](@ref).
+must be a array of the same shape as the dimension(s) it indexes into, or it
+must be the only index provided and match the shape of the one-dimensional
+reshaped view of the array it indexes into. It is generally more efficient
+to use boolean arrays as indices directly instead of first calling [`findall`](@ref).
 
 ```jldoctest
-julia> x = reshape(1:16, 4, 4)
-4×4 reshape(::UnitRange{Int64}, 4, 4) with eltype Int64:
- 1 5 9 13
- 2 6 10 14
- 3 7 11 15
- 4 8 12 16
+julia> x = reshape(1:12, 2, 3, 2)
+2×3×2 reshape(::UnitRange{Int64}, 2, 3, 2) with eltype Int64:
+[:, :, 1] =
+ 1 3 5
+ 2 4 6
 
-julia> x[[false, true, true, false], :]
-2×4 Matrix{Int64}:
- 2 6 10 14
- 3 7 11 15
+[:, :, 2] =
+ 7 9 11
+ 8 10 12
+
+julia> x[:, [true false; false true; true false]]
+2×3 Matrix{Int64}:
+ 1 5 9
+ 2 6 10
 
 julia> mask = map(ispow2, x)
-4×4 Matrix{Bool}:
- 1 0 0 0
- 1 0 0 0
- 0 0 0 0
- 1 1 0 1
+2×3×2 Array{Bool, 3}:
+[:, :, 1] =
+ 1 0 0
+ 1 1 0
+
+[:, :, 2] =
+ 0 0 0
+ 1 0 0
 
 julia> x[mask]
-5-element Vector{Int64}:
- 1
- 2
- 4
- 8
- 16
+4-element Vector{Int64}:
+ 1
+ 2
+ 4
+ 8
+
+julia> x[vec(mask)] == x[mask] # we can also index with a single Boolean vector
+true
 ```
 
 ### Number of indices