From 8333be5d2dd48729d0eceb59c6555e70ff3ea23e Mon Sep 17 00:00:00 2001
From: filifa <filifa.breath652@8alias.com>
Date: Thu, 21 Aug 2025 19:52:05 -0400
Subject: [PATCH] have lib function take map instead of list

---
 cmd/primitiveRoot.go | 20 ++++++++++++++++++--
 internal/lib/lib.go  | 15 +++++++--------
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/cmd/primitiveRoot.go b/cmd/primitiveRoot.go
index 2a799c7..80a7b48 100644
--- a/cmd/primitiveRoot.go
+++ b/cmd/primitiveRoot.go
@@ -28,21 +28,37 @@ import (
 var modulus string
 var tpf []string
 
+func count(tpf []string) map[string]*big.Int {
+	counter := make(map[string]*big.Int)
+	for _, s := range tpf {
+		count, ok := counter[s]
+		if ok {
+			count.Add(count, big.NewInt(1))
+		} else {
+			counter[s] = big.NewInt(1)
+		}
+	}
+
+	return counter
+}
+
 func primitiveRoot(cmd *cobra.Command, args []string) {
 	m, ok := new(big.Int).SetString(modulus, 10)
 	if !ok {
 		log.Fatal("invalid input " + modulus)
 	}
 
+	factors := count(tpf)
+
 	root := new(big.Int)
 	var err error
-	if len(tpf) == 0 {
+	if len(factors) == 0 {
 		root, err = lib.PrimitiveRoot(m)
 		if err != nil {
 			log.Fatal(err)
 		}
 	} else {
-		root, err = lib.PrimitiveRootFast(m, tpf)
+		root, err = lib.PrimitiveRootFast(m, factors)
 		if err != nil {
 			log.Fatal(err)
 		}
diff --git a/internal/lib/lib.go b/internal/lib/lib.go
index 9718f7f..2297625 100644
--- a/internal/lib/lib.go
+++ b/internal/lib/lib.go
@@ -141,17 +141,16 @@ func PrimitiveRoot(modulus *big.Int) (*big.Int, error) {
 	return nil, errors.New("no primitive root")
 }
 
-func PrimitiveRootFast(modulus *big.Int, tpf []string) (*big.Int, error) {
+func PrimitiveRootFast(modulus *big.Int, tpf map[string]*big.Int) (*big.Int, error) {
 	phi := big.NewInt(1)
-	factors := make(map[string]bool)
-	for _, s := range tpf {
-		p, ok := new(big.Int).SetString(s, 10)
+	for p, exp := range tpf {
+		pow, ok := new(big.Int).SetString(p, 10)
 		if !ok {
-			return nil, errors.New("invalid input " + s)
+			return nil, errors.New("invalid factor " + p)
 		}
 
-		phi.Mul(phi, p)
-		factors[p.Text(10)] = true
+		pow.Exp(pow, exp, nil)
+		phi.Mul(phi, pow)
 	}
 
 	for g := big.NewInt(1); g.Cmp(modulus) == -1; g.Add(g, big.NewInt(1)) {
@@ -161,7 +160,7 @@ func PrimitiveRootFast(modulus *big.Int, tpf []string) (*big.Int, error) {
 		}
 
 		isPrimitive := true
-		for p := range factors {
+		for p := range tpf {
 			e := new(big.Int)
 			f, _ := new(big.Int).SetString(p, 10)
 			k := new(big.Int).Div(phi, f)