Add encoding and decoding to bytes from source

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

Author: stratospher

test/functional/test_framework/ellsq.py

@@ -8,7 +8,7 @@
 import random
 import unittest
 
-from .key import modsqrt, SECP256K1, SECP256K1_G, SECP256K1_ORDER
+from .key import modsqrt, SECP256K1, SECP256K1_FIELD_SIZE, SECP256K1_G, SECP256K1_ORDER
 
 class fe:
     """Prime field over 2^256 - 2^32 - 977"""
@@ -101,6 +101,38 @@ def r(x,y,i):
     else:
         return -u
 
+def encode(P):
+    while True:
+        u = field_random()
+        T = curve_negate(f(u))
+        Q = curve_add(T, P)
+        if is_infinity(Q): Q = T
+        j = secrets.choice([1,2,3,4])
+        v = r(Q, j)
+        if v is not Nothing: return (u, v)
+
+def decode(u, v):
+    T = f(u)
+    P = curve_add(T, f(v))
+    if is_infinity(P): P = T
+    return P
+
+P = SECP256K1_FIELD_SIZE
+FIELD_BITS = P.bit_length()
+FIELD_BYTES = (FIELD_BITS + 7) // 8
+PAD_BITS = FIELD_BYTES*8 - FIELD_BITS
+
+def encode_bytes(P):
+    u, v = encode(P)
+    up = u + secrets.randbits(PAD_BITS) << FIELD_BITS
+    vp = v + secrets.randbits(PAD_BITS) << FIELD_BITS
+    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(u, v)
+
 def SECP256K1_FE_CONST(d7, d6, d5, d4, d3, d2, d1, d0):
     n = []
     n.append(d0 | (((d1) & 0xFFFFF) << 32))