From 45cf35ae1b0090c14375d50f254885c31ca10649 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Fri, 14 Jun 2024 09:41:58 +0200
Subject: [PATCH] Add test

---
 src/arithmetic.jl      | 31 ++++++++++++++++---------------
 src/mstructures.jl     | 13 ++++++++++---
 src/mtables.jl         |  3 +++
 test/monoid_algebra.jl |  1 +
 4 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/src/arithmetic.jl b/src/arithmetic.jl
index 12d64ed..2a45339 100644
--- a/src/arithmetic.jl
+++ b/src/arithmetic.jl
@@ -1,14 +1,12 @@
-function _preallocate_output(X::AlgebraElement, a::Number, op)
-    T = MA.promote_operation(op, eltype(X), typeof(a))
-    return similar(X, T)
-end
+_coeff_type(X::AlgebraElement) = eltype(X)
+_coeff_type(a) = typeof(a)
 
-function _preallocate_output(X::AlgebraElement, Y::AlgebraElement, op)
-    T = MA.promote_operation(op, eltype(X), eltype(Y))
-    if coeffs(Y) isa DenseArray # what a hack :)
-        return similar(Y, T)
+function _preallocate_output(op, args::Vararg{Any,N}) where {N}
+    T = MA.promote_operation(op, _coeff_type.(args)...)
+    if args[2] isa AlgebraElement && coeffs(args[2]) isa DenseArray # what a hack :)
+        return similar(args[2], T)
     end
-    return similar(X, T)
+    return similar(args[1], T)
 end
 
 # module structure:
@@ -18,23 +16,26 @@ Base.:(/)(X::AlgebraElement, a::Number) = inv(a) * X
 Base.:(//)(X::AlgebraElement, a::Number) = 1 // a * X
 
 function Base.:-(X::AlgebraElement)
-    return MA.operate_to!(_preallocate_output(X, -1, *), -, X)
+    return MA.operate_to!(_preallocate_output(*, X, -1), -, X)
 end
 function Base.:*(a::Number, X::AlgebraElement)
-    return MA.operate_to!(_preallocate_output(X, a, *), *, X, a)
+    return MA.operate_to!(_preallocate_output(*, X, a), *, X, a)
 end
 function Base.:div(X::AlgebraElement, a::Number)
-    return MA.operate_to!(_preallocate_output(X, a, div), div, X, a)
+    return MA.operate_to!(_preallocate_output(div, X, a), div, X, a)
 end
 
 function Base.:+(X::AlgebraElement, Y::AlgebraElement)
-    return MA.operate_to!(_preallocate_output(X, Y, +), +, X, Y)
+    return MA.operate_to!(_preallocate_output(+, X, Y), +, X, Y)
 end
 function Base.:-(X::AlgebraElement, Y::AlgebraElement)
-    return MA.operate_to!(_preallocate_output(X, Y, -), -, X, Y)
+    return MA.operate_to!(_preallocate_output(-, X, Y), -, X, Y)
 end
 function Base.:*(X::AlgebraElement, Y::AlgebraElement)
-    return MA.operate_to!(_preallocate_output(X, Y, *), *, X, Y)
+    return MA.operate_to!(_preallocate_output(*, X, Y), *, X, Y)
+end
+function Base.:*(args::Vararg{AlgebraElement,N}) where {N}
+    return MA.operate_to!(_preallocate_output(*, args...), *, args...)
 end
 Base.:^(a::AlgebraElement, p::Integer) = Base.power_by_squaring(a, p)
 
diff --git a/src/mstructures.jl b/src/mstructures.jl
index 70d220e..1fab9d1 100644
--- a/src/mstructures.jl
+++ b/src/mstructures.jl
@@ -54,17 +54,24 @@ function MA.operate_to!(res, ms::MultiplicativeStructure, args::Vararg{Any,N}) w
     return res
 end
 
-function operate!(::UnsafeAddMul, res, c)
+function MA.operate!(::UnsafeAddMul, res, c)
     for (k, v) in nonzero_pairs(c)
         unsafe_push!(res, k, v)
     end
     return res
 end
 
-function operate!(op::UnsafeAddMul, res, b, c, args::Vararg{Any, N}) where {N}
+function MA.operate!(op::UnsafeAddMul, res, b, c, args::Vararg{Any, N}) where {N}
     for (kb, vb) in nonzero_pairs(b)
         for (kc, vc) in nonzero_pairs(c)
-            operate!(op, res, SparseCoefficients((op.structure(kb, kc),), (vb * vc,)), args...)
+            for (k, v) in nonzero_pairs(op.structure(kb, kc))
+                MA.operate!(
+                    op,
+                    res,
+                    SparseCoefficients((_key(op.structure, k),), (vb * vc * v,)),
+                    args...,
+                )
+            end
         end
     end
     return res
diff --git a/src/mtables.jl b/src/mtables.jl
index 6b9fac0..20b5bee 100644
--- a/src/mtables.jl
+++ b/src/mtables.jl
@@ -92,6 +92,9 @@ function complete!(mt::MTable)
     return mt
 end
 
+_key(_, k) = k
+_key(mstr::MTable, k) = mstr[k]
+
 function MA.operate!(
     ms::UnsafeAddMul{<:MTable},
     res::AbstractCoefficients,
diff --git a/test/monoid_algebra.jl b/test/monoid_algebra.jl
index 92b22cc..92f376c 100644
--- a/test/monoid_algebra.jl
+++ b/test/monoid_algebra.jl
@@ -18,6 +18,7 @@
     @test iszero(zero(fRG))
     @test zero(g) == zero(fRG)
     @test iszero(0 * g)
+    @test isone(*(g, g, g))
 
     @testset "Translations between bases" begin
         Z = zero(RG)