Add class for field elements and ellsq mapping functions

- source: https://github.com/sipa/writeups/tree/main/elligator-square-for-bn
- f maps maps every field element to a curve point
- r is a partial reverse function which can map a curve point to a field element

Author: stratospher

test/functional/test_framework/ellsq.py

@@ -0,0 +1,91 @@
+# Source: https://github.com/sipa/writeups/tree/main/elligator-square-for-bn
+"""Test-only Elligator Squared implementation
+
+WARNING: This code is slow and uses bad randomness.
+Do not use for anything but tests."""
+
+from .key import fe
+
+C1 = fe(-3).sqrt()
+C2 = (C1 - fe(1)) / fe(2)
+B = fe(7)
+
+def forward_map(u):
+    """Forward mapping function
+
+    Parameters:
+        u (of type fe) : any field element
+    Returns:
+        fe, fe : affine X and Y coordinates of a point on the secp256k1 curve
+    """
+    s = u**2
+    x1 = C2 - C1*s / (fe(1)+B+s)
+    g1 = x1**3 + B
+    if g1.is_square():
+        x, g = x1, g1
+    else:
+        x2 = -x1 - fe(1)
+        g2 = x2**3 + B
+        if g2.is_square():
+            x, g = x2, g2
+        else:
+            x3 = fe(1) - (fe(1)+B+s)**2 / (fe(3)*s)
+            g3 = x3**3 + B
+            x, g = x3, g3
+    y = g.sqrt()
+    if y.is_odd() == u.is_odd():
+        return x, y
+    else:
+        return x, -y
+
+def reverse_map(x, y, i):
+    """Reverse mapping function
+
+    Parameters:
+        fe, fe : X and Y coordinates of a point on the secp256k1 curve
+        i      : integer in range [0,3]
+    Returns:
+        u (of type fe) : such that forward_map(u) = (x,y), or None.
+
+        - There can be up to 4 such inverses, and i selects which formula to use.
+        - Each i can independently from other i values return a value or None.
+        - All non-None values returned across all 4 i values are guaranteed to be distinct.
+        - Together they will cover all inverses of (x,y) under forward_map.
+    """
+    if i == 0 or i == 1:
+        z = fe(2)*x + fe(1)
+        t1 = C1 - z
+        t2 = C1 + z
+        if not (t1*t2).is_square():
+            return None
+        if i == 0:
+            if t2 == fe(0):
+                return None
+            if t1 == fe(0) and y.is_odd():
+                return None
+            u = ((fe(1)+B)*t1/t2).sqrt()
+        else:
+            x1 = -x-fe(1)
+            if (x1**3 + B).is_square():
+                return None
+            u = ((fe(1)+B)*t2/t1).sqrt()
+    else:
+        z = fe(2) - fe(4)*B - fe(6)*x
+        if not (z**2 - fe(16)*(B+fe(1))**2).is_square():
+            return None
+        if i == 2:
+            s = (z + (z**2 - fe(16)*(B+fe(1))**2).sqrt()) / fe(4)
+        else:
+            if z**2 == fe(16)*(B+fe(1))**2:
+                return None
+            s = (z - (z**2 - fe(16)*(B+fe(1))**2).sqrt()) / fe(4)
+        if not s.is_square():
+            return None
+        x1 = C2 - C1*s / (fe(1)+B+s)
+        if (x1**3 + B).is_square():
+            return None
+        u = s.sqrt()
+    if y.is_odd() == u.is_odd():
+        return u
+    else:
+        return -u

test/functional/test_framework/key.py

@@ -60,6 +60,34 @@ def modsqrt(a, p):
         return sqrt
     return None
 
+SECP256K1_FIELD_SIZE = 2**256 - 2**32 - 977
+
+class fe:
+    """Prime field over 2^256 - 2^32 - 977"""
+    def __init__(self, x):
+        self.val = x % SECP256K1_FIELD_SIZE
+
+    def __add__     (self, o): return fe(self.val + o.val)
+    def __eq__      (self, o): return self.val == o.val
+    def __hash__    (self   ): return id(self)
+    def __mul__     (self, o): return fe(self.val * o.val)
+    def __neg__     (self   ): return fe(-self.val)
+    def __pow__     (self, s): return fe(pow(self.val, s, SECP256K1_FIELD_SIZE))
+    def __sub__     (self, o): return fe(self.val - o.val)
+    def __truediv__ (self, o): return fe(self.val * o.invert().val)
+
+    def invert      (self   ):
+        return fe(modinv(self.val, SECP256K1_FIELD_SIZE))
+    def is_odd(self): return (self.val & 1) != 0
+    def is_square(self):
+        return jacobi_symbol(self.val, SECP256K1_FIELD_SIZE) >= 0
+    def sqrt(self):
+        return fe(modsqrt(self.val, SECP256K1_FIELD_SIZE))
+
+    @staticmethod
+    def from_bytes(b): return fe(int.from_bytes(b, 'big'))
+    def to_bytes(self): return self.val.to_bytes(32, 'big')
+
 class EllipticCurve:
     def __init__(self, p, a, b):
         """Initialize elliptic curve y^2 = x^3 + a*x + b over GF(p)."""
@@ -216,7 +244,6 @@ def mul(self, ps):
                     r = self.add(r, p)
         return r
 
-SECP256K1_FIELD_SIZE = 2**256 - 2**32 - 977
 SECP256K1 = EllipticCurve(SECP256K1_FIELD_SIZE, 0, 7)
 SECP256K1_G = (0x79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798, 0x483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8, 1)
 SECP256K1_ORDER = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141