From 319741761d4fb9fe31f57b842a3b7a8fb7914a8c Mon Sep 17 00:00:00 2001
From: Gerald Teschl <gerald.teschl@univie.ac.at>
Date: Tue, 12 Nov 2024 18:39:42 +0100
Subject: [PATCH] Cleanup in ec

---
 doc/kryptools.ipynb |  2 +-
 kryptools/ec.py     |  5 +----
 tests/test_ec.py    | 52 +++++++++++++++++++++++++++++----------------
 3 files changed, 36 insertions(+), 23 deletions(-)

diff --git a/doc/kryptools.ipynb b/doc/kryptools.ipynb
index 12d9a40..8130a06 100644
--- a/doc/kryptools.ipynb
+++ b/doc/kryptools.ipynb
@@ -887,7 +887,7 @@
     {
      "data": {
       "text/plain": [
-       "270"
+       "266"
       ]
      },
      "execution_count": 29,
diff --git a/kryptools/ec.py b/kryptools/ec.py
index 11a951c..6800cd0 100644
--- a/kryptools/ec.py
+++ b/kryptools/ec.py
@@ -129,9 +129,7 @@ def dbl(self, x, y):
 
     def mult(self, j: int, x, y):  # Addition-subtraction ladder
         "Point multiplication"
-        if x is None:
-            return None, None
-        if j == 0:
+        if j == 0 or x is None:
             return None, None
         if j < 0:
             y = -y
@@ -166,7 +164,6 @@ def random(self):
 
     def psi(self, n: int):
         """The x-part of the n'th division polynomial."""
-
         if len(self.psi_list) < 5:
             self.psi_list = [ Poly([i], ring=self.gf) for i in range(3)]
             self.psi_list += [ Poly([-self.a * self.a, 12 * self.b, 6 * self.a, 0, 3], ring=self.gf) ]
diff --git a/tests/test_ec.py b/tests/test_ec.py
index dd1c0db..a9bfb5e 100644
--- a/tests/test_ec.py
+++ b/tests/test_ec.py
@@ -1,28 +1,44 @@
 import pytest
 from random import randint, seed
+from math import gcd
 from kryptools import EC_Weierstrass
 seed(0)
 
 
 def test_EC():
     ec = EC_Weierstrass(239, 3, 1)
+    assert len(list(ec)) == ec.order()
     O = ec.inf()  # point at infinity
     assert O in ec
-    P = ec.random()
-    assert P in ec
-    Q = ec.random()
-    assert Q in ec
-    assert P + O == P
-    assert O + P == P
-    assert P - P == O
-    assert P + P == 2 * P
-    R = O
-    for i in range(5):
-        assert i * P == R
-        R += P
-    assert P + Q == Q + P
-    assert Q.order() * Q == O
-    k = randint(1, Q.order())
-    R = k * Q
-    assert R.dlog(Q) == k
-    assert len(list(ec)) == ec.order()
+    assert O + O == O
+    assert O - O == O
+    assert 0 * O == O
+    assert 3 * O == O
+    assert O.order() == 1
+    assert O.dlog(O) == 0
+    for _ in range(100):
+        P = ec.random()
+        assert P in ec
+        assert P + O == P
+        assert O + P == P
+        assert P - O == P
+        assert O - P == -P
+        assert P - P == O
+        assert P + P == 2 * P
+        Q = O
+        order_P = P.order()
+        assert order_P * P == O
+        for i in range(9):
+            assert i * P == Q
+            assert Q.order() == order_P // gcd(order_P, i)
+            Q += P
+        Q = ec.random()
+        assert Q in ec
+        assert P + Q == Q + P
+        R = ec.random()
+        assert R in ec
+        assert (P + Q) + R == P + (Q + R)
+        for _ in range(10):
+            k = randint(0, Q.order()-1)
+            R = k * Q
+            assert R.dlog(Q) == k