From eebadbcfe98bd7ffc4a549f051bf95642c038014 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Fri, 8 Feb 2019 10:56:47 -0500
Subject: [PATCH] fix overflow in CartesianIndices iteration

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 <mayingbo5@gmail.com>
---
 base/multidimensional.jl |  7 +++++--
 test/cartesian.jl        | 27 +++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/base/multidimensional.jl b/base/multidimensional.jl
index f03e641af10fc..648943dc7457a 100644
--- a/base/multidimensional.jl
+++ b/base/multidimensional.jl
@@ -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
diff --git a/test/cartesian.jl b/test/cartesian.jl
index 40badf6bb24bb..e630159824405 100644
--- a/test/cartesian.jl
+++ b/test/cartesian.jl
@@ -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(Int),))
+    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((1:typemax(Int), 1:typemax(Int)))
+    i = last(I)
+    @test iterate(I, i) === nothing
+
+    i = CartesianIndex(typemax(Int), 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