From d87736e5732d79e1fb6616dc1dd2c5ae03b533cc Mon Sep 17 00:00:00 2001
From: John Lapeyre <jlapeyre@users.noreply.github.com>
Date: Sat, 1 Dec 2018 17:34:58 +0100
Subject: [PATCH] Implement faster "iterate" for Eye (#45)

---
 src/FillArrays.jl | 8 +++++++-
 test/runtests.jl  | 9 ++++++++-
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/FillArrays.jl b/src/FillArrays.jl
index eebe6277..5d4f5a95 100644
--- a/src/FillArrays.jl
+++ b/src/FillArrays.jl
@@ -200,7 +200,13 @@ const Eye{T, Axes} = Diagonal{T, Ones{T,1,Tuple{Axes}}}
 Eye{T}(n::Integer) where T = Diagonal(Ones{T}(n))
 Eye(n::Integer) = Diagonal(Ones(n))
 
-
+function Base.iterate(iter::Eye{T}, istate = (1, 1)) where T
+    (i::Int, j::Int) = istate
+    m = size(iter, 1)
+    return i > m ? nothing :
+        ((@inbounds getindex(iter, i, j)),
+         j == m ? (i + 1, 1) : (i, j + 1))
+end
 
 @deprecate Eye(n::Integer, m::Integer) view(Eye(max(n,m)), 1:n, 1:m)
 @deprecate Eye{T}(n::Integer, m::Integer) where T view(Eye{T}(max(n,m)), 1:n, 1:m)
diff --git a/test/runtests.jl b/test/runtests.jl
index 2f142af1..a1ac6e2f 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -35,7 +35,7 @@ import FillArrays: AbstractFill
                 @test convert(AbstractArray,Z) ≡ convert(AbstractFill,Z) ≡ Z
                 @test convert(AbstractArray{T},Z) ≡ convert(AbstractFill{T},Z) ≡ AbstractArray{T}(Z) ≡ Z
                 @test convert(AbstractMatrix{T},Z) ≡ convert(AbstractFill{T,2},Z) ≡ AbstractMatrix{T}(Z) ≡ Z
-                
+
                 @test_throws Exception convert(Fill{Float64}, [1,1,2])
                 @test_throws Exception convert(Fill, [])
                 @test convert(Fill{Float64}, [1,1,1]) ≡ Fill(1.0, 3)
@@ -469,3 +469,10 @@ end
     @test (1:10) .* Zeros{Int}(10) ≡ Zeros{Int}(10)
     @test_throws DimensionMismatch (1:11) .* Zeros(10)
 end
+
+@testset "iterate" begin
+    for d in (0, 1, 2, 100)
+        m = Eye(d)
+        @test [x for x in m] == m
+    end
+end