Add functions for encoding and decoding to bytes

Source: https://github.com/sipa/writeups/tree/main/elligator-square-for-bn

Author: stratospher

test/functional/test_framework/ellsq.py

@@ -6,9 +6,10 @@
 anything but tests."""
 
 import random
+import secrets
 import unittest
 
-from .key import fe, SECP256K1, SECP256K1_G, SECP256K1_ORDER
+from .key import fe, SECP256K1, SECP256K1_FIELD_SIZE, SECP256K1_G, SECP256K1_ORDER
 
 C1 = fe(-3).sqrt()
 C2 = (C1 - fe(1)) / fe(2)
@@ -124,6 +125,51 @@ def r(x,y,i):
     [(fe(0x3498662504b73c7c8cecb6c33cd493bdfc190e0f87d913d7ff9ad42e222bfe95), fe(0x245b3a61b8d46997f14f2fea2874899691eb32542b9907d65eb9d21d42454021)), [fe(0x7f556282c3dd9d263390d6bbddada698ab8fd7c7d1a06498f42b30437c8361ad), None                                                                  , None                                                                  , None                                                                  ]]
 ]
 
+def encode(P):
+    while True:
+        u = fe(random.randrange(1, SECP256K1_ORDER))
+        fe1 = f(u)
+        # convert fe1 to jacobian form for EC operations
+        fe1 = (fe1[0].val, fe1[1].val, 1)
+        T = SECP256K1.negate(fe1)
+        Q = SECP256K1.add(T, P)
+        if SECP256K1.is_infinity(Q):
+            Q = T
+        j = secrets.choice([0,1,2,3])
+        Q = SECP256K1.affine(Q)
+        v = r(fe(Q[0]), fe(Q[1]), j)
+        if v is not None:
+            return (u, v)
+
+def decode(u, v):
+    fe1 = f(u)
+    fe2 = f(v)
+    # convert fe1 and fe2 to jacobian form for EC operations
+    jac1 = (fe1[0].val, fe1[1].val, 1)
+    jac2 = (fe2[0].val, fe2[1].val, 1)
+    T = jac1
+    S = jac2
+    P = SECP256K1.add(T, S)
+    if SECP256K1.is_infinity(P):
+        P = T
+    P = SECP256K1.affine(P)
+    return (fe(P[0]), fe(P[1]))
+
+P = SECP256K1_FIELD_SIZE
+FIELD_BITS = P.bit_length()
+FIELD_BYTES = (FIELD_BITS + 7) // 8
+
+def encode_bytes(P):
+    u, v = encode(P)
+    up = u.val
+    vp = v.val
+    return up.to_bytes(FIELD_BYTES, 'big') + vp.to_bytes(FIELD_BYTES, 'big')
+
+def decode_bytes(enc):
+    u = (int.from_bytes(enc[:FIELD_BYTES], 'big') & ((1 << FIELD_BITS) - 1)) % P
+    v = (int.from_bytes(enc[FIELD_BYTES:], 'big') & ((1 << FIELD_BITS) - 1)) % P
+    return decode(fe(u), fe(v))
+
 class TestFrameworkEllsq(unittest.TestCase):
     def test_fe_to_ge_to_fe(self):
         for i in range(100):