From 6fd078925e750717abd8f65e07664f3c85cc950c Mon Sep 17 00:00:00 2001
From: Sam Wilson <sam@binarycake.ca>
Date: Thu, 5 Jun 2025 13:49:11 -0400
Subject: [PATCH] Cache subclasses of FQ

---
 py_ecc/fields/field_elements.py | 33 ++++++++++++++++++++++++++++++---
 1 file changed, 30 insertions(+), 3 deletions(-)

diff --git a/py_ecc/fields/field_elements.py b/py_ecc/fields/field_elements.py
index f55de8c..c8df623 100644
--- a/py_ecc/fields/field_elements.py
+++ b/py_ecc/fields/field_elements.py
@@ -12,6 +12,8 @@
     Union,
     cast,
 )
+from weakref import WeakValueDictionary
+from threading import Lock
 
 from py_ecc.utils import (
     deg,
@@ -201,6 +203,33 @@ def zero(cls: Type[T_FQ]) -> T_FQ:
 int_types_or_FQ = (int, FQ)
 
 
+_FQ_TYPES = WeakValueDictionary()
+_FQ_LOCK = Lock()
+
+
+def _make_fq(field_modulus: int) -> Type[FQ]:
+    try:
+        return _FQ_TYPES[field_modulus]
+    except KeyError:
+        pass
+
+    _FQ_LOCK.acquire()
+    try:
+        try:
+            return _FQ_TYPES[field_modulus]
+        except KeyError:
+            pass
+
+        new_type = type(
+            "FQP_corresponding_FQ_class", (FQ,), {"field_modulus": field_modulus}
+        )
+        _FQ_TYPES[field_modulus] = new_type
+        return new_type
+    finally:
+        _FQ_LOCK.release()
+
+
+
 class FQP:
     """
     A class for elements in polynomial extension fields
@@ -218,9 +247,7 @@ def __init__(
         if len(coeffs) != len(modulus_coeffs):
             raise Exception("coeffs and modulus_coeffs aren't of the same length")
         # Encoding all coefficients in the corresponding type FQ
-        self.FQP_corresponding_FQ_class = type(
-            "FQP_corresponding_FQ_class", (FQ,), {"field_modulus": self.field_modulus}
-        )
+        self.FQP_corresponding_FQ_class = _make_fq(self.field_modulus)
         self.coeffs: Tuple[IntOrFQ, ...] = tuple(
             self.FQP_corresponding_FQ_class(c) for c in coeffs
         )