various fixes

* correct overflow checks * correct broadcasting
JuliaLang · Oct 3, 2020 · 4b47f63 · 4b47f63
1 parent 113201f
commit 4b47f63
Show file tree

Hide file tree

Showing 3 changed files with 185 additions and 42 deletions.
diff --git 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
@@ -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
@@ -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