Define in for CartesianIndex ranges

Author: jishnub

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}
     """

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