diff --git a/NEWS.md b/NEWS.md index 2b7c281a8df31..c8dda2673ad9f 100644 --- a/NEWS.md +++ b/NEWS.md @@ -104,6 +104,9 @@ Standard library changes * `first` and `last` functions now accept an integer as second argument to get that many leading or trailing elements of any iterable ([#34868]). * `intersect` on `CartesianIndices` now returns `CartesianIndices` instead of `Vector{<:CartesianIndex}` ([#36643]). +* `CartesianIndices` now supports step different from `1`. It can also be constructed from three + `CartesianIndex`es `I`, `S`, `J` using `I:S:J`. `step` for `CartesianIndices` now returns a + `CartesianIndex`. ([#37829]) * `push!(c::Channel, v)` now returns channel `c`. Previously, it returned the pushed value `v` ([#34202]). * `RegexMatch` objects can now be probed for whether a named capture group exists within it through `haskey()` ([#36717]). * For consistency `haskey(r::RegexMatch, i::Integer)` has also been added and returns if the capture group for `i` exists ([#37300]). diff --git a/base/broadcast.jl b/base/broadcast.jl index 12bddaf531e27..6229e29d56118 100644 --- a/base/broadcast.jl +++ b/base/broadcast.jl @@ -1112,7 +1112,7 @@ broadcasted(::typeof(+), j::CartesianIndex{N}, I::CartesianIndices{N}) where N = broadcasted(::typeof(-), I::CartesianIndices{N}, j::CartesianIndex{N}) where N = CartesianIndices(map((rng, offset)->rng .- offset, I.indices, Tuple(j))) function broadcasted(::typeof(-), j::CartesianIndex{N}, I::CartesianIndices{N}) where N - diffrange(offset, rng) = range(offset-last(rng), length=length(rng)) + diffrange(offset, rng) = range(offset-last(rng), length=length(rng), step=step(rng)) Iterators.reverse(CartesianIndices(map(diffrange, Tuple(j), I.indices))) end diff --git a/base/multidimensional.jl b/base/multidimensional.jl index f19afd78d481b..66164c334a43d 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -11,6 +11,7 @@ module IteratorsMD using .Base: IndexLinear, IndexCartesian, AbstractCartesianIndex, fill_to_length, tail, ReshapedArray, ReshapedArrayLF, OneTo using .Base.Iterators: Reverse, PartitionIterator + using .Base: @propagate_inbounds export CartesianIndex, CartesianIndices @@ -149,13 +150,13 @@ module IteratorsMD function Base.nextind(a::AbstractArray{<:Any,N}, i::CartesianIndex{N}) where {N} iter = CartesianIndices(axes(a)) # might overflow - I = inc(i.I, first(iter).I, last(iter).I) + I = inc(i.I, iter.indices) return I end function Base.prevind(a::AbstractArray{<:Any,N}, i::CartesianIndex{N}) where {N} iter = CartesianIndices(axes(a)) # might underflow - I = dec(i.I, last(iter).I, first(iter).I) + I = dec(i.I, iter.indices) return I end @@ -169,15 +170,15 @@ module IteratorsMD # Iteration """ CartesianIndices(sz::Dims) -> R - CartesianIndices((istart:istop, jstart:jstop, ...)) -> R + CartesianIndices((istart:[istep:]istop, jstart:[jstep:]jstop, ...)) -> R Define a region `R` spanning a multidimensional rectangular range of integer indices. These are most commonly encountered in the context of iteration, where `for I in R ... end` will return [`CartesianIndex`](@ref) indices `I` equivalent to the nested loops - for j = jstart:jstop - for i = istart:istop + for j = jstart:jstep:jstop + for i = istart:istep:istop ... end end @@ -190,6 +191,10 @@ module IteratorsMD As a convenience, constructing a `CartesianIndices` from an array makes a range of its indices. + !!! compat "Julia 1.6" + The step range method `CartesianIndices((istart:istep:istop, jstart:[jstep:]jstop, ...))` + requires at least Julia 1.6. + # Examples ```jldoctest julia> foreach(println, CartesianIndices((2, 2, 2))) @@ -222,6 +227,15 @@ module IteratorsMD julia> cartesian[4] CartesianIndex(1, 2) + + julia> cartesian = CartesianIndices((1:2:5, 1:2)) + 3×2 CartesianIndices{2, Tuple{StepRange{Int64, Int64}, UnitRange{Int64}}}: + CartesianIndex(1, 1) CartesianIndex(1, 2) + CartesianIndex(3, 1) CartesianIndex(3, 2) + CartesianIndex(5, 1) CartesianIndex(5, 2) + + julia> cartesian[2, 2] + CartesianIndex(3, 2) ``` ## Broadcasting @@ -248,29 +262,37 @@ module IteratorsMD For cartesian to linear index conversion, see [`LinearIndices`](@ref). """ - struct CartesianIndices{N,R<:NTuple{N,AbstractUnitRange{Int}}} <: AbstractArray{CartesianIndex{N},N} + struct CartesianIndices{N,R<:NTuple{N,OrdinalRange{Int, Int}}} <: AbstractArray{CartesianIndex{N},N} indices::R end CartesianIndices(::Tuple{}) = CartesianIndices{0,typeof(())}(()) - CartesianIndices(inds::NTuple{N,AbstractUnitRange{<:Integer}}) where {N} = - CartesianIndices(map(r->convert(AbstractUnitRange{Int}, r), inds)) + function CartesianIndices(inds::NTuple{N,OrdinalRange{<:Integer, <:Integer}}) where {N} + indices = map(r->convert(OrdinalRange{Int, Int}, r), inds) + CartesianIndices{N, typeof(indices)}(indices) + end CartesianIndices(index::CartesianIndex) = CartesianIndices(index.I) - CartesianIndices(sz::NTuple{N,<:Integer}) where {N} = CartesianIndices(map(Base.OneTo, sz)) - CartesianIndices(inds::NTuple{N,Union{<:Integer,AbstractUnitRange{<:Integer}}}) where {N} = - CartesianIndices(map(i->first(i):last(i), inds)) + CartesianIndices(inds::NTuple{N,Union{<:Integer,OrdinalRange{<:Integer}}}) where {N} = + CartesianIndices(map(_convert2ind, inds)) CartesianIndices(A::AbstractArray) = CartesianIndices(axes(A)) + _convert2ind(sz::Integer) = Base.OneTo(sz) + _convert2ind(sz::AbstractUnitRange) = first(sz):last(sz) + _convert2ind(sz::OrdinalRange) = first(sz):step(sz):last(sz) + """ - (:)(I::CartesianIndex, J::CartesianIndex) + (:)(start::CartesianIndex, [step::CartesianIndex], stop::CartesianIndex) - Construct [`CartesianIndices`](@ref) from two `CartesianIndex`. + Construct [`CartesianIndices`](@ref) from two `CartesianIndex` and an optional step. !!! compat "Julia 1.1" This method requires at least Julia 1.1. + !!! compat "Julia 1.6" + The step range method start:step:stop requires at least Julia 1.6. + # Examples ```jldoctest julia> I = CartesianIndex(2,1); @@ -281,17 +303,26 @@ module IteratorsMD 2×3 CartesianIndices{2, Tuple{UnitRange{Int64}, UnitRange{Int64}}}: CartesianIndex(2, 1) CartesianIndex(2, 2) CartesianIndex(2, 3) CartesianIndex(3, 1) CartesianIndex(3, 2) CartesianIndex(3, 3) + + julia> I:CartesianIndex(1, 2):J + 2×2 CartesianIndices{2, Tuple{StepRange{Int64, Int64}, StepRange{Int64, Int64}}}: + CartesianIndex(2, 1) CartesianIndex(2, 3) + CartesianIndex(3, 1) CartesianIndex(3, 3) ``` """ (:)(I::CartesianIndex{N}, J::CartesianIndex{N}) where N = CartesianIndices(map((i,j) -> i:j, Tuple(I), Tuple(J))) + (:)(I::CartesianIndex{N}, S::CartesianIndex{N}, J::CartesianIndex{N}) where N = + CartesianIndices(map((i,s,j) -> i:s:j, Tuple(I), Tuple(S), Tuple(J))) promote_rule(::Type{CartesianIndices{N,R1}}, ::Type{CartesianIndices{N,R2}}) where {N,R1,R2} = CartesianIndices{N,Base.indices_promote_type(R1,R2)} convert(::Type{Tuple{}}, R::CartesianIndices{0}) = () - convert(::Type{NTuple{N,AbstractUnitRange{Int}}}, R::CartesianIndices{N}) where {N} = - R.indices + for RT in (OrdinalRange{Int, Int}, StepRange{Int, Int}, AbstractUnitRange{Int}) + @eval convert(::Type{NTuple{N,$RT}}, R::CartesianIndices{N}) where {N} = + map(x->convert($RT, x), R.indices) + end convert(::Type{NTuple{N,AbstractUnitRange}}, R::CartesianIndices{N}) where {N} = convert(NTuple{N,AbstractUnitRange{Int}}, R) convert(::Type{NTuple{N,UnitRange{Int}}}, R::CartesianIndices{N}) where {N} = @@ -318,13 +349,8 @@ module IteratorsMD # AbstractArray implementation Base.axes(iter::CartesianIndices{N,R}) where {N,R} = map(Base.axes1, iter.indices) Base.IndexStyle(::Type{CartesianIndices{N,R}}) where {N,R} = IndexCartesian() - @inline function Base.getindex(iter::CartesianIndices{N,<:NTuple{N,Base.OneTo}}, I::Vararg{Int, N}) where {N} - @boundscheck checkbounds(iter, I...) - CartesianIndex(I) - end - @inline function Base.getindex(iter::CartesianIndices{N,R}, I::Vararg{Int, N}) where {N,R} - @boundscheck checkbounds(iter, I...) - CartesianIndex(I .- first.(Base.axes1.(iter.indices)) .+ first.(iter.indices)) + @propagate_inbounds function Base.getindex(iter::CartesianIndices{N,R}, I::Vararg{Int, N}) where {N,R} + CartesianIndex(getindex.(iter.indices, I)) end ndims(R::CartesianIndices) = ndims(typeof(R)) @@ -344,47 +370,65 @@ module IteratorsMD IteratorSize(::Type{<:CartesianIndices{N}}) where {N} = Base.HasShape{N}() @inline function iterate(iter::CartesianIndices) - iterfirst, iterlast = first(iter), last(iter) - if any(map(>, iterfirst.I, iterlast.I)) + iterfirst = first(iter) + if !all(map(in, iterfirst.I, iter.indices)) return nothing end iterfirst, iterfirst end @inline function iterate(iter::CartesianIndices, state) - valid, I = __inc(state.I, first(iter).I, last(iter).I) + valid, I = __inc(state.I, iter.indices) valid || return nothing return CartesianIndex(I...), CartesianIndex(I...) end # increment & carry - @inline function inc(state, start, stop) - _, I = __inc(state, start, stop) + @inline function inc(state, indices) + _, I = __inc(state, indices) return CartesianIndex(I...) end - # increment post check to avoid integer overflow - @inline __inc(::Tuple{}, ::Tuple{}, ::Tuple{}) = false, () - @inline function __inc(state::Tuple{Int}, start::Tuple{Int}, stop::Tuple{Int}) - valid = state[1] < stop[1] - return valid, (state[1]+1,) + # Unlike ordinary ranges, CartesianIndices continues the iteration in the next column when the + # current column is consumed. The implementation is written recursively to achieve this. + # `iterate` returns `Union{Nothing, Tuple}`, we explicitly pass a `valid` flag to eliminate + # the type instability inside the core `__inc` logic, and this gives better runtime performance. + __inc(::Tuple{}, ::Tuple{}) = false, () + @inline function __inc(state::Tuple{Int}, indices::Tuple{<:OrdinalRange}) + rng = indices[1] + I = state[1] + step(rng) + valid = __is_valid_range(I, rng) && state[1] != last(rng) + return valid, (I, ) + end + @inline function __inc(state, indices) + rng = indices[1] + I = state[1] + step(rng) + if __is_valid_range(I, rng) && state[1] != last(rng) + return true, (I, tail(state)...) + end + valid, I = __inc(tail(state), tail(indices)) + return valid, (first(rng), I...) end - @inline function __inc(state, start, stop) - if state[1] < stop[1] - return true, (state[1]+1, tail(state)...) + @inline __is_valid_range(I, rng::AbstractUnitRange) = I in rng + @inline function __is_valid_range(I, rng::OrdinalRange) + if step(rng) > 0 + lo, hi = first(rng), last(rng) + else + lo, hi = last(rng), first(rng) end - valid, I = __inc(tail(state), tail(start), tail(stop)) - return valid, (start[1], I...) + lo <= I <= hi end # 0-d cartesian ranges are special-cased to iterate once and only once iterate(iter::CartesianIndices{0}, done=false) = done ? nothing : (CartesianIndex(), true) - size(iter::CartesianIndices) = map(dimlength, first(iter).I, last(iter).I) - dimlength(start, stop) = stop-start+1 + size(iter::CartesianIndices) = map(length, iter.indices) length(iter::CartesianIndices) = prod(size(iter)) + # make CartesianIndices a multidimensional range + Base.step(iter::CartesianIndices) = CartesianIndex(map(step, iter.indices)) + first(iter::CartesianIndices) = CartesianIndex(map(first, iter.indices)) last(iter::CartesianIndices) = CartesianIndex(map(last, iter.indices)) @@ -395,11 +439,8 @@ module IteratorsMD @inline to_indices(A, inds, I::Tuple{CartesianIndices{0},Vararg{Any}}) = (first(I), to_indices(A, inds, tail(I))...) - @inline function in(i::CartesianIndex{N}, r::CartesianIndices{N}) where {N} - _in(true, i.I, first(r).I, last(r).I) - end - _in(b, ::Tuple{}, ::Tuple{}, ::Tuple{}) = b - @inline _in(b, i, start, stop) = _in(b & (start[1] <= i[1] <= stop[1]), tail(i), tail(start), tail(stop)) + @inline in(i::CartesianIndex, r::CartesianIndices) = false + @inline in(i::CartesianIndex{N}, r::CartesianIndices{N}) where {N} = all(map(in, i.I, r.indices)) simd_outer_range(iter::CartesianIndices{0}) = iter function simd_outer_range(iter::CartesianIndices) @@ -410,8 +451,8 @@ module IteratorsMD simd_inner_length(iter::CartesianIndices, I::CartesianIndex) = Base.length(iter.indices[1]) simd_index(iter::CartesianIndices{0}, ::CartesianIndex, I1::Int) = first(iter) - @inline function simd_index(iter::CartesianIndices, Ilast::CartesianIndex, I1::Int) - CartesianIndex((I1+first(iter.indices[1]), Ilast.I...)) + @propagate_inbounds function simd_index(iter::CartesianIndices, Ilast::CartesianIndex, I1::Int) + CartesianIndex(getindex(iter.indices[1], I1+first(Base.axes1(iter.indices[1]))), Ilast.I...) end # Split out the first N elements of a tuple @@ -440,44 +481,79 @@ module IteratorsMD # reversed CartesianIndices iteration + Base.reverse(iter::CartesianIndices) = CartesianIndices(reverse.(iter.indices)) + @inline function iterate(r::Reverse{<:CartesianIndices}) - iterfirst, iterlast = last(r.itr), first(r.itr) - if any(map(<, iterfirst.I, iterlast.I)) + iterfirst = last(r.itr) + if !all(map(in, iterfirst.I, r.itr.indices)) return nothing end iterfirst, iterfirst end @inline function iterate(r::Reverse{<:CartesianIndices}, state) - valid, I = __dec(state.I, last(r.itr).I, first(r.itr).I) + valid, I = __dec(state.I, r.itr.indices) valid || return nothing return CartesianIndex(I...), CartesianIndex(I...) end # decrement & carry - @inline function dec(state, start, stop) - _, I = __dec(state, start, stop) + @inline function dec(state, indices) + _, I = __dec(state, indices) return CartesianIndex(I...) end # decrement post check to avoid integer overflow - @inline __dec(::Tuple{}, ::Tuple{}, ::Tuple{}) = false, () - @inline function __dec(state::Tuple{Int}, start::Tuple{Int}, stop::Tuple{Int}) - valid = state[1] > stop[1] - return valid, (state[1]-1,) + @inline __dec(::Tuple{}, ::Tuple{}) = false, () + @inline function __dec(state::Tuple{Int}, indices::Tuple{<:OrdinalRange}) + rng = indices[1] + I = state[1] - step(rng) + valid = __is_valid_range(I, rng) && state[1] != first(rng) + return valid, (I,) end - @inline function __dec(state, start, stop) - if state[1] > stop[1] - return true, (state[1]-1, tail(state)...) + @inline function __dec(state, indices) + rng = indices[1] + I = state[1] - step(rng) + if __is_valid_range(I, rng) && state[1] != first(rng) + return true, (I, tail(state)...) end - valid, I = __dec(tail(state), tail(start), tail(stop)) - return valid, (start[1], I...) + valid, I = __dec(tail(state), tail(indices)) + return valid, (last(rng), I...) end # 0-d cartesian ranges are special-cased to iterate once and only once iterate(iter::Reverse{<:CartesianIndices{0}}, state=false) = state ? nothing : (CartesianIndex(), true) - Base.LinearIndices(inds::CartesianIndices{N,R}) where {N,R} = LinearIndices{N,R}(inds.indices) + function Base.LinearIndices(inds::CartesianIndices{N,R}) where {N,R<:NTuple{N, AbstractUnitRange}} + LinearIndices{N,R}(inds.indices) + end + function Base.LinearIndices(inds::CartesianIndices) + indices = inds.indices + if all(x->step(x)==1, indices) + indices = map(rng->first(rng):last(rng), indices) + LinearIndices{length(indices), typeof(indices)}(indices) + else + # Given the fact that StepRange 1:2:4 === 1:2:3, we lost the original size information + # and thus cannot calculate the correct linear indices when the steps are not 1. + throw(ArgumentError("LinearIndices for $(typeof(inds)) with non-1 step size is not yet supported.")) + end + end + + # This is currently needed because converting to LinearIndices is only available when steps are + # all 1 + # NOTE: this is only a temporary patch and could be possibly removed when StepRange support to + # LinearIndices is done + function Base.collect(inds::CartesianIndices{N, R}) where {N,R<:NTuple{N, AbstractUnitRange}} + Base._collect_indices(axes(inds), inds) + end + function Base.collect(inds::CartesianIndices) + dest = Array{eltype(inds), ndims(inds)}(undef, size(inds)) + i = 0 + @inbounds for a in inds + dest[i+=1] = a + end + dest + end # array operations Base.intersect(a::CartesianIndices{N}, b::CartesianIndices{N}) where N = @@ -501,7 +577,7 @@ module IteratorsMD end @inline function iterate(iter::CartesianPartition, (state, n)) n >= length(iter) && return nothing - I = IteratorsMD.inc(state.I, first(iter.parent.parent).I, last(iter.parent.parent).I) + I = IteratorsMD.inc(state.I, iter.parent.parent.indices) return I, (I, n+1) end diff --git a/base/range.jl b/base/range.jl index 5f30a31a0bde5..2fd5866dc5d6b 100644 --- a/base/range.jl +++ b/base/range.jl @@ -960,6 +960,11 @@ AbstractUnitRange{T}(r::AbstractUnitRange{T}) where {T} = r AbstractUnitRange{T}(r::UnitRange) where {T} = UnitRange{T}(r) AbstractUnitRange{T}(r::OneTo) where {T} = OneTo{T}(r) +OrdinalRange{T1, T2}(r::StepRange) where {T1, T2<: Integer} = StepRange{T1, T2}(r) +OrdinalRange{T1, T2}(r::AbstractUnitRange{T1}) where {T1, T2<:Integer} = r +OrdinalRange{T1, T2}(r::UnitRange) where {T1, T2<:Integer} = UnitRange{T1}(r) +OrdinalRange{T1, T2}(r::OneTo) where {T1, T2<:Integer} = OneTo{T1}(r) + promote_rule(::Type{StepRange{T1a,T1b}}, ::Type{StepRange{T2a,T2b}}) where {T1a,T1b,T2a,T2b} = el_same(promote_type(T1a,T2a), # el_same only operates on array element type, so just promote second type parameter diff --git a/test/cartesian.jl b/test/cartesian.jl index e769d03ae8035..32c3e1fb0a4e6 100644 --- a/test/cartesian.jl +++ b/test/cartesian.jl @@ -7,6 +7,274 @@ ex = Base.Cartesian.exprresolve(:(if 5 > 4; :x; else :y; end)) @test Base.Cartesian.lreplace!("val_col", Base.Cartesian.LReplace{String}(:col, "col", 1)) == "val_1" @test Base.setindex(CartesianIndex(1,5,4),3,2) == CartesianIndex(1, 3, 4) +@testset "CartesianIndices constructions" begin + @testset "AbstractUnitRange" begin + for oinds in [ + (2, 3), + (UInt8(2), 3), + (2, UInt8(3)), + (2, 1:3), + (Base.OneTo(2), 1:3) + ] + R = CartesianIndices(oinds) + @test size(R) == (2, 3) + @test axes(R) == (Base.OneTo(2), Base.OneTo(3)) + @test step.(R.indices) == (1, 1) + @test step(R) == CartesianIndex(1, 1) + + @test R[begin] == CartesianIndex(1, 1) + @test R[2] == CartesianIndex(2, 1) + @test R[1, 2] == CartesianIndex(1, 2) + @test R[end] == CartesianIndex(2, 3) + end + @test CartesianIndices((2, 3)) == CartesianIndex(1, 1):CartesianIndex(2, 3) + + R = CartesianIndices((0:5, 0:5)) + @test R[begin] == R[1] == first(R) == CartesianIndex(0, 0) + @test R[2, 1] == R[2] == CartesianIndex(1, 0) + @test R[1, 2] == R[7] == CartesianIndex(0, 1) + @test R[end] == R[length(R)] == last(R) == CartesianIndex(5, 5) + + for oinds in [(2, ), (2, 3), (2, 3, 4)] + R = CartesianIndices(oinds) + @test eltype(R) == CartesianIndex{length(oinds)} + @test ndims(R) == length(oinds) + @test size(R) == oinds + end + + # generic iterators doesn't have axes interface + iter = Iterators.repeated([1 2], 4) + @test_throws MethodError CartesianIndices(iter) + end + + @testset "Step Range" begin + for oinds in [ + (2, 1:2:6), + (Base.OneTo(2), 1:2:6), + (UInt8(2), 1:2:6), + (2, UInt8(1):UInt8(2):UInt8(6)) + ] + R = CartesianIndices(oinds) + @test size(R) == (2, 3) + @test axes(R) == (Base.OneTo(2), Base.OneTo(3)) + @test step.(R.indices) == (1, 2) + @test step(R) == CartesianIndex(1, 2) + + @test R[begin] == CartesianIndex(1, 1) + @test R[2] == CartesianIndex(2, 1) + @test R[1, 2] == CartesianIndex(1, 3) + @test R[end] == CartesianIndex(2, 5) + end + + @test CartesianIndices((1:2:5, 1:3:7)) == CartesianIndex(1, 1):CartesianIndex(2,3):CartesianIndex(5,7) + + R = CartesianIndex(0, 0):CartesianIndex(2, 3):CartesianIndex(5, 7) + @test R[begin] == R[1] == first(R) == CartesianIndex(0, 0) + @test R[2, 1] == R[2] == CartesianIndex(2, 0) + @test R[1, 2] == R[4] == CartesianIndex(0, 3) + @test R[end] == R[length(R)] == last(R) == CartesianIndex(4, 6) + + for oinds in [(1:2:5, ), (1:2:5, 1:3:7), (1:2:5, 1:3:7, 1:4:11)] + R = CartesianIndices(oinds) + @test eltype(R) == CartesianIndex{length(oinds)} + @test ndims(R) == length(oinds) + @test size(R) == length.(oinds) + end + + R = CartesianIndices((1:2:5, 7:-3:1)) + @test R == CartesianIndex(1, 7):CartesianIndex(2,-3):CartesianIndex(5, 1) + @test step.(R.indices) == (2, -3) + @test R[begin] == R[1] == first(R) == CartesianIndex(1, 7) + @test R[2, 1] == R[2] == CartesianIndex(3, 7) + @test R[1, 2] == R[4] == CartesianIndex(1, 4) + @test R[end] == R[length(R)] == last(R) == CartesianIndex(5, 1) + end + + @testset "IdentityUnitRange" begin + function _collect(A) + rst = eltype(A)[] + for i in A + push!(rst, i) + end + rst + end + function _simd_collect(A) + rst = eltype(A)[] + @simd for i in A + push!(rst, i) + end + rst + end + + for oinds in [ + (Base.IdentityUnitRange(0:1),), + (Base.IdentityUnitRange(0:1), Base.IdentityUnitRange(0:2)), + (Base.IdentityUnitRange(0:1), Base.OneTo(3)), + ] + R = CartesianIndices(oinds) + @test axes(R) === oinds + @test _collect(R) == _simd_collect(R) == vec(collect(R)) + end + R = CartesianIndices((Base.IdentityUnitRange(0:1), 0:1)) + @test axes(R) == (Base.IdentityUnitRange(0:1), Base.OneTo(2)) + + end + + for oinds in [(2, 3), (0:1, 0:2), (0:1:1, 0:1:2), (Base.IdentityUnitRange(0:1), Base.IdentityUnitRange(0:2)) ] + R = CartesianIndices(oinds) + @test vec(LinearIndices(R)) == 1:6 + end + # TODO: non-1 steps are not supported yet, but may change in the future + @test_throws ArgumentError LinearIndices(CartesianIndices((1:2:5, ))) + @test_throws ArgumentError LinearIndices(CartesianIndices((1:1:5, 1:2:5))) +end + +module TestOffsetArray + isdefined(Main, :OffsetArrays) || @eval Main include("testhelpers/OffsetArrays.jl") + using .Main.OffsetArrays + using Test + + A = OffsetArray(rand(2, 3), -1, -1) + R = CartesianIndices(A) + @test R == CartesianIndices((0:1, 0:2)) + @test axes(R) == (0:1, 0:2) + for i in eachindex(A) + @test A[i] == A[R[i]] + end + for i in R + @test A[i] == A[Tuple(i)...] + end +end + +@testset "CartesianIndices getindex" begin + @testset "AbstractUnitRange" begin + for oinds in [(2, ), (2, 3), (2, 3, 4)] + A = rand(1:10, oinds) + R = CartesianIndices(A) + @test R == CartesianIndices(oinds) + + @test A[R] == A + @test axes(A) == axes(R) + @test all(i->A[i]==A[R[i]], eachindex(A)) + @test all(i->A[i]==A[R[i]], R) + @test all(i->A[i]==A[R[i]], collect(R)) + @test all(i->i in R, collect(R)) + end + end + + @testset "StepRange" begin + for oinds in [(1:2:5, ), (1:2:5, 1:3:7), (1:2:5, 1:3:7, 1:4:11)] + A = rand(1:10, last.(oinds)) + R = CartesianIndices(A) + + SR = CartesianIndex(first.(oinds)):CartesianIndex(step.(oinds)):CartesianIndex(last.(oinds)) + @test A[oinds...] == A[SR] + @test A[SR] == A[R[SR]] + + # TODO: A[SR] == A[Linearindices(SR)] should hold for StepRange CartesianIndices + @test_broken A[SR] == A[LinearIndices(SR)] + end + end +end + +@testset "range interface" begin + for (I, i, i_next) in [ + (CartesianIndices((1:2:5, )), CartesianIndex(2, ), CartesianIndex(4, )), + (1:2:5, 2, 4), + ] + # consistent with ranges behavior + @test !(i in I) + @test iterate(I, i) == (i_next, i_next) + end + + # check iteration behavior on boundary + R = CartesianIndex(1, 1):CartesianIndex(2, 3):CartesianIndex(4, 5) + @test R.indices == (1:2:3, 1:3:4) + i = CartesianIndex(4, 1) + i_next = CartesianIndex(1, 4) + @test !(i in R) && iterate(R, i) == (i_next, i_next) + + for R in [ + CartesianIndices((1:-1:-1, 1:2:5)), + CartesianIndices((2, 3)), + CartesianIndex(1, 2) .- CartesianIndices((1:-1:-1, 1:2:5)), + CartesianIndex(1, 2) .- CartesianIndices((2, 3)), + ] + Rc = collect(R) + @test all(map(==, R, Rc)) + end +end + +@testset "Cartesian simd/broadcasting" begin + @testset "AbstractUnitRange" begin + A = rand(-5:5, 64, 64) + @test abs.(A) == map(abs, A) + + function test_simd(f, @nospecialize(A); init=zero(eltype(A))) + val_simd = init + @simd for i in CartesianIndices(A) + val_simd = f(val_simd, A[i]) + end + + val_iter = init + for i in CartesianIndices(A) + val_iter = f(val_iter, A[i]) + end + + @test val_iter == reduce(f, A, init=init) + @test val_iter ≈ val_simd + end + + test_simd(+, A) + end + + R = CartesianIndex(-1, -1):CartesianIndex(6, 7) + @test R .+ CartesianIndex(1, 2) == CartesianIndex(0, 1):CartesianIndex(7, 9) + @test R .- CartesianIndex(1, 2) == CartesianIndex(-2, -3):CartesianIndex(5, 5) + # 37867: collect is needed + @test collect(CartesianIndex(1, 2) .- R) == CartesianIndex(2, 3):CartesianIndex(-1, -1):CartesianIndex(-5, -5) + + R = CartesianIndex(-1, -1):CartesianIndex(2, 3):CartesianIndex(6, 7) + @test R .+ CartesianIndex(2, 2) == CartesianIndex(1, 1):CartesianIndex(2, 3):CartesianIndex(8, 9) + @test R .- CartesianIndex(2, 2) == CartesianIndex(-3, -3):CartesianIndex(2, 3):CartesianIndex(4, 5) + # 37867: collect is needed + @test collect(CartesianIndex(1, 1) .- R) == CartesianIndex(2, 2):CartesianIndex(-2, -3):CartesianIndex(-4, -4) +end + +@testset "Iterators" begin + @testset "Reverse" begin + R = CartesianIndices((0:5, 0:5)) + RR = Iterators.Reverse(R) + rR = reverse(R) + @test rR == collect(RR) + @test rR.indices == (5:-1:0, 5:-1:0) + + @test eltype(RR) == CartesianIndex{2} + @test size(RR) == size(R) + @test axes(RR) == axes(R) + + @test first(RR) == last(R) == CartesianIndex(5, 5) + @test last(RR) == first(R) == CartesianIndex(0, 0) + RRR = collect(Iterators.Reverse(collect(RR))) + @test R == RRR + end + + @testset "collect" begin + for oinds in [(0:5, ), (2:2:7, ), (2:-1:0, ), + (0:5, 2:8), (2:2:7, 3:3:10), (2:-1:0, 2:7),] + R = CartesianIndices(oinds) + @test collect(R) == R + end + end +end + +@testset "set operations" begin + R1 = CartesianIndices((3, 4, 5)) + R2 = CartesianIndices((-2:2, -3:3, -4:4)) + R = CartesianIndices((2, 3, 4)) + @test intersect(R1, R2) == R +end + # test conversions for CartesianIndex @testset "CartesianIndex Conversions" begin @@ -19,25 +287,89 @@ ex = Base.Cartesian.exprresolve(:(if 5 > 4; :x; else :y; end)) end @testset "CartesianIndices overflow" begin - I = CartesianIndices((1:typemax(Int),)) - i = last(I) - @test iterate(I, i) === nothing + @testset "incremental steps" begin + I = CartesianIndices((1:typemax(Int),)) + i = last(I) + @test iterate(I, i) === nothing + + I = CartesianIndices((1:2:typemax(Int), )) + i = CartesianIndex(typemax(Int)-1) + @test iterate(I, i) === nothing + + I = CartesianIndices((1:(typemax(Int)-1),)) + i = CartesianIndex(typemax(Int)) + @test iterate(I, i) === nothing + + I = CartesianIndices((1:2:typemax(Int)-1, )) + i = CartesianIndex(typemax(Int)-1) + @test iterate(I, i) === nothing + + I = CartesianIndices((1:typemax(Int), 1:typemax(Int))) + i = last(I) + @test iterate(I, i) === nothing + + I = CartesianIndices((1:2:typemax(Int), 1:2:typemax(Int))) + i = CartesianIndex(typemax(Int)-1, typemax(Int)-1) + @test iterate(I, i) === nothing + + I = CartesianIndices((1:typemax(Int), 1:typemax(Int))) + i = CartesianIndex(typemax(Int), 1) + @test iterate(I, i) === (CartesianIndex(1, 2), CartesianIndex(1,2)) + + I = CartesianIndices((1:2:typemax(Int), 1:2:typemax(Int))) + i = CartesianIndex(typemax(Int)-1, 1) + @test iterate(I, i) === (CartesianIndex(1, 3), CartesianIndex(1, 3)) + + I = CartesianIndices((typemin(Int):(typemin(Int)+3),)) + i = last(I) + @test iterate(I, i) === nothing - I = CartesianIndices((1:(typemax(Int)-1),)) - i = CartesianIndex(typemax(Int)) - @test iterate(I, i) === nothing + I = CartesianIndices(((typemin(Int):2:typemin(Int)+3), )) + i = CartesianIndex(typemin(Int)+2) + @test iterate(I, i) === nothing + end + + @testset "decremental steps" begin + I = Iterators.Reverse(CartesianIndices((typemin(Int):typemin(Int)+10, ))) + i = last(I) + @test iterate(I, i) === nothing + + I = Iterators.Reverse(CartesianIndices((typemin(Int):2:typemin(Int)+10, ))) + i = last(I) + @test iterate(I, i) === nothing + + I = Iterators.Reverse(CartesianIndices((typemin(Int):typemin(Int)+10, ))) + i = CartesianIndex(typemin(Int)) + @test iterate(I, i) === nothing + + I = Iterators.Reverse(CartesianIndices((typemin(Int):2:typemin(Int)+10, ))) + i = CartesianIndex(typemin(Int)) + @test iterate(I, i) === nothing - I = CartesianIndices((1:typemax(Int), 1:typemax(Int))) - i = last(I) - @test iterate(I, i) === nothing + I = Iterators.Reverse(CartesianIndices((typemin(Int):typemin(Int)+10, typemin(Int):typemin(Int)+10))) + i = last(I) + @test iterate(I, i) === nothing - i = CartesianIndex(typemax(Int), 1) - @test iterate(I, i) === (CartesianIndex(1, 2), CartesianIndex(1,2)) + I = Iterators.Reverse(CartesianIndices((typemin(Int):2:typemin(Int)+10, typemin(Int):2:typemin(Int)+10))) + i = CartesianIndex(typemin(Int), typemin(Int)) + @test iterate(I, i) === nothing - # reverse cartesian indices - I = CartesianIndices((typemin(Int):(typemin(Int)+3),)) - i = last(I) - @test iterate(I, i) === nothing + I = Iterators.Reverse(CartesianIndices((typemin(Int):typemin(Int)+10, typemin(Int):typemin(Int)+10))) + i = CartesianIndex(typemin(Int), typemin(Int)+1) + @test iterate(I, i) === (CartesianIndex(typemin(Int)+10, typemin(Int)), CartesianIndex(typemin(Int)+10, typemin(Int))) + + I = Iterators.Reverse(CartesianIndices((typemin(Int):2:typemin(Int)+10, typemin(Int):2:typemin(Int)+10))) + i = CartesianIndex(typemin(Int), typemin(Int)+2) + @test iterate(I, i) === (CartesianIndex(typemin(Int)+10, typemin(Int)), CartesianIndex(typemin(Int)+10, typemin(Int))) + + I = CartesianIndices((typemax(Int):-1:typemax(Int)-10, )) + i = last(I) + @test iterate(I, i) === nothing + + I = CartesianIndices((typemax(Int):-2:typemax(Int)-10, )) + i = last(I) + @test iterate(I, i) === nothing + end end @testset "CartesianIndices iteration" begin