9 files changed, 210 insertions, 38 deletions

diff --git a/.github/workflows/perf.yml b/.github/workflows/perf.yml
index 6211ea1..2b9b81c 100644
--- a/.github/workflows/perf.yml
+++ b/.github/workflows/perf.yml
@@ -6,6 +6,7 @@ env:
   LLVM_CONFIG: /usr/bin/llvm-config-10
   PS_PACKAGES: libps4000 libps5000 libps6000
   GMP_PACKAGES: libgmp-dev libmpfr-dev libmpc-dev
+  FLINT_PACKAGES: libflint-dev
   OTHER_PACKAGES: swig gcc libpcsclite-dev llvm-10 libllvm10 llvm-10-dev libpari-dev pari-gp pari-seadata
 
 jobs:
@@ -14,10 +15,10 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.9", "3.10", "3.11"]
-        gmp: [0, 1]
+        mod: ["python", "gmp", "flint"]
     env:
       PYTHON: ${{ matrix.python-version }}
-      USE_GMP: ${{ matrix.gmp }}
+      MOD_IMPL: ${{ matrix.mod }}
     steps:
       - uses: actions/checkout@v4
         with:
@@ -25,10 +26,10 @@ jobs:
       - uses: actions/cache@v4
         with:
           path: ~/.cache/pip
-          key: pip-${{ runner.os }}-${{ matrix.gmp }}-${{ matrix.python-version }}-${{ hashFiles('pyproject.toml') }}
+          key: pip-${{ runner.os }}-${{ matrix.mod }}-${{ matrix.python-version }}-${{ hashFiles('pyproject.toml') }}
           restore-keys: |
-            pip-${{ runner.os }}-${{ matrix.gmp }}-${{ matrix.python-version }}-
-            pip-${{ runner.os }}-${{ matrix.gmp }}-
+            pip-${{ runner.os }}-${{ matrix.mod }}-${{ matrix.python-version }}-
+            pip-${{ runner.os }}-${{ matrix.mod }}-
             pip-${{ runner.os }}-
       - name: Setup Python ${{ matrix.python-version }}
         uses: actions/setup-python@v5
@@ -42,22 +43,24 @@ jobs:
       - name: Install system dependencies
         run: |
           sudo apt-get install -y $PS_PACKAGES $OTHER_PACKAGES
-          if [ $USE_GMP == 1 ]; then sudo apt-get install -y $GMP_PACKAGES; fi
+          if [ $MOD_IMPL == "gmp" ]; then sudo apt-get install -y $GMP_PACKAGES; fi
+          if [ $MOD_IMPL == "flint" ]; then sudo apt-get install -y $FLINT_PACKAGES; fi
       - name: Install picoscope bindings
         run: |
           python -m pip install -U pip setuptools wheel
           git clone https://github.com/colinoflynn/pico-python && cd pico-python && pip install . && cd ..
           git clone https://github.com/picotech/picosdk-python-wrappers && cd picosdk-python-wrappers && pip install . && cd ..
-      - name: Install dependencies
+      - name: Install
         run: |
-          if [ $USE_GMP == 1 ]; then pip install -e ".[picoscope_sdk, picoscope_alt, chipwhisperer, smartcard, pari, leia, gmp, test, dev]"; fi
-          if [ $USE_GMP == 0 ]; then pip install -e ".[picoscope_sdk, picoscope_alt, chipwhisperer, smartcard, pari, leia, test, dev]"; fi
+          if [ $MOD_IMPL == "gmp" ];    then pip install -e ".[picoscope_sdk, picoscope_alt, chipwhisperer, smartcard, pari, leia, gmp, test, dev]"; fi
+          if [ $MOD_IMPL == "flint" ];  then pip install -e ".[picoscope_sdk, picoscope_alt, chipwhisperer, smartcard, pari, leia, flint, test, dev]"; fi
+          if [ $MOD_IMPL == "python" ]; then pip install -e ".[picoscope_sdk, picoscope_alt, chipwhisperer, smartcard, pari, leia, test, dev]"; fi
       - name: Perf
         run: |
           make perf
       - name: Archive perf results
         uses: actions/upload-artifact@v4
         with:
-          name: perf-results-${{ matrix.gmp }}-${{ matrix.python-version }}
+          name: perf-results-${{ matrix.mod }}-${{ matrix.python-version }}
           path:
             .perf
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 8a31b88..1134792 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -6,6 +6,7 @@ env:
   LLVM_CONFIG: /usr/bin/llvm-config-10
   PS_PACKAGES: libps4000 libps5000 libps6000
   GMP_PACKAGES: libgmp-dev libmpfr-dev libmpc-dev
+  FLINT_PACKAGES: libflint-dev
   OTHER_PACKAGES: swig gcc libpcsclite-dev llvm-10 libllvm10 llvm-10-dev libpari-dev pari-gp pari-seadata
 
 jobs:
@@ -14,10 +15,10 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.9", "3.10", "3.11"]
-        gmp: [0, 1]
+        mod: ["python", "gmp", "flint"]
     env:
       PYTHON: ${{ matrix.python-version }}
-      USE_GMP: ${{ matrix.gmp }}
+      MOD_IMPL: ${{ matrix.mod }}
     steps:
       - uses: actions/checkout@v4
         with:
@@ -25,10 +26,10 @@ jobs:
       - uses: actions/cache@v4
         with:
           path: ~/.cache/pip
-          key: pip-${{ runner.os }}-${{ matrix.gmp }}-${{ matrix.python-version }}-${{ hashFiles('pyproject.toml') }}
+          key: pip-${{ runner.os }}-${{ matrix.mod }}-${{ matrix.python-version }}-${{ hashFiles('pyproject.toml') }}
           restore-keys: |
-            pip-${{ runner.os }}-${{ matrix.gmp }}-${{ matrix.python-version }}-
-            pip-${{ runner.os }}-${{ matrix.gmp }}-
+            pip-${{ runner.os }}-${{ matrix.mod }}-${{ matrix.python-version }}-
+            pip-${{ runner.os }}-${{ matrix.mod }}-
             pip-${{ runner.os }}-
       - name: Setup Python ${{ matrix.python-version }}
         uses: actions/setup-python@v5
@@ -42,16 +43,18 @@ jobs:
       - name: Install system dependencies
         run: |
           sudo apt-get install -y $PS_PACKAGES $OTHER_PACKAGES
-          if [ $USE_GMP == 1 ]; then sudo apt-get install -y $GMP_PACKAGES; fi
+          if [ $MOD_IMPL == "gmp" ]; then sudo apt-get install -y $GMP_PACKAGES; fi
+          if [ $MOD_IMPL == "flint" ]; then sudo apt-get install -y $FLINT_PACKAGES; fi
       - name: Install picoscope bindings
         run: |
           python -m pip install -U pip setuptools wheel
           git clone https://github.com/colinoflynn/pico-python && cd pico-python && pip install . && cd ..
           git clone https://github.com/picotech/picosdk-python-wrappers && cd picosdk-python-wrappers && pip install . && cd ..
-      - name: Install dependencies
+      - name: Install
         run: |
-          if [ $USE_GMP == 1 ]; then pip install -e ".[picoscope_sdk, picoscope_alt, chipwhisperer, smartcard, pari, leia, gmp, test, dev]"; fi
-          if [ $USE_GMP == 0 ]; then pip install -e ".[picoscope_sdk, picoscope_alt, chipwhisperer, smartcard, pari, leia, test, dev]"; fi
+          if [ $MOD_IMPL == "gmp" ];    then pip install -e ".[picoscope_sdk, picoscope_alt, chipwhisperer, smartcard, pari, leia, gmp, test, dev]"; fi
+          if [ $MOD_IMPL == "flint" ];  then pip install -e ".[picoscope_sdk, picoscope_alt, chipwhisperer, smartcard, pari, leia, flint, test, dev]"; fi
+          if [ $MOD_IMPL == "python" ]; then pip install -e ".[picoscope_sdk, picoscope_alt, chipwhisperer, smartcard, pari, leia, test, dev]"; fi
       - name: Test
         run: |
           make test
diff --git a/pyecsca/ec/mod.py b/pyecsca/ec/mod.py
index 2aa4bdd..8453ed3 100644
--- a/pyecsca/ec/mod.py
+++ b/pyecsca/ec/mod.py
@@ -28,8 +28,21 @@ except ImportError:
     gmpy2 = None
 
 
+has_flint = False
+try:
+    import flint
+
+    _major, _minor, *_ = flint.__version__.split(".")
+    if (int(_major), int(_minor)) >= (0, 5):
+        has_flint = True
+    else:
+        flint = None
+except ImportError:
+    flint = None
+
+
 @public
-def gcd(a, b):
+def gcd(a: int, b: int) -> int:
     """Euclid's greatest common denominator algorithm."""
     if abs(a) < abs(b):
         return gcd(b, a)
@@ -42,7 +55,7 @@ def gcd(a, b):
 
 
 @public
-def extgcd(a, b):
+def extgcd(a: int, b: int) -> Tuple[int, int, int]:
     """Compute the extended Euclid's greatest common denominator algorithm."""
     if abs(b) > abs(a):
         x, y, d = extgcd(b, a)
@@ -138,6 +151,7 @@ class RandomModAction(ResultAction):
 
 
 _mod_classes: Dict[str, Type] = {}
+_mod_order = ["flint", "gmp", "python"]
 
 
 @public
@@ -156,7 +170,10 @@ class Mod:
         selected_class = getconfig().ec.mod_implementation
         if selected_class not in _mod_classes:
             # Fallback to something
-            selected_class = next(iter(_mod_classes.keys()))
+            for fallback in _mod_order:
+                if fallback in _mod_classes:
+                    selected_class = fallback
+                    break
         return _mod_classes[selected_class].__new__(
             _mod_classes[selected_class], *args, **kwargs
         )
@@ -232,9 +249,8 @@ class Mod:
         return ~self * other
 
     @_check
-    def __divmod__(self, divisor) -> Tuple["Mod", "Mod"]:
-        q, r = divmod(self.x, divisor.x)
-        return self.__class__(q, self.n), self.__class__(r, self.n)
+    def __divmod__(self, divisor):
+        raise NotImplementedError
 
     def __bytes__(self) -> bytes:
         raise NotImplementedError
@@ -667,11 +683,6 @@ if has_gmp:
         def __mul__(self, other) -> "GMPMod":
             return GMPMod((self.x * other.x) % self.n, self.n, ensure=False)
 
-        @_check
-        def __divmod__(self, divisor) -> Tuple["GMPMod", "GMPMod"]:
-            q, r = gmpy2.f_divmod(self.x, divisor.x)
-            return GMPMod(q, self.n, ensure=False), GMPMod(r, self.n, ensure=False)
-
         def __bytes__(self):
             return int(self.x).to_bytes((self.n.bit_length() + 7) // 8, byteorder="big")
 
@@ -708,3 +719,158 @@ if has_gmp:
             )
 
     _mod_classes["gmp"] = GMPMod
+
+
+if has_flint:
+
+    @lru_cache
+    def _fmpz_ctx(n):
+        if type(n) == flint.fmpz_mod_ctx:
+            return n
+        return flint.fmpz_mod_ctx(n)
+
+    @public
+    class FlintMod(Mod):
+        """An element x of ℤₙ. Implemented by GMP."""
+
+        x: flint.fmpz_mod
+        _ctx: flint.fmpz_mod_ctx
+        __slots__ = ("x", "_ctx")
+
+        def __new__(cls, *args, **kwargs):
+            return object.__new__(cls)
+
+        def __init__(
+            self,
+            x: Union[int, flint.fmpz_mod],
+            n: Union[int, flint.fmpz_mod_ctx],
+            ensure: bool = True,
+        ):
+            if ensure:
+                self._ctx = _fmpz_ctx(n)
+                self.x = self._ctx(x)
+            else:
+                self._ctx = n
+                self.x = x
+
+        @property
+        def n(self) -> flint.fmpz:
+            return self._ctx.modulus()
+
+        def bit_length(self):
+            return int(self.x).bit_length()
+
+        def inverse(self) -> "FlintMod":
+            if self.x == 0:
+                raise_non_invertible()
+            if self.x == 1:
+                return FlintMod(self._ctx(1), self._ctx, ensure=False)
+            try:
+                res = self.x.inverse()
+            except ZeroDivisionError:
+                raise_non_invertible()
+                res = self._ctx(0)
+            return FlintMod(res, self._ctx, ensure=False)
+
+        def is_residue(self) -> bool:
+            if not self.n.is_prime():
+                raise NotImplementedError
+            if self.x == 0:
+                return True
+            if self.n == 2:
+                return self.x in (0, 1)
+            legendre_symbol = jacobi(int(self.x), int(self.n))
+            return legendre_symbol == 1
+
+        def sqrt(self) -> "FlintMod":
+            if not self.n.is_prime():
+                raise NotImplementedError
+            if self.x == 0:
+                return FlintMod(self._ctx(0), self._ctx, ensure=False)
+            if not self.is_residue():
+                raise_non_residue()
+            mod = self.n
+            if mod % 4 == 3:
+                return self ** int((mod + 1) // 4)
+            q = mod - 1
+            s = 0
+            while q % 2 == 0:
+                q //= 2
+                s += 1
+
+            z = self._ctx(2)
+            while FlintMod(z, self._ctx, ensure=False).is_residue():
+                z += 1
+
+            m = s
+            c = FlintMod(z, self._ctx, ensure=False) ** int(q)
+            t = self ** int(q)
+            r_exp = (q + 1) // 2
+            r = self ** int(r_exp)
+
+            while t != 1:
+                i = 1
+                while not (t ** (2**i)) == 1:
+                    i += 1
+                two_exp = m - (i + 1)
+                b = c ** int(FlintMod(self._ctx(2), self._ctx, ensure=False) ** two_exp)
+                m = int(FlintMod(self._ctx(i), self._ctx, ensure=False))
+                c = b**2
+                t *= c
+                r *= b
+            return r
+
+        @_check
+        def __add__(self, other) -> "FlintMod":
+            return FlintMod(self.x + other.x, self._ctx, ensure=False)
+
+        @_check
+        def __sub__(self, other) -> "FlintMod":
+            return FlintMod(self.x - other.x, self._ctx, ensure=False)
+
+        def __neg__(self) -> "FlintMod":
+            return FlintMod(-self.x, self._ctx, ensure=False)
+
+        @_check
+        def __mul__(self, other) -> "FlintMod":
+            return FlintMod(self.x * other.x, self._ctx, ensure=False)
+
+        def __bytes__(self):
+            return int(self.x).to_bytes(
+                (int(self.n).bit_length() + 7) // 8, byteorder="big"
+            )
+
+        def __int__(self):
+            return int(self.x)
+
+        def __eq__(self, other):
+            if type(other) is int:
+                return self.x == other
+            if type(other) is not FlintMod:
+                return False
+            try:
+                return self.x == other.x
+            except ValueError:
+                return False
+
+        def __ne__(self, other):
+            return not self == other
+
+        def __repr__(self):
+            return str(int(self.x))
+
+        def __hash__(self):
+            return hash(("FlintMod", self.x, self.n))
+
+        def __pow__(self, n) -> "FlintMod":
+            if type(n) not in (int, flint.fmpz):
+                raise TypeError
+            if n == 0:
+                return FlintMod(self._ctx(1), self._ctx, ensure=False)
+            if n < 0:
+                return self.inverse() ** (-n)
+            if n == 1:
+                return FlintMod(self.x, self._ctx, ensure=False)
+            return FlintMod(self.x**n, self._ctx, ensure=False)
+
+    _mod_classes["flint"] = FlintMod
diff --git a/pyecsca/misc/cfg.py b/pyecsca/misc/cfg.py
index 05f89b8..eac5174 100644
--- a/pyecsca/misc/cfg.py
+++ b/pyecsca/misc/cfg.py
@@ -18,7 +18,7 @@ class ECConfig:
     _non_residue_action: str = "error"
     _unsatisfied_formula_assumption_action: str = "error"
     _unsatisfied_coordinate_assumption_action: str = "error"
-    _mod_implementation: str = "gmp"
+    _mod_implementation: str = "flint"
 
     @property
     def no_inverse_action(self) -> str:
@@ -109,6 +109,7 @@ class ECConfig:
 
         One of:
 
+         - ``"flint"``: Requires the flint library and `python-flint` package.
          - ``"gmp"``: Requires the GMP library and `gmpy2` package.
          - ``"python"``: Doesn't require anything.
          - ``"symbolic"``: Requires sympy.
@@ -117,8 +118,8 @@ class ECConfig:
 
     @mod_implementation.setter
     def mod_implementation(self, value: str):
-        if value not in ("python", "gmp", "symbolic"):
-            raise ValueError("Bad Mod implementaiton, can be one of 'python', 'gmp' or 'symbolic'.")
+        if value not in ("python", "gmp", "flint", "symbolic"):
+            raise ValueError("Bad Mod implementaiton, can be one of 'python', 'gmp', 'flint' or 'symbolic'.")
         self._mod_implementation = value
 
 
diff --git a/test/ec/perf_formula.py b/test/ec/perf_formula.py
index 0a5f2a0..7adf369 100755
--- a/test/ec/perf_formula.py
+++ b/test/ec/perf_formula.py
@@ -12,7 +12,7 @@ from test.utils import Profiler
 @click.option(
     "-m",
     "--mod",
-    type=click.Choice(("python", "gmp")),
+    type=click.Choice(("python", "gmp", "flint")),
     default="gmp" if has_gmp else "python",
 )
 @click.option("-o", "--operations", type=click.INT, default=5000)
diff --git a/test/ec/perf_mod.py b/test/ec/perf_mod.py
index 6a83314..f77c6bb 100755
--- a/test/ec/perf_mod.py
+++ b/test/ec/perf_mod.py
@@ -11,7 +11,7 @@ from test.utils import Profiler
 @click.option(
     "-m",
     "--mod",
-    type=click.Choice(("python", "gmp")),
+    type=click.Choice(("python", "gmp", "flint")),
     default="gmp" if has_gmp else "python",
 )
 @click.option("-o", "--operations", type=click.INT, default=100000)
diff --git a/test/ec/perf_mult.py b/test/ec/perf_mult.py
index d5e6a83..2038b57 100755
--- a/test/ec/perf_mult.py
+++ b/test/ec/perf_mult.py
@@ -16,7 +16,7 @@ from test.utils import Profiler
 @click.option(
     "-m",
     "--mod",
-    type=click.Choice(("python", "gmp")),
+    type=click.Choice(("python", "gmp", "flint")),
     default="gmp" if has_gmp else "python",
 )
 @click.option("-o", "--operations", type=click.INT, default=50)
diff --git a/test/ec/test_mod.py b/test/ec/test_mod.py
index 7234816..9a6cef0 100644
--- a/test/ec/test_mod.py
+++ b/test/ec/test_mod.py
@@ -157,7 +157,6 @@ def test_other():
     assert 5 // b == Mod(4, 7)
     assert a / 3 == Mod(4, 7)
     assert a // 3 == Mod(4, 7)
-    assert divmod(a, b) == (Mod(1, 7), Mod(2, 7))
     assert a + b == Mod(1, 7)
     assert 5 + b == Mod(1, 7)
     assert a + 3 == Mod(1, 7)
diff --git a/test/utils.py b/test/utils.py
index d4893eb..0939528 100644
--- a/test/utils.py
+++ b/test/utils.py
@@ -70,6 +70,6 @@ class Profiler:
         if self._state != "out":
             raise ValueError
         if self._prof_type == "py":
-            return self._root_frame.time
+            return self._root_frame.time  # type: ignore
         else:
             return pstats.Stats(self._prof).total_tt  # type: ignore