From 74d8d1bfa5f16ff9566a04ef89558ffb4baee62e Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Wed, 15 May 2024 15:51:39 +0200
Subject: [PATCH] fix iteration on AugmentedBasis and add tests

---
 src/diracs_augmented.jl  | 13 +++++++++++--
 test/perm_grp_algebra.jl |  8 ++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/src/diracs_augmented.jl b/src/diracs_augmented.jl
index 09b887a..904a874 100644
--- a/src/diracs_augmented.jl
+++ b/src/diracs_augmented.jl
@@ -51,6 +51,7 @@ struct AugmentedBasis{T,I,A<:Augmented{T},B<:AbstractBasis{T,I}} <:
 end
 
 function AugmentedBasis(basis::DiracBasis{T,I}) where {T,I}
+    @assert one(object(basis)) in basis
     return AugmentedBasis{T,I,Augmented{T},typeof(basis)}(basis)
 end
 
@@ -67,10 +68,18 @@ function Base.length(ab::AugmentedBasis)
 end
 
 function Base.iterate(ab::AugmentedBasis)
-    return ((v, st) = iterate(object(ab)); (Augmented(v), st))
+    (v, st) = iterate(object(ab))
+    isone(v) && return iterate(ab, st)
+    return Augmented(v), st
 end
+
 function Base.iterate(ab::AugmentedBasis, st)
-    return ((v, st) = iterate(object(ab), st); (Augmented(v), st))
+    (v, st) = let k = iterate(object(ab), st)
+        isnothing(k) && return nothing
+        k
+    end
+    isone(v) && return iterate(ab, st)
+    return Augmented(v), st
 end
 
 Base.in(g, ab::AugmentedBasis) = false
diff --git a/test/perm_grp_algebra.jl b/test/perm_grp_algebra.jl
index 950431f..12013a0 100644
--- a/test/perm_grp_algebra.jl
+++ b/test/perm_grp_algebra.jl
@@ -34,6 +34,10 @@
         ad = SA.AugmentedBasis(db)
         @test SA.mstructure(ad) == SA.AugmentedMStructure(SA.mstructure(db))
         @test ad[SA.Augmented(h)] isa SA.Augmented
+        @test sprint(show, ad[SA.Augmented(h)]) == "(-1·()+1·(1,2,4,5))"
+
+        @test !(h in ad)
+        @test SA.Augmented(h) in ad
 
         IG = SA.StarAlgebra(G, ad)
 
@@ -47,6 +51,10 @@
         @test coeffs(ax * ay) == SA.coeffs(x * y, basis(IG))
         @test coeffs(ax * az) == SA.coeffs(x * z, basis(IG))
         @test SA.aug(ax) == 0
+        @test star(ax) * star(ay) == star(ay) * star(ax)
+
+        @test length(ad) == length(db) - 1
+        @test Set(ad) == Set(SA.Augmented(g) for g in db if !isone(g))
     end
 
     @testset "Random elements" begin