From 39b286c46dbf547b7c94b0b4e0a9df04fafc2c2c Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Mon, 24 Feb 2025 10:25:49 +0530 Subject: [PATCH 1/3] Specialize indexing a CartesianIndices with a StepRangeLen{<:CartesianIndex} --- base/multidimensional.jl | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/base/multidimensional.jl b/base/multidimensional.jl index 7add7b9e74205..f029e40bc4912 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -403,6 +403,12 @@ module IteratorsMD getindex(iter, C.indices...) end @inline Base.getindex(iter::CartesianIndices{0}, ::CartesianIndices{0}) = iter + @inline function Base.getindex(iter::CartesianIndices{N}, r::StepRangeLen{CartesianIndex{N},CartesianIndex{N},CartesianIndex{N},<:Integer}) where {N} + @boundscheck checkbounds(iter, r) + start = first(iter) + first(r) - CartesianIndex(first.(axes(iter))) + stepsz = CartesianIndex(Tuple(step(iter)) .* Tuple(step(r))) + StepRangeLen(start, stepsz, length(r)) + end # If dimensions permit, we may index into a CartesianIndices directly instead of constructing a SubArray wrapper @propagate_inbounds function Base.view(c::CartesianIndices{N}, r::Vararg{Union{OrdinalRange{<:Integer, <:Integer}, Colon},N}) where {N} @@ -743,7 +749,7 @@ end @inline checkindex(::Type{Bool}, inds::Tuple, I::CartesianIndex) = checkbounds_indices(Bool, inds, I.I) @inline checkindex(::Type{Bool}, inds::Tuple, i::AbstractRange{<:CartesianIndex}) = - isempty(i) | (checkindex(Bool, inds, first(i)) & checkindex(Bool, inds, last(i))) + isempty(i) || (checkindex(Bool, inds, first(i)) & checkindex(Bool, inds, last(i))) # Indexing into Array with mixtures of Integers and CartesianIndices is # extremely performance-sensitive. While the abstract fallbacks support this, From f44252fc7e4173d167ead827e3ff77d177101f1b Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 25 Feb 2025 23:09:47 +0530 Subject: [PATCH 2/3] Fix starting point and add tests --- base/multidimensional.jl | 2 +- test/abstractarray.jl | 28 ++++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/base/multidimensional.jl b/base/multidimensional.jl index f029e40bc4912..5a1e5c3a82171 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -405,7 +405,7 @@ module IteratorsMD @inline Base.getindex(iter::CartesianIndices{0}, ::CartesianIndices{0}) = iter @inline function Base.getindex(iter::CartesianIndices{N}, r::StepRangeLen{CartesianIndex{N},CartesianIndex{N},CartesianIndex{N},<:Integer}) where {N} @boundscheck checkbounds(iter, r) - start = first(iter) + first(r) - CartesianIndex(first.(axes(iter))) + start = first(iter) + CartesianIndex((Tuple(first(r)) .- first.(axes(iter))) .* Tuple(step(iter))) stepsz = CartesianIndex(Tuple(step(iter)) .* Tuple(step(r))) StepRangeLen(start, stepsz, length(r)) end diff --git a/test/abstractarray.jl b/test/abstractarray.jl index bbac7812174da..9d9fcd0716672 100644 --- a/test/abstractarray.jl +++ b/test/abstractarray.jl @@ -285,6 +285,34 @@ end end end +@testset "Indexing CartesianIndices as an nD range" begin + for (m, n) in [(4, 6), (4, 4), (6,4)] + C = CartesianIndices((m, n)) + for k in -7:7 + dlen = length = max(0, k <= 0 ? min(m+k, n) : min(m, n-k)) + dind = StepRangeLen(CartesianIndex(1+max(0,-k),1+max(0,k)), CartesianIndex(1,1), dlen) + @test C[dind] === dind + @test C[dind[1:2:end]] === dind[1:2:end] + end + end + for C in [CartesianIndices((20:4:100,)), + CartesianIndices((20:-4:-100,)), + CartesianIndices((20:100,)), + CartesianIndices((Base.IdentityUnitRange(20:100),))] + r = StepRangeLen(CartesianIndex(firstindex(C)+1), CartesianIndex(3), 4) + @test C[r] == C[collect(r)] + end + C = CartesianIndices((3:8, 3:8, 3:8)) + r = StepRangeLen(CartesianIndex(1,1,1), CartesianIndex(1,1,1), 6) + @test C[r] == StepRangeLen(CartesianIndex(3,3,3), CartesianIndex(1,1,1), 6) + r = StepRangeLen(CartesianIndex(1,1,1), CartesianIndex(1,0,0), 6) + @test C[r] == StepRangeLen(CartesianIndex(3,3,3), CartesianIndex(1,0,0), 6) + + C = CartesianIndices((3:8, 3:8, Base.IdentityUnitRange(3:8))) + r = StepRangeLen(CartesianIndex(1,1,3), CartesianIndex(1,1,1), 6) + @test C[r] == StepRangeLen(CartesianIndex(3,3,3), CartesianIndex(1,1,1), 6) +end + @testset "LinearIndices" begin @testset "constructors" begin for oinds in [ From 6319f2b75be823e9fcfb0dd405bafddc2855981d Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Mon, 16 Jun 2025 23:41:33 +0530 Subject: [PATCH 3/3] Relax signature to `AbstractRange` --- base/multidimensional.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/base/multidimensional.jl b/base/multidimensional.jl index 5a1e5c3a82171..1cc1f53575a32 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -403,7 +403,7 @@ module IteratorsMD getindex(iter, C.indices...) end @inline Base.getindex(iter::CartesianIndices{0}, ::CartesianIndices{0}) = iter - @inline function Base.getindex(iter::CartesianIndices{N}, r::StepRangeLen{CartesianIndex{N},CartesianIndex{N},CartesianIndex{N},<:Integer}) where {N} + @inline function Base.getindex(iter::CartesianIndices{N}, r::AbstractRange{CartesianIndex{N}}) where {N} @boundscheck checkbounds(iter, r) start = first(iter) + CartesianIndex((Tuple(first(r)) .- first.(axes(iter))) .* Tuple(step(iter))) stepsz = CartesianIndex(Tuple(step(iter)) .* Tuple(step(r)))