From e6d67ca3f78e3c91ff90841e79fc44c9674668c7 Mon Sep 17 00:00:00 2001
From: kalmarek <kalmar@amu.edu.pl>
Date: Wed, 8 Aug 2018 14:12:55 +0200
Subject: [PATCH] add DirectPowerIter struct to iterate over DirectProduct

---
 src/DirectProducts.jl | 41 ++++++++++++++++++++++++++++++++++-------
 src/WreathProducts.jl |  5 +++--
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/src/DirectProducts.jl b/src/DirectProducts.jl
index ff7510c..97c1273 100644
--- a/src/DirectProducts.jl
+++ b/src/DirectProducts.jl
@@ -277,18 +277,45 @@ end
 #
 ###############################################################################
 
+import Base: size, length, start, next, done, eltype
+
+struct DirectPowerIter{Gr<:AbstractAlgebra.Group, GrEl<:AbstractAlgebra.GroupElem}
+   G::Gr
+   N::Int
+   elts::Vector{GrEl}
+   totalorder::Int
+   orderG::Int
+end
+
+function DirectPowerIter(G::Gr, N::Integer) where {Gr<:AbstractAlgebra.Group}
+   return DirectPowerIter{Gr, elem_type(G)}(
+      G,
+      N,
+      vec(collect(elements(G))),
+      Int(order(G))^N,
+      Int(order(G))
+      )
+end
+
+Base.size(DPIter::DirectPowerIter) = ntuple(i -> DPIter.orderG, DPIter.N)
+Base.length(DPIter::DirectPowerIter) = DPIter.totalorder
+Base.start(::DirectPowerIter) = 0
+
+function Base.next(DPIter::DirectPowerIter, state)
+    idx = ind2sub(size(DPIter), state+1)
+    return DirectProductGroupElem([DPIter.elts[i] for i in idx]), state+1
+end
+
+Base.done(DPIter::DirectPowerIter, state) = DPIter.totalorder <= state
+
+Base.eltype(::Type{DirectPowerIter{Gr, GrEl}}) where {Gr, GrEl} = DirectProductGroupElem{GrEl}
+
 doc"""
     elements(G::DirectProductGroup)
 > Returns `generator` that produces all elements of group `G` (provided that
 > `G.group` implements the `elements` method).
 """
-# TODO: can Base.product handle generators?
-#   now it returns nothing's so we have to collect ellements...
-function elements(G::DirectProductGroup)
-   elts = collect(elements(G.group))
-   cartesian_prod = Base.product([elts for _ in 1:G.n]...)
-   return (DirectProductGroupElem([elt...]) for elt in cartesian_prod)
-end
+elements(G::DirectProductGroup) = DirectPowerIter(G.group, G.n)
 
 doc"""
     order(G::DirectProductGroup)
diff --git a/src/WreathProducts.jl b/src/WreathProducts.jl
index 1e03ae7..1f80d4d 100644
--- a/src/WreathProducts.jl
+++ b/src/WreathProducts.jl
@@ -201,8 +201,9 @@ end
 matrix_repr(g::WreathProductElem) = Any[matrix_repr(g.p) g.n]
 
 function elements(G::WreathProduct)
-   iter = Base.product(collect(elements(G.N)), collect(elements(G.P)))
-   return (WreathProductElem(n, p, false) for (n,p) in iter)
+   Nelts = collect(elements(G.N))
+   Pelts = collect(elements(G.P))
+   return (WreathProductElem(n, p, false) for n in Nelts, p in Pelts)
 end
 
 order(G::WreathProduct) = order(G.P)*order(G.N)