diff --git a/base/broadcast.jl b/base/broadcast.jl index 12bddaf531e27..b1e708612261f 100644 --- a/base/broadcast.jl +++ b/base/broadcast.jl @@ -1112,7 +1112,8 @@ 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)) + # TODO: 8ns overhead by adding step=step(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 608058bfb9ba3..76b30c4dcb471 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -368,38 +368,69 @@ module IteratorsMD eltype(::Type{CartesianIndices{N,TT}}) where {N,TT} = CartesianIndex{N} 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)) + @propagate_inbounds function iterate(iter::CartesianIndices) + iterfirst = first(iter) + if !all(map(in, iterfirst.I, iter.indices)) return nothing end iterfirst, iterfirst end - @inline function iterate(iter::CartesianIndices, state) + @propagate_inbounds function iterate(iter::CartesianIndices, state) valid, I = __inc(state.I, iter.indices) valid || return nothing return CartesianIndex(I...), CartesianIndex(I...) end # increment & carry - @inline function inc(state, indices) + @propagate_inbounds function inc(state, indices) _, I = __inc(state, indices) return CartesianIndex(I...) end __inc(::Tuple{}, ::Tuple{}) = false, () - @inline function __inc(state::Tuple{Int}, indices::Tuple{<:OrdinalRange}) - I = state[1] + step(indices[1]) - valid = I <= last(indices[1]) && typemax(Int) - step(indices[1]) >= state[1] + @propagate_inbounds function __inc(state::Tuple{Int}, indices::Tuple{<:OrdinalRange}) + rng = indices[1] + I = state[1] + step(rng) + valid = __isvalid_inc_range(I, rng) && __has_inc_overflow(state[1], rng) return valid, (I, ) end - @inline function __inc(state, indices) - I = state[1] + step(indices[1]) - if I <= last(indices[1]) && typemax(Int) - step(indices[1]) >= state[1] + @propagate_inbounds function __inc(state, indices) + rng = indices[1] + I = state[1] + step(rng) + if __isvalid_inc_range(I, rng) && __has_inc_overflow(state[1], rng) return true, (I, tail(state)...) end valid, I = __inc(tail(state), tail(indices)) - return valid, (first(indices[1]), I...) + return valid, (first(rng), I...) + end + + @inline __isvalid_inc_range(I, rng) = __isvalid_range(I, rng) + @inline __isvalid_inc_range(I, rng::AbstractUnitRange) = I <= last(rng) + + @propagate_inbounds __has_inc_overflow(I, rng) = __has_overflow(I, rng) + @inline function __has_inc_overflow(I::T, rng::AbstractUnitRange) where T + @boundscheck return I < last(rng) + return true + end + + @inline function __isvalid_range(I, rng) + if step(rng) > 0 + lo, hi = first(rng), last(rng) + elseif step(rng) < 0 + lo, hi = last(rng), first(rng) + end + lo <= I <= hi + end + @inline function __has_overflow(I::T, rng) where T + @boundscheck begin + inc_step = step(rng) + if inc_step > 0 && I > typemax(T) - inc_step + return false + elseif inc_step < 0 && I < typemin(T) - inc_step + return false + end + end + return true end # 0-d cartesian ranges are special-cased to iterate once and only once @@ -463,45 +494,57 @@ module IteratorsMD 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)) + @propagate_inbounds function iterate(r::Reverse{<:CartesianIndices}) + 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) + @propagate_inbounds function iterate(r::Reverse{<:CartesianIndices}, state) valid, I = __dec(state.I, r.itr.indices) valid || return nothing return CartesianIndex(I...), CartesianIndex(I...) end # decrement & carry - @inline function dec(state, indices) + @propagate_inbounds function dec(state, indices) _, I = __dec(state, indices) return CartesianIndex(I...) end # decrement post check to avoid integer overflow - @inline __dec(::Tuple{}, ::Tuple{}) = false, () - @inline function __dec(state::Tuple{Int}, indices::Tuple{<:OrdinalRange}) - I = state[1] - step(indices[1]) - valid = I >= first(indices[1]) && typemin(Int) + step(indices[1]) <= state[1] + @propagate_inbounds __dec(::Tuple{}, ::Tuple{}) = false, () + @propagate_inbounds function __dec(state::Tuple{Int}, indices::Tuple{<:OrdinalRange}) + rng = indices[1] + I = state[1] - step(rng) + valid = __isvalid_dec_range(I, rng) && __has_dec_overflow(I, rng) return valid, (I,) end - @inline function __dec(state, indices) - I = state[1] - step(indices[1]) - if I >= first(indices[1]) && typemin(Int) + step(indices[1]) <= state[1] + @propagate_inbounds function __dec(state, indices) + rng = indices[1] + I = state[1] - step(rng) + if __isvalid_dec_range(I, rng) && __has_dec_overflow(I, rng) return true, (I, tail(state)...) end valid, I = __dec(tail(state), tail(indices)) - return valid, (last(indices[1]), I...) + return valid, (last(rng), I...) + end + + @inline __isvalid_dec_range(I, rng) = __isvalid_range(I, rng) + @inline __isvalid_dec_range(I, rng::AbstractUnitRange) = I >= first(rng) + + @propagate_inbounds __has_dec_overflow(I, rng) = __has_overflow(I, rng) + @inline function __has_dec_overflow(I::T, rng::AbstractUnitRange) where T + @boundscheck return I <= last(rng) + return true 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) + # FIXME: A[SR] == A[Linearindices[SR]] should hold for StepRange CartesianIndices Base.LinearIndices(inds::CartesianIndices{N,R}) where {N,R} = LinearIndices(axes(inds)) # array operations diff --git a/test/cartesian.jl b/test/cartesian.jl index 26b7899363824..3ff31064886d2 100644 --- a/test/cartesian.jl +++ b/test/cartesian.jl @@ -130,11 +130,34 @@ end SR = CartesianIndex(first.(oinds)):CartesianIndex(step.(oinds)):CartesianIndex(last.(oinds)) @test A[oinds...] == A[SR] @test A[SR] == A[R[SR]] + # FIXME: A[SR] == A[Linearindices[SR]] should hold for StepRange CartesianIndices + @test_broken A[SR] == A[LinearIndices[SR]] end end end -@testset "Cartesian simd" begin +@testset "range interface" begin + for (I, i, i_next) in [ + (CartesianIndices((1:2:5, )), CartesianIndex(2, ), CartesianIndex(4, )), + (1:2:5, 2, 4), + ] + # consistant to ranges + @test !(i in I) + @test iterate(I, i) == (i_next, i_next) + end + + 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) @@ -156,6 +179,18 @@ 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 @@ -195,25 +230,89 @@ end end @testset "CartesianIndices overflow" begin - I = CartesianIndices((1:typemax(Int),)) - i = last(I) - @test iterate(I, i) === nothing + @testset "incremental" 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(((typemin(Int):2:typemin(Int)+3), )) + i = CartesianIndex(typemin(Int)+2) + @test iterate(I, i) === nothing + end - I = CartesianIndices((1:(typemax(Int)-1),)) - i = CartesianIndex(typemax(Int)) - @test iterate(I, i) === nothing + @testset "decremental" begin + I = Iterators.Reverse(CartesianIndices((typemin(Int):typemin(Int)+10, ))) + i = last(I) + @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):2: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):typemin(Int)+10, ))) + i = CartesianIndex(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):2:typemin(Int)+10, ))) + i = CartesianIndex(typemin(Int)) + @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 = 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 + + 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