/ sage / square_root_bls12_377.sage
square_root_bls12_377.sage
  1  # Constantine
  2  # Copyright (c) 2018-2019    Status Research & Development GmbH
  3  # Copyright (c) 2020-Present Mamy André-Ratsimbazafy
  4  # Licensed and distributed under either of
  5  #   * MIT license (license terms in the root directory or at http://opensource.org/licenses/MIT).
  6  #   * Apache v2 license (license terms in the root directory or at http://www.apache.org/licenses/LICENSE-2.0).
  7  # at your option. This file may not be copied, modified, or distributed except according to those terms.
  8  
  9  # ############################################################
 10  #
 11  #                         BLS12-377
 12  #                 Constant-time Square Root
 13  #
 14  # ############################################################
 15  
 16  # Parameters
 17  x = 3 * 2^46 * (7 * 13 * 499) + 1
 18  p = (x - 1)^2 * (x^4 - x^2 + 1)//3 + x
 19  r = x^4 - x^2 + 1
 20  t = x + 1
 21  print('x  : ' + x.hex())
 22  print('p  : ' + p.hex())
 23  print('r  : ' + r.hex())
 24  print('t  : ' + t.hex())
 25  
 26  def modCheck(p, pow):
 27      ## Find q mod 2^s != 1
 28      q = p^pow
 29      s = 4
 30      while q % s == 1:
 31          s *= 2
 32          if s > q:
 33              raise ValueError('Uh Oh')
 34      if pow == 1:
 35        print(f'Found: p mod {s} = {q % s}')
 36      else:
 37        print(f'Found: p^{pow} mod {s} = {q % s}')
 38  
 39  modCheck(p, 1) # Found: p mod 140737488355328 = 70368744177665
 40  modCheck(p, 2) # Found: p^2 mod 281474976710656 = 140737488355329
 41  
 42  # On Fp
 43  # a^((p-70368744177665+140737488355328)/140737488355328)
 44  # would lead to a square root but there would be
 45  # log2(140737488355328)-1 candidates
 46  # which must be checked constant time
 47  
 48  def precomp_tonelli_shanks(p):
 49      ## Precompute constants for
 50      ## constant-time Tonelli Shanks algorithm
 51      ## with q = p^pow returns:
 52      ## 1. c1, the largest integer such that 2^c1 divides q - 1.
 53      ## 2. c2 = (q - 1) / (2^c1) in ℕ
 54      ## 3. c3 = (c2 - 1) / 2 in ℕ
 55      ## 4. c4, a non-square value in Fq
 56      ## 5. c5 = c4^c2 in Fq
 57      q = p
 58      c1 = 0
 59      c2 = q-1
 60      while c2 & 1 == 0:
 61          c2 >>= 1
 62          c1 += 1
 63      c3 = (c2 - 1) // 2
 64      c4 = 1
 65      while kronecker(c4, q) == 1:
 66          c4 += 1
 67      c5 = GF(p)(c4)^c2
 68      return (c1,c2,c3,c4,c5)
 69  
 70  def ccopy(a, b, ctl):
 71      ## `b` if `ctl` is true, `a` if false
 72      return int(not(bool(ctl)))*a + int(bool(ctl))*b
 73  
 74  def sqrt_tonelli_shanks(x, p):
 75      ## Returns z = x² (p^pow)
 76      (c1, c2, c3, c4, c5) = precomp_tonelli_shanks(p)
 77  
 78      x = GF(p)(x)
 79  
 80      z = x^c3
 81      t = z*z*x
 82      z *= x
 83      b = t
 84      c = c5
 85      for i in range(c1, 1, -1): # c1 ... 2
 86          for j in range(1, i-1): # 1 ... i-2
 87              b *= b
 88          z = ccopy(z, z*c, b != 1)
 89          c *= c
 90          t = ccopy(t, t*c, b != 1)
 91          b = t
 92      return z
 93  
 94  # for a in range(2, 30):
 95  #     if kronecker(a, p) != 1:
 96  #         continue
 97  #     # print(f'{a}^(p-1)/2 = ' + str(GF(p)(a)^((p-1)/2)))
 98  #     print(f'{a} is a quadratic residue mod p')
 99  #     b = sqrt_tonelli_shanks(a, p)
100  #     # print(f'{b}² = {a} mod p')
101  #     # print('b*b = ' + str(b*b))
102  #     assert b*b == a
103  
104  # Optimized Tonelli Shanks
105  # --------------------------------------------------------
106  
107  # Finite fields
108  Fp       = GF(p)
109  K2.<u>  = PolynomialRing(Fp)
110  Fp2.<beta>  = Fp.extension(u^2+5)
111  
112  def precomp_ts(Fq):
113      ## From q = p^m with p the prime characteristic of the field Fp^m
114      ##
115      ## Returns (s, e) such as
116      ## q == s * 2^e + 1
117      s = Fq.order() - 1
118      e = 0
119      while s & 1 == 0:
120          s >>= 1
121          e += 1
122      return s, e
123  
124  def find_any_qnr(Fq):
125      ## Find a quadratic Non-Residue
126      ## in GF(p^m)
127      qnr = Fq(Fq.gen())
128      r = Fq.order()
129      while qnr.is_square():
130          qnr += 1
131      return qnr
132  
133  def sqrt_exponent_precomp(Fq, e):
134      ## Returns precomputation a^((q-1-2^e)/(2*2^e))
135      ##
136      ## With 2^e the largest power of 2 that divides q-1
137      ##
138      ## For all sqrt related functions
139      ## - legendre symbol
140      ## - SQRT
141      ## - inverse SQRT
142      r = Fq.order()
143      precomp = (r - 1) >> e       # (q-1) / 2^e
144      precomp = (precomp - 1) >> 1 # ((q-1) / 2^e) - 1) / 2 = (q-1-2^e)/2^e / 2
145      return precomp
146  
147  s, e = precomp_ts(Fp)
148  qnr = find_any_qnr(Fp)
149  root_unity = qnr^s
150  exponent = sqrt_exponent_precomp(Fp, e)
151  
152  # print('tonelli        s: 0x' + Integer(s).hex())
153  print('tonelli e (2-adicity): ' + str(e))
154  print('tonelli          root: 0x' + Integer(root_unity).hex())
155  print('tonelli      exponent: 0x' + Integer(exponent).hex())
156  
157  def legendre_symbol_impl(a, e, a_pre_exp):
158      ## Legendre symbol χ(a) = a^(q-1)/2
159      ## -1 if a is non-square
160      ## 0 if a is 0
161      ## 1 if a is square
162      ##
163      ## a_pre_exp = a^((q-1-2^e)/(2*2^e))
164      ## with
165      ##  s and e, precomputed values
166      ##  such as q == s * 2^e + 1
167      ##
168      ## a_pre_exp is used in square root
169      ## and or inverse square root computation
170      ##
171      ## for fused operations
172      r = a_pre_exp * a_pre_exp # a^((q-1-2^e)/2^e) = a^((q-1)/2^e - 1)
173      r *= a                    # a^((q-1)/2^e)
174      for i in range(0, e-1):
175          r *= r                # a^((q-1)/2)
176  
177      return r
178  
179  def legendre_symbol(a):
180      a_pre_exp = a^exponent
181      return legendre_symbol_impl(a, e, a_pre_exp)
182  
183  for a in range(20):
184      assert kronecker(a, p) == legendre_symbol(GF(p)(a))
185  
186  def sqrt_tonelli_shanks_impl(a, a_pre_exp, s, e, root_of_unity):
187      ## Square root for any `a` in a field of prime characteristic p
188      ##
189      ## a_pre_exp = a^((q-1-2^e)/(2*2^e))
190      ## with
191      ##  s and e, precomputed values
192      ##  such as q == s * 2^e + 1
193      z = a_pre_exp
194      t = z*z*a
195      r = z * a
196      b = t
197      root = root_of_unity
198      for i in range(e, 1, -1):   # e .. 2
199          for j in range(1, i-1): # 1 .. i-2
200              b *= b
201          doCopy = b != 1
202          r = ccopy(r, r * root, doCopy)
203          root *= root
204          t = ccopy(t, t * root, doCopy)
205          b = t
206      return r
207  
208  def sqrt_tonelli_shanks_opt(a):
209      a_pre_exp = a^exponent
210      return sqrt_tonelli_shanks_impl(a, a_pre_exp, s, e, root_unity)
211  
212  # for a in range(2, 30):
213  #     if kronecker(a, p) != 1:
214  #         continue
215  #     # print(f'{a}^(p-1)/2 = ' + str(GF(p)(a)^((p-1)/2)))
216  #     print(f'{a} is a quadratic residue mod p')
217  #     b = sqrt_tonelli_shanks_opt(GF(p)(a))
218  #     # print(f'{b}² = {a} mod p')
219  #     # print('b*b = ' + str(b*b))
220  #     assert b*b == a
221  
222  def sqrt_inv_sqrt_tonelli_shanks_impl(a, a_pre_exp, s, e, root_of_unity):
223      ## Square root and inverse square root for any `a` in a field of prime characteristic p
224      ##
225      ## a_pre_exp = a^((q-1-2^e)/(2*2^e))
226      ## with
227      ##  s and e, precomputed values
228      ##  such as q == s * 2^e + 1
229  
230      # Implementation
231      # 1/√a * a = √a
232      # Notice that in Tonelli Shanks, the result `r` is bootstrapped by "z*a"
233      # We bootstrap it instead by just z to get invsqrt for free
234  
235      z = a_pre_exp
236      t = z*z*a
237      r = z
238      b = t
239      root = root_of_unity
240      for i in range(e, 1, -1):   # e .. 2
241          for j in range(1, i-1): # 1 .. i-2
242              b *= b
243          doCopy = b != 1
244          r = ccopy(r, r * root, doCopy)
245          root *= root
246          t = ccopy(t, t * root, doCopy)
247          b = t
248      return r*a, r
249  
250  def sqrt_invsqrt_tonelli_shanks_opt(a):
251      a_pre_exp = a^exponent
252      return sqrt_inv_sqrt_tonelli_shanks_impl(a, a_pre_exp, s, e, root_unity)
253  
254  for a in range(2, 30):
255      if kronecker(a, p) != 1:
256          continue
257      # print(f'{a}^(p-1)/2 = ' + str(GF(p)(a)^((p-1)/2)))
258      print(f'{a} is a quadratic residue mod p')
259      b, invb = sqrt_invsqrt_tonelli_shanks_opt(GF(p)(a))
260      # print(f'{b}² = {a} mod p')
261      # print('b*b = ' + str(b*b))
262      assert b*b == a
263      assert invb*a == b