From 38c0ef1301ce3bc4e3319b7af88f8a5195d478e8 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 3 Jun 2025 17:33:47 +0530 Subject: [PATCH] Define `in` for `CartesianIndex` ranges --- base/multidimensional.jl | 18 ++++++++++++++++++ test/cartesian.jl | 39 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/base/multidimensional.jl b/base/multidimensional.jl index 7add7b9e74205..fd87c7afcf80f 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -189,6 +189,24 @@ module IteratorsMD step(r), ", ", length(r), ")") end + Base.in(x::CartesianIndex, r::AbstractRange{<:CartesianIndex}) = false + function Base.in(x::CartesianIndex{N}, r::AbstractRange{CartesianIndex{N}}) where {N} + isempty(r) && return false + f, st, l = first(r), step(r), last(r) + # The n-th element of the range is a CartesianIndex + # whose elements are the n-th along each dimension + # Find the first dimension along which the index is changing, + # so that n may be uniquely determined + for i in 1:N + iszero(st[i]) && continue + n = findfirst(==(x[i]), f[i]:st[i]:l[i]) + isnothing(n) && return false + return r[n] == x + end + # if the step is zero, the elements are identical, so compare with the first + return x == f + end + # Iteration const OrdinalRangeInt = OrdinalRange{Int, Int} """ diff --git a/test/cartesian.jl b/test/cartesian.jl index f9a0c8f976956..6097d4ca3770a 100644 --- a/test/cartesian.jl +++ b/test/cartesian.jl @@ -588,3 +588,42 @@ end @test I[begin] == I[1] @test I[end] == I[2] end + +@testset "in for a CartesianIndex StepRangeLen" begin + @testset for l in [0, 1, 4], r in Any[ + StepRangeLen(CartesianIndex(), CartesianIndex(), l), + StepRangeLen(CartesianIndex(1), CartesianIndex(0), l), + StepRangeLen(CartesianIndex(1), CartesianIndex(1), l), + StepRangeLen(CartesianIndex(1), CartesianIndex(4), l), + StepRangeLen(CartesianIndex(1), CartesianIndex(-4), l), + StepRangeLen(CartesianIndex(-1, 2), CartesianIndex(0, 0), l), + StepRangeLen(CartesianIndex(-1, 2), CartesianIndex(0, 4), l), + StepRangeLen(CartesianIndex(-1, 2), CartesianIndex(0, -4), l), + StepRangeLen(CartesianIndex(-1, 2), CartesianIndex(4, 0), l), + StepRangeLen(CartesianIndex(-1, 2), CartesianIndex(-4, 0), l), + StepRangeLen(CartesianIndex(-1, 2), CartesianIndex(4, 2), l), + StepRangeLen(CartesianIndex(-1, 2), CartesianIndex(-4, 2), l), + StepRangeLen(CartesianIndex(-1, 2), CartesianIndex(4, -2), l), + StepRangeLen(CartesianIndex(-1, 2), CartesianIndex(-4, -2), l), + StepRangeLen(CartesianIndex(-1, 2, 0), CartesianIndex(0, 0, 0), l), + StepRangeLen(CartesianIndex(-1, 2, 0), CartesianIndex(0, 0, -2), l), + ] + + if length(r) == 0 + @test !(first(r) in r) + @test !(last(r) in r) + end + for x in r + @test x in r + if step(r) != oneunit(x) + @test !((x + oneunit(x)) in r) + end + end + @test !(CartesianIndex(ntuple(x->0, ndims(r))) in r) + @test !(CartesianIndex(ntuple(x->typemax(Int), ndims(r))) in r) + @test !(CartesianIndex(ntuple(x->typemin(Int), ndims(r))) in r) + if ndims(r) > 1 + @test !(CartesianIndex(ntuple(x->0, ndims(r)-1)...) in r) + end + end +end