Skip to content

Commit

Permalink
fix overflow in CartesianIndices iteration
Browse files Browse the repository at this point in the history
This allows LLVM to vectorize the 1D CartesianIndices case,
as well as fixing an overflow bug for:

```julia
CartesianIndices(((typemax(Int64)-2):typemax(Int64),))
```

Co-authored-by: Yingbo Ma <[email protected]>
  • Loading branch information
vchuravy and YingboMa committed Feb 8, 2019
1 parent a0bb006 commit 1aac477
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 2 deletions.
7 changes: 5 additions & 2 deletions base/multidimensional.jl
Original file line number Diff line number Diff line change
Expand Up @@ -334,9 +334,12 @@ module IteratorsMD
iterfirst, iterfirst
end
@inline function iterate(iter::CartesianIndices, state)
# If we increment before the condition check, we run the
# risk of an integer overflow.
idxend = map(last, iter.indices)
all(map(>=, state.I, idxend)) && return nothing
nextstate = CartesianIndex(inc(state.I, first(iter).I, last(iter).I))
nextstate.I[end] > last(iter.indices[end]) && return nothing
nextstate, nextstate
return nextstate, nextstate
end

# increment & carry
Expand Down
27 changes: 27 additions & 0 deletions test/cartesian.jl
Original file line number Diff line number Diff line change
Expand Up @@ -15,3 +15,30 @@ ex = Base.Cartesian.exprresolve(:(if 5 > 4; :x; else :y; end))
# can't convert higher-dimensional indices to Int
@test_throws MethodError convert(Int, CartesianIndex(42, 1))
end

@testset "CartesianIndices overflow" begin
I = CartesianIndices((1:typemax(Int64),))
i = last(I)
@test iterate(I, i) === nothing

I = CartesianIndices((1:(typemax(Int64)-1),))
i = CartesianIndex(typemax(Int64))
@test iterate(I, i) === nothing

I = CartesianIndices((1:typemax(Int64), 1:typemax(Int64)))
i = last(I)
@test iterate(I, i) === nothing

i = CartesianIndex(typemax(Int64), 1)
@test iterate(I, i) === (CartesianIndex(1, 2), CartesianIndex(1,2))
end

@testset "CartesianRange iteration" begin
I = CartesianIndices(map(Base.Slice, (2:4, 0:1, 1:1, 3:5)))
indices = Vector{eltype(I)}()
for i in I
push!(indices, i)
end
@test length(I) == length(indices)
@test vec(I) == indices
end

0 comments on commit 1aac477

Please sign in to comment.