More efficient indexing by CartesianIndices (#31214)

Treat CartesianIndices as its tuple of ranges when indexing. This allows views that use CartesianIndices as indices to be strided and changes the complexity of boundschecking from O(n^d) to O(d)

Co-authored-by: Kristoffer Carlsson <kristoffer.carlsson@chalmers.se>

Author: mbauman

base/multidimensional.jl

@@ -4,7 +4,7 @@
 module IteratorsMD
     import .Base: eltype, length, size, first, last, in, getindex,
                  setindex!, IndexStyle, min, max, zero, oneunit, isless, eachindex,
-                 ndims, IteratorSize, convert, show, iterate, promote_rule
+                 ndims, IteratorSize, convert, show, iterate, promote_rule, to_indices
 
     import .Base: +, -, *, (:)
     import .Base: simd_outer_range, simd_inner_length, simd_index
@@ -381,6 +381,10 @@ module IteratorsMD
     first(iter::CartesianIndices) = CartesianIndex(map(first, iter.indices))
     last(iter::CartesianIndices)  = CartesianIndex(map(last, iter.indices))
 
+    # When used as indices themselves, CartesianIndices can simply become its tuple of ranges
+    @inline to_indices(A, inds, I::Tuple{CartesianIndices, Vararg{Any}}) =
+        to_indices(A, inds, (I[1].indices..., tail(I)...))
+
     @inline function in(i::CartesianIndex{N}, r::CartesianIndices{N}) where {N}
         _in(true, i.I, first(r).I, last(r).I)
     end
@@ -573,6 +577,7 @@ function checkindex(::Type{Bool}, inds::Tuple, I::AbstractArray{<:CartesianIndex
     end
     b
 end
+checkindex(::Type{Bool}, inds::Tuple, I::CartesianIndices) = all(checkindex.(Bool, inds, I.indices))
 
 # combined count of all indices, including CartesianIndex and
 # AbstractArray{CartesianIndex}

test/arrayops.jl

@@ -2648,6 +2648,21 @@ end
 # Fix oneunit bug for unitful arrays
 @test oneunit([Second(1) Second(2); Second(3) Second(4)]) == [Second(1) Second(0); Second(0) Second(1)]
 
+@testset "indexing by CartesianIndices" begin
+    A = rand(10,10)
+    for (I,Rs) in ((keys(A), (1:10, 1:10)),
+                   (CartesianIndex(2,2):CartesianIndex(9,9), (2:9, 2:9)),
+                   (CartesianIndex(5,3):CartesianIndex(6,7), (5:6, 3:7)))
+        @test A[I] == A[Rs...] == @view(A[I]) == @view(A[Rs...])
+        @test @view(A[I]) isa StridedArray
+        @test !checkbounds(Bool, [], I)
+        @test !checkbounds(Bool, fill(2,1,1,1), :, I)
+        @test !checkbounds(Bool, fill(2,1,1,1), I, :)
+    end
+    @test !checkbounds(Bool, rand(3,3,3), :, CartesianIndex(0,0):CartesianIndex(1,1))
+    @test !checkbounds(Bool, rand(3,3,3), CartesianIndex(0,0):CartesianIndex(1,1), :)
+end
+
 # Throws ArgumentError for negative dimensions in Array
 @test_throws ArgumentError fill('a', -10)