clean-up iteration

Project.toml

@@ -5,10 +5,10 @@ version = "0.5.2"
 
 [deps]
 AbstractAlgebra = "c3fe647b-3220-5bb0-a1ea-a7954cac585d"
-DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 GroupsCore = "d5909c97-4eac-4ecc-a3dc-fdd0858a4120"
 KnuthBendix = "c2604015-7b3d-4a30-8a26-9074551ec60a"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ThreadsX = "ac1d9e8a-700a-412c-b207-f0111f4b6c0d"
 

src/Groups.jl

@@ -23,7 +23,7 @@ include("arithmetic.jl")
 include("findreplace.jl")
 
 module New
-using DataStructures
+import OrderedCollections: OrderedSet
 
 include("new_types.jl")
 include("new_hashing.jl")

src/iteration.jl

@@ -1,53 +1,55 @@
-mutable struct FPIterState{GEl,T}
-    elts::OrderedSet{GEl}
-    u::GEl
-    v::GEl
-    elts_iter_state::T
-end
-
-Base.in(x, itr::FPIterState) = x in itr.elts
-Base.push!(itr::FPIterState, x) = push!(itr.elts, x)
-function FPIterState(G::AbstractFPGroup)
-    S = OrderedSet([one(G)])
-    elt, state = iterate(S)
-    return elt, FPIterState(S, one(G), one(G), state)
+mutable struct FPGroupIter{S, T, GEl}
+    seen::S
+    seen_iter_state::T
+    current::GEl
+    gen_idx::Int
+    u_tmp::GEl
+    v_tmp::GEl
 end
 
 function Base.iterate(G::AbstractFPGroup)
-    elt, state = FPIterState(G)
-    return one(G), (state, elt, 1)
+    seen = OrderedSet([one(G)])
+    current, seen_state = iterate(seen)
+    gr_iter = FPGroupIter(seen, seen_state, current, 1, one(G), one(G))
+    return one(G), gr_iter
+end
+
+function _next_elt!(itr::FPGroupIter)
+    res = iterate(itr.seen, itr.seen_iter_state)
+    res === nothing && return true
+
+    itr.current = first(res)
+    itr.seen_iter_state = last(res)
+    itr.gen_idx = 1
+
+    return false
 end
 
 function Base.iterate(G::AbstractFPGroup, state)
-    iter, elt, gen_idx = state
-
-    if gen_idx > length(alphabet(G))
-        res = iterate(iter.elts, iter.elts_iter_state)
-        res === nothing && return nothing
-        gen_idx = 1
-        elt = first(res)
-        iter.elts_iter_state = last(res)
+    if state.gen_idx > length(alphabet(G))
+        finished = _next_elt!(state)
+        finished && return nothing
     end
 
-    res = let (u, v) = (iter.u, iter.v), elt = elt
-        copyto!(v, elt) # this invalidates normalform of v
-        push!(word(v), gen_idx)
-        resize!(word(u), 0)
-
+    elt = let u = state.u_tmp, v = state.v_tmp, current = state.current
+        copyto!(v, current)
+        push!(word(v), state.gen_idx)
         _setnormalform!(v, false)
         _setvalidhash!(v, false)
-
         @assert !isnormalform(v)
+
+        resize!(word(u), 0)
         normalform!(u, v)
     end
 
-    if res in iter
-        return iterate(G, (iter, elt, gen_idx + 1))
+    state.gen_idx += 1
+
+    if elt in state.seen
+        return iterate(G, state)
     else
-        w = deepcopy(res)
-        @assert isnormalform(w)
-        push!(iter, w)
-        return w, (iter, elt, gen_idx + 1)
+        @assert isnormalform(elt)
+        push!(state.seen, deepcopy(elt))
+        return elt, state
     end
 end
 

test/free_groups.jl

@@ -33,11 +33,12 @@
             return w
         end
 
-        k = test_iteration(F3, 10)
+        @time k = test_iteration(F3, 10)
         @test k == a*b^-1
 
-        @time k = test_iteration(F3, 1000)
+        k = test_iteration(F3, 1000)
         @test k == (a^2)*c^2*a^-1
+        @time k = test_iteration(F3, 1000)
     end
 
     @testset "wl_ball" begin
@@ -48,7 +49,7 @@
                 isnothing(res) && break
                 g, state = res
             end
-            elts = collect(first(state).elts)
+            elts = collect(state.seen)
             elts = resize!(elts, length(elts)-1)
             return elts
         end