1 /*
2     Copyright (C) 2012 Sebastian Pancratz
3 
4     This file is part of FLINT.
5 
6     FLINT is free software: you can redistribute it and/or modify it under
7     the terms of the GNU Lesser General Public License (LGPL) as published
8     by the Free Software Foundation; either version 2.1 of the License, or
9     (at your option) any later version.  See <http://www.gnu.org/licenses/>.
10 */
11 
12 #include <assert.h>
13 
14 #include "nmod_mat.h"
15 
16 #include "fmpz_mod_poly.h"
17 #include "qadic.h"
18 
19 /*
20     FILE DOCUMENTATION.
21 
22     This file contains routines for computing square roots in
23     finite fields and Hensel lifting.
24 
25         - _find_nonresidue()
26         - _artin_schreier_preimage()
27         - _fmpz_mod_poly_sqrtmod_p()
28         - _fmpz_mod_poly_sqrtmod_2()
29         - _qadic_sqrt_p()
30         - _qadic_sqrt()
31         - qadic_sqrt()
32  */
33 
34 /*
35     Sets \code{(rop,d)} to a non-residue in $\mathbf{F_q}$ for odd $p$
36     and $d \geq 2$.
37  */
_find_nonresidue(fmpz * rop,const fmpz * a,const slong * j,slong lena,const fmpz_t p)38 static void _find_nonresidue(fmpz *rop,
39                              const fmpz *a, const slong *j, slong lena,
40                              const fmpz_t p)
41 {
42     const slong d = j[lena - 1];
43 
44     fmpz *w;
45     fmpz_t pm1, z;
46     slong i;
47 
48     w = _fmpz_vec_init(2 * d - 1);
49     fmpz_init(pm1);
50     fmpz_init(z);
51 
52     fmpz_sub_ui(pm1, p, 1);
53     fmpz_pow_ui(z, p, d);
54     fmpz_sub_ui(z, z, 1);
55     fmpz_fdiv_q_2exp(z, z, 1);
56 
57     /*
58         Find a non-residue g;
59         uses stack-based depth first traversal starting from [0,...,0,1]
60      */
61     _fmpz_vec_zero(rop, d);
62     fmpz_one(rop + (d - 1));
63 
64     for (i = d; i > 0; )
65     {
66         if (i == d)
67         {
68             /* Consider this element, rop^{(q-1)/2} == -1 ? */
69             _qadic_pow(w, rop, d, z, a, j, lena, p);
70             if (fmpz_equal(w + 0, pm1) && _fmpz_vec_is_zero(w + 1, d - 1))
71                 break;
72 
73             /* Backtrace, find the next element */
74             for (i--; i >= 0 && fmpz_equal(rop + i, pm1); i--) ;
75             if (i >= 0)
76             {
77                 fmpz_add_ui(rop + i, rop + i, 1);
78                 i++;
79             }
80         }
81         else
82         {
83             _fmpz_vec_zero(rop + i, d - i);
84             i = d;
85         }
86     }
87 
88     _fmpz_vec_clear(w, 2 * d - 1);
89     fmpz_clear(pm1);
90     fmpz_clear(z);
91 }
92 
93 /*
94     Given an element $d$ as \code{(op,len)}, returns whether it has a
95     preimage $u$ under the Artin Schreier map $Y \mapsto Y^2 + Y$, and
96     if so sets \code{(rop,d)} to $u$.
97 
98     In this case, the other preimage is given by $u + 1$.  This
99     completes the set of preimages as the kernel of the Artin Schreier
100     map is $\mathbf{F}_2$.
101 
102     The value of \code{(rop,d)}$ is undefined when the return value
103     is zero.
104  */
105 FLINT_DLL int
_artin_schreier_preimage(fmpz * rop,const fmpz * op,slong len,const fmpz * a,const slong * j,slong lena)106 _artin_schreier_preimage(fmpz *rop, const fmpz *op, slong len,
107                          const fmpz *a, const slong *j, slong lena)
108 {
109     const slong d   = j[lena - 1];
110     const fmpz_t p = {WORD(2)};
111 
112     int ans;
113 
114     fmpz *e, *f;
115     nmod_mat_t A;
116     slong i, k, rk, *P;
117 
118     e = _fmpz_vec_init(d);
119     f = _fmpz_vec_init(2 * d - 1);
120 
121     nmod_mat_init(A, d, d, 2);
122     P  = flint_malloc(d * sizeof(slong));
123 
124     /* Construct Artin Schreier matrix ------------------------------------- */
125 
126     for (i = 0; i < d; i++)
127     {
128         fmpz_one(e + i);
129 
130         _fmpz_poly_sqr(f, e, i+1);
131         _fmpz_poly_reduce(f, 2*(i+1)-1, a, j, lena);
132         fmpz_add_ui(f + i, f + i, 1);
133         _fmpz_vec_scalar_mod_fmpz(f, f, d, p);
134 
135         for (k = 0; k < d; k++)
136         {
137             nmod_mat_entry(A, k, i) = (mp_limb_t) f[k];
138         }
139         fmpz_zero(e + i);
140     }
141 
142     /* Perform LUP decomposition ------------------------------------------- */
143 
144     /*
145         Write LU = PA and r = rank(A) so that U is in row echelon form
146         with only the top r rows non-zero, and L is lower unit triangular
147         truncated to r columns.
148 
149         We know that Y^2 + Y = 0 if and only if Y is in the base field,
150         i.e., dim ker(A) = 1 and rank(A) = d-1.
151 
152         Consider the linear system A x = b, which we can then write as
153         LU x = Pb, hence L y = Pb and U x = y.
154      */
155 
156     rk = nmod_mat_lu(P, A, 0);
157 
158     assert(rk == d-1);
159 
160     /* Solve for a preimage of (op,len) ------------------------------------ */
161 
162     _fmpz_vec_zero(rop, d);
163 
164     for (i = 0; i < d; i++)
165     {
166         rop[i] = P[i] < len ? op[P[i]] : 0;
167         for (k = 0; k < i; k++)
168             rop[i] ^= nmod_mat_entry(A, i, k) & rop[k];
169     }
170 
171     ans = (rop[d-1] == 0);
172 
173     if (ans)
174     {
175         slong c;
176 
177         for (c = 0; c < d; c++)
178             if (nmod_mat_entry(A, c, c) == 0)
179                 break;
180 
181         for (k = d - 1; k > c; k--)
182         {
183             rop[k] = rop[k-1];
184 
185             if ((rop[k]))
186                 for (i = k - 2; i >= 0; i--)
187                     rop[i] ^= nmod_mat_entry(A, i, k);
188         }
189         rop[k] = 0;
190         for (k--; k >= 0; k--)
191         {
192             if ((rop[k]))
193                 for (i = k - 1; i >= 0; i--)
194                     rop[i] ^= nmod_mat_entry(A, i, k);
195         }
196     }
197 
198     /* Clean-up ------------------------------------------------------------ */
199 
200     _fmpz_vec_clear(e, d);
201     _fmpz_vec_clear(f, 2 * d - 1);
202     nmod_mat_clear(A);
203     flint_free(P);
204 
205     return ans;
206 }
207 
208 /*
209     Returns whether the non-zero element \code{(op, len)} is a square,
210     and if so sets \code{(rop, 2 * d - 1)} to its square root.
211 
212     Assumes that $p$ is an odd prime.
213 
214     Assumes that $d \geq 1$.
215 
216     Does not support aliasing.
217  */
218 static int
_fmpz_mod_poly_sqrtmod_p(fmpz * rop,const fmpz * op,slong len,const fmpz * a,const slong * j,slong lena,const fmpz_t p)219 _fmpz_mod_poly_sqrtmod_p(fmpz *rop, const fmpz *op, slong len,
220                          const fmpz *a, const slong *j, slong lena,
221                          const fmpz_t p)
222 {
223     const slong d = j[lena - 1];
224     int ans;
225 
226     /*
227         When $q \equiv 3 \pmod{4}$...
228 
229         A non-zero element $x$ is a square if and only if $x^{(q-1)/2} = 1$,
230         and in this case one of its square roots is given by $x^{(q+1)/4}$.
231 
232         To avoid recomputation of powers of $x$, we compute $x^{(q-3)/4}$,
233         multiply this by $x$ to obtain the potential square root $x^{(q+1)/4}$,
234         and then combine these two powers to find $x^{(q-1)/2}$.
235      */
236     if (fmpz_fdiv_ui(p, 4) == 3 && (d & WORD(1)))
237     {
238         fmpz_t z;
239         fmpz *v, *w;
240 
241         fmpz_init(z);
242         v = _fmpz_vec_init(4 * d - 2);
243         w = v + (2 * d - 1);
244 
245         fmpz_pow_ui(z, p, d);
246         fmpz_sub_ui(z, z, 3);
247         fmpz_fdiv_q_2exp(z, z, 2);
248 
249         _qadic_pow(v, op, len, z, a, j, lena, p);
250 
251         _fmpz_poly_mul(rop, v, d, op, len);
252         _fmpz_vec_zero(rop + d + len - 1, d - len);
253         _fmpz_mod_poly_reduce(rop, d + len - 1, a, j, lena, p);
254 
255         _fmpz_poly_mul(w, rop, d, v, d);
256         _fmpz_mod_poly_reduce(w, 2 * d - 1, a, j, lena, p);
257         ans = fmpz_is_one(w + 0);
258 
259         fmpz_clear(z);
260         _fmpz_vec_clear(v, 4 * d - 2);
261     }
262 
263     /*
264         When $q \equiv 5 \pmod{8}$...
265 
266         A non-zero element $x$ is a square if and only if $y = x^{(q-1)/4}$
267         is $\pm 1$.  If $y = +1$, a square root is given by $x^{(q+3)/8}$,
268         otherwise it is $(2x) (4x)^{(q-5)/8}$.
269 
270         We begin by computing $v = x^{(q-5)/8}$, $w = x v = x^{(q+3)/8}$
271         and $y = v w = x^{(q-1)/4}$.  If $y = +1$ we return $w$, and if
272         $y = -1$ we return $2^{(q-1)/4} x^{(q+3)/8} = 2^{(q-1)/4} w$.
273      */
274     else if (fmpz_fdiv_ui(p, 8) == 5 && (d & WORD(1)))
275     {
276         fmpz_t f, g, pm1;
277 
278         fmpz *v;  /* v = x^{(q-5)/8} */
279         fmpz *w;  /* w = x^{(q+3)/8} */
280         fmpz *y;  /* y = x^{(q-1)/4} */
281 
282         fmpz_init(f);
283         fmpz_init(g);
284         fmpz_init(pm1);
285         fmpz_sub_ui(pm1, p, 1);
286 
287         v = _fmpz_vec_init(2 * d - 1);
288         w = _fmpz_vec_init(2 * d - 1);
289         y = _fmpz_vec_init(2 * d - 1);
290 
291         fmpz_pow_ui(g, p, d);
292         fmpz_sub_ui(g, g, 5);
293         fmpz_fdiv_q_2exp(g, g, 3);
294         _qadic_pow(v, op, len, g, a, j, lena, p);
295 
296         _fmpz_poly_mul(w, v, d, op, len);
297         _fmpz_mod_poly_reduce(w, d + len - 1, a, j, lena, p);
298 
299         _fmpz_poly_mul(y, v, d, w, d);
300         _fmpz_mod_poly_reduce(y, 2 * d - 1, a, j, lena, p);
301 
302         if ((fmpz_is_one(y + 0) || fmpz_equal(y + 0, pm1))  /* q.r. */
303             && _fmpz_vec_is_zero(y + 1, d - 1))
304         {
305             if (fmpz_is_one(y + 0))
306             {
307                 _fmpz_vec_set(rop, w, d);
308             }
309             else
310             {
311                 fmpz_t two = {WORD(2)};
312 
313                 fmpz_zero(g);
314                 fmpz_pow_ui(g, p, d);
315                 fmpz_sub_ui(g, g, 1);
316                 fmpz_fdiv_q_2exp(g, g, 2);
317                 _fmpz_vec_set(rop, v, d);
318 
319                 fmpz_powm(f, two, g, p);
320 
321                 _fmpz_mod_poly_scalar_mul_fmpz(rop, w, d, f, p);
322             }
323             _fmpz_vec_zero(rop + d, d - 1);
324             ans = 1;
325         }
326         else  /* q.n.r. */
327         {
328             ans = 0;
329         }
330 
331         fmpz_clear(f);
332         fmpz_clear(g);
333         fmpz_clear(pm1);
334 
335         _fmpz_vec_clear(v, 2 * d - 1);
336         _fmpz_vec_clear(w, 2 * d - 1);
337         _fmpz_vec_clear(y, 2 * d - 1);
338     }
339 
340     /*
341         General case for odd $q$...
342 
343         TODO:  Find a better way to integrate the check for square-ness
344         into the computation of a potential square root.
345      */
346     else
347     {
348         slong i, s;
349 
350         fmpz_t t, pm1, qm1, z;
351 
352         fmpz *b, *g, *bpow, *gpow, *w;
353 
354         fmpz_init(t);
355         fmpz_init(pm1);
356         fmpz_init(qm1);
357         fmpz_init(z);
358 
359         fmpz_sub_ui(pm1, p, 1);
360         fmpz_pow_ui(qm1, p, d);
361         fmpz_sub_ui(qm1, qm1, 1);
362 
363         b    = _fmpz_vec_init(2 * d - 1);
364         g    = _fmpz_vec_init(2 * d - 1);
365         bpow = _fmpz_vec_init(2 * d - 1);
366         gpow = _fmpz_vec_init(2 * d - 1);
367         w    = _fmpz_vec_init(2 * d - 1);
368 
369         /* Check whether op is a square, i.e. op^{(q-1}/2} == 1 */
370         fmpz_fdiv_q_2exp(z, qm1, 1);
371         _qadic_pow(w, op, len, z, a, j, lena, p);
372         ans = fmpz_is_one(w);
373 
374         if (ans)
375         {
376             /* Find a non-residue g */
377             _find_nonresidue(g, a, j, lena, p);
378 
379             /* Write q - 1 = 2^s t */
380             for (s = 0, fmpz_set(t, qm1); fmpz_is_even(t); s++)
381                 fmpz_fdiv_q_2exp(t, t, 1);
382 
383             /* Set g = g^t */
384             _qadic_pow(w, g, d, t, a, j, lena, p);
385             _fmpz_vec_set(g, w, d);
386 
387             /* Set rop = op^{(t+1)/2} */
388             fmpz_add_ui(z, t, 1);
389             fmpz_fdiv_q_2exp(z, z, 1);
390             _qadic_pow(rop, op, len, z, a, j, lena, p);
391 
392             /* Set b = op^t */
393             _qadic_pow(b, op, len, t, a, j, lena, p);
394 
395             while (!_fmpz_poly_is_one(b, d))
396             {
397                 slong k;
398 
399                 _fmpz_vec_set(bpow, b, d);
400                 for (k = 1; (k < s) && !_fmpz_poly_is_one(bpow, d); k++)
401                 {
402                     _fmpz_poly_sqr(w, bpow, d);
403                     _fmpz_mod_poly_reduce(w, 2 * d - 1, a, j, lena, p);
404                     _fmpz_vec_scalar_mod_fmpz(bpow, w, d, p);
405                 }
406 
407                 _fmpz_vec_set(gpow, g, d);
408                 for (i = 1; i < s - k; i++)
409                 {
410                     _fmpz_poly_sqr(w, gpow, d);
411                     _fmpz_mod_poly_reduce(w, 2 * d - 1, a, j, lena, p);
412                     _fmpz_vec_scalar_mod_fmpz(gpow, w, d, p);
413                 }
414 
415                 _fmpz_poly_mul(w, rop, d, gpow, d);
416                 _fmpz_mod_poly_reduce(w, 2 * d - 1, a, j, lena, p);
417                 _fmpz_vec_scalar_mod_fmpz(rop, w, d, p);
418 
419                 _fmpz_poly_sqr(w, gpow, d);
420                 _fmpz_mod_poly_reduce(w, 2 * d - 1, a, j, lena, p);
421                 _fmpz_vec_scalar_mod_fmpz(gpow, w, d, p);
422 
423                 _fmpz_poly_mul(w, b, d, gpow, d);
424                 _fmpz_mod_poly_reduce(w, 2 * d - 1, a, j, lena, p);
425                 _fmpz_vec_scalar_mod_fmpz(b, w, d, p);
426 
427                 s = k;
428             }
429         }
430 
431         fmpz_clear(t);
432         fmpz_clear(pm1);
433         fmpz_clear(qm1);
434         fmpz_clear(z);
435         _fmpz_vec_clear(b,    2 * d - 1);
436         _fmpz_vec_clear(g,    2 * d - 1);
437         _fmpz_vec_clear(bpow, 2 * d - 1);
438         _fmpz_vec_clear(gpow, 2 * d - 1);
439         _fmpz_vec_clear(w,    2 * d - 1);
440     }
441     return ans;
442 }
443 
444 /*
445     Sets \code{(rop, 2 * d - 1)} to the square root of
446     \code{(op, len)}.
447 
448     Note that the group of units of $\mathbf{F}_q$ is cyclic of order $q - 1$,
449     which is odd if $p = 2$.  We have $x^{q-1} = 1$ for every non-zero $x$, so
450     $x^q = x$ for every $x$ and $u = x^{q/2}$ satisfies $u^2 = x$.
451 
452     Assumes that $d \geq 1$.
453  */
454 static void
_fmpz_mod_poly_sqrtmod_2(fmpz * rop,const fmpz * op,slong len,const fmpz * a,const slong * j,slong lena)455 _fmpz_mod_poly_sqrtmod_2(fmpz *rop, const fmpz *op, slong len,
456                          const fmpz *a, const slong *j, slong lena)
457 {
458     const fmpz_t p = {WORD(2)};
459     const slong d   = j[lena - 1];
460 
461     fmpz_t z;
462 
463     fmpz_init(z);
464     fmpz_setbit(z, d - 1);
465     _qadic_pow(rop, op, len, z, a, j, lena, p);
466     fmpz_clear(z);
467 }
468 
469 /*
470     Returns whether \code{(op, len)} is a square, and if so
471     sets \code{(rop, 2 * d - 1)} to a square root mod $p^N$.
472 
473     Assumes that \code{(op, len)} has valuation $0$.
474  */
475 static int
_qadic_sqrt_p(fmpz * rop,const fmpz * op,slong len,const fmpz * a,const slong * j,slong lena,const fmpz_t p,slong N)476 _qadic_sqrt_p(fmpz *rop, const fmpz *op, slong len,
477               const fmpz *a, const slong *j, slong lena,
478               const fmpz_t p, slong N)
479 {
480     const slong d = j[lena - 1];
481     int ans;
482 
483     if (N == 1)
484     {
485         ans = _fmpz_mod_poly_sqrtmod_p(rop, op, len, a, j, lena, p);
486         return ans;
487     }
488     else
489     {
490         slong *e, i, k, n;
491         fmpz *pow, *u;
492         fmpz *r, *s, *t, *f;
493 
494         n = FLINT_CLOG2(N) + 1;
495 
496         /* Compute sequence of exponents */
497         e = flint_malloc(n * sizeof(slong));
498         for (e[i = 0] = N; e[i] > 1; i++)
499             e[i + 1] = (e[i] + 1) / 2;
500 
501         pow = _fmpz_vec_init(n);
502         u   = _fmpz_vec_init(len * n);
503         r   = _fmpz_vec_init(2 * d - 1);
504         s   = _fmpz_vec_init(2 * d - 1);
505         t   = _fmpz_vec_init(2 * d - 1);
506         f   = _fmpz_vec_init(d + 1);
507 
508         /* Compute powers of p */
509         {
510             fmpz_one(t);
511             fmpz_set(pow + i, p);
512         }
513         for (i--; i >= 1; i--)
514         {
515             if (e[i] & WORD(1))
516             {
517                 fmpz_mul(pow + i, t, pow + (i + 1));
518                 fmpz_mul(t, t, t);
519             }
520             else
521             {
522                 fmpz_mul(t, t, pow + (i + 1));
523                 fmpz_mul(pow + i, pow + (i + 1), pow + (i + 1));
524             }
525         }
526         {
527             if (e[i] & WORD(1))
528                 fmpz_mul(pow + i, t, pow + (i + 1));
529             else
530                 fmpz_mul(pow + i, pow + (i + 1), pow + (i + 1));
531         }
532 
533         /* Compute reduced units */
534         {
535             _fmpz_vec_scalar_mod_fmpz(u + 0 * len, op, len, pow + 0);
536         }
537         for (i = 1; i < n; i++)
538         {
539             _fmpz_vec_scalar_mod_fmpz(u + i * len, u + (i - 1) * len, len, pow + i);
540         }
541 
542         /* Run Newton iteration */
543         i = n - 1;
544         {
545             ans = _fmpz_mod_poly_sqrtmod_p(t, u + i * len, len, a, j, lena, p);
546             if (!ans)
547                 goto exit;
548 
549             /* Dense copy of f, used for inversion */
550             for (k = 0; k < lena; k++)
551                 fmpz_set(f + j[k], a + k);
552             _fmpz_mod_poly_invmod(rop, t, d, f, d + 1, p);
553         }
554         for (i--; i >= 1; i--)  /* z := z - z (a z^2 - 1) / 2 */
555         {
556             _fmpz_poly_sqr(s, rop, d);
557             _fmpz_poly_reduce(s, 2 * d - 1, a, j, lena);
558             _fmpz_poly_mul(t, s, d, u + i * len, len);
559             _fmpz_poly_reduce(t, d + len - 1, a, j, lena);
560             fmpz_sub_ui(t, t, 1);
561 
562             for (k = 0; k < d; k++)
563             {
564                 if (fmpz_is_odd(t + k))
565                     fmpz_add(t + k, t + k, pow + i);
566                 fmpz_fdiv_q_2exp(t + k, t + k, 1);
567             }
568 
569             _fmpz_poly_mul(s, t, d, rop, d);
570             _fmpz_poly_reduce(s, 2 * d - 1, a, j, lena);
571             _fmpz_poly_sub(rop, rop, d, s, d);
572             _fmpz_vec_scalar_mod_fmpz(rop, rop, d, pow + i);
573         }
574         {
575             _fmpz_poly_mul(s, rop, d, u + 1 * len, len);
576             _fmpz_poly_reduce(s, d + len - 1, a, j, lena);
577             _fmpz_poly_sqr(t, s, d);
578             _fmpz_poly_reduce(t, 2 * d - 1, a, j, lena);
579             _fmpz_poly_sub(t, u + 0 * len, len, t, d);
580 
581             for (k = 0; k < d; k++)
582             {
583                 if (fmpz_is_odd(t + k))
584                     fmpz_add(t + k, t + k, pow + 0);
585                 fmpz_fdiv_q_2exp(t + k, t + k, 1);
586             }
587 
588             _fmpz_poly_mul(r, rop, d, t, d);
589             _fmpz_poly_reduce(r, 2 * d - 1, a, j, lena);
590             _fmpz_poly_add(rop, r, d, s, d);
591             _fmpz_vec_scalar_mod_fmpz(rop, rop, d, pow + 0);
592         }
593 
594       exit:
595 
596         _fmpz_vec_clear(pow, n);
597         _fmpz_vec_clear(u, len * n);
598         _fmpz_vec_clear(r, 2 * d - 1);
599         _fmpz_vec_clear(s, 2 * d - 1);
600         _fmpz_vec_clear(t, 2 * d - 1);
601         _fmpz_vec_clear(f, d + 1);
602         flint_free(e);
603 
604         return ans;
605     }
606 }
607 
608 /*
609     Returns whether \code{(op, len)} is a square, and if so
610     sets \code{(rop, 2 * d - 1)} to a square root mod $2^N$.
611 
612     Assumes that \code{(op, len)} has valuation $0$.
613  */
614 static int
_qadic_sqrt_2(fmpz * rop,const fmpz * op,slong len,const fmpz * a,const slong * j,slong lena,slong N)615 _qadic_sqrt_2(fmpz *rop, const fmpz *op, slong len,
616               const fmpz *a, const slong *j, slong lena, slong N)
617 {
618     int ans;
619 
620     const slong d    = j[lena - 1];
621     const fmpz_t p  = {WORD(2)};
622 
623     fmpz *g, *r, *s, *t;
624     slong i;
625 
626     g = _fmpz_vec_init(2 * d - 1);
627     r = _fmpz_vec_init(2 * d - 1);
628     s = _fmpz_vec_init(2 * d - 1);
629     t = _fmpz_vec_init(2 * d - 1);
630 
631     /* Compute r = 1/op (mod 8) */
632     _qadic_inv(r, op, len, a, j, lena, p, 3);
633 
634     /* Compute s = invsqrt(op) (mod 2) */
635     _fmpz_vec_scalar_fdiv_r_2exp(t, r, d, 1);
636     _fmpz_mod_poly_sqrtmod_2(s, r, d, a, j, lena);
637 
638     /* Compute t = (1/op - s^2) / 4 (mod 2) if that makes sense */
639     _fmpz_poly_sqr(t, s, d);
640     _fmpz_poly_reduce(t, 2*d - 1, a, j, lena);
641     _fmpz_vec_sub(t, r, t, d);
642 
643     _fmpz_vec_zero(rop, 2 * d - 1);
644     ans = 1;
645     for (i = 0; i < d; i++)
646     {
647         if (fmpz_val2(t + i) == 1)
648             ans = 0;
649     }
650 
651     if (ans) {
652         _fmpz_vec_scalar_fdiv_q_2exp(t, t, d, 2);
653         _fmpz_vec_scalar_fdiv_r_2exp(t, t, d, 1);
654 
655         /* Check whether X^2 + s X = t (mod 2) has a solution */
656         /*
657             u = (op,len) is a square in Zq iff it is a square modulo 8,
658             which is the case iff X^2 + s X + t == 0 (mod 2) is soluble.
659             Let g be t/s^2.  Then the above quadratic is soluble iff
660             Y^2 + Y + d = 0 is soluble.
661 
662             We can observe that 1/s^2 = op (mod 2).
663          */
664         _fmpz_vec_scalar_fdiv_r_2exp(r, op, len, 1);
665         _fmpz_poly_mul(g, t, d, r, len);
666         _fmpz_mod_poly_reduce(g, 2 * d - 1, a, j, lena, p);
667 
668         ans = _artin_schreier_preimage(r, g, d, a, j, lena);
669 
670         if (ans)
671         {
672             if (N == 1)
673             {
674                 _fmpz_mod_poly_sqrtmod_2(rop, op, len, a, j, lena);
675             }
676             else
677             {
678 
679                 /* Set t = r s, a root of X^2 + s X = t (mod 2) */
680                 _fmpz_poly_mul(t, r, d, s, d);
681                 _fmpz_mod_poly_reduce(t, 2 * d - 1, a, j, lena, p);
682 
683                 /* Now s + 2t is an invsqrt of (op, len) to precision 4. */
684                 _fmpz_vec_scalar_addmul_si(s, t, d, 2);
685 
686                 if (N == 2)
687                 {
688                     _qadic_inv(rop, s, d, a, j, lena, p, 2);
689                 }
690                 else  /* N >= 3 */
691                 {
692                     slong *e, n;
693                     fmpz *u;
694 
695                     n = FLINT_CLOG2(N - 1) + 1;
696 
697                     /* Compute sequence of exponents */
698                     e = flint_malloc(n * sizeof(slong));
699                     for (e[i = 0] = N; e[i] > 2; i++)
700                         e[i + 1] = e[i] / 2 + 1;
701 
702                     u = _fmpz_vec_init(len * n);
703 
704                     /* Compute reduced units */
705                     {
706                         _fmpz_vec_scalar_fdiv_r_2exp(u + 0 * len, op, len, e[0]);
707                     }
708                     for (i = 1; i < n; i++)
709                     {
710                         _fmpz_vec_scalar_fdiv_r_2exp(u + i * len, u + (i - 1) * len, len, e[i] + 1);
711                     }
712 
713                     /* Run Newton iteration */
714                     {
715                         _fmpz_vec_set(rop, s, d);
716                     }
717                     for (i = n - 2; i >= 1; i--)  /* z := z - z (a z^2 - 1) / 2 */
718                     {
719                         _fmpz_poly_sqr(s, rop, d);
720                         _fmpz_poly_reduce(s, 2 * d - 1, a, j, lena);
721                         _fmpz_poly_mul(t, s, d, u + i * len, len);
722                         _fmpz_poly_reduce(t, d + len - 1, a, j, lena);
723                         fmpz_sub_ui(t, t, 1);
724                         _fmpz_vec_scalar_fdiv_q_2exp(t, t, d, 1);
725                         _fmpz_poly_mul(s, t, d, rop, d);
726                         _fmpz_poly_reduce(s, 2 * d - 1, a, j, lena);
727                         _fmpz_poly_sub(rop, rop, d, s, d);
728                         _fmpz_vec_scalar_fdiv_r_2exp(rop, rop, d, e[i]);
729                     }
730                     /* z := a z + z (a - (a z)^2) / 2 */
731                     {
732                         _fmpz_poly_mul(s, rop, d, u + len, len);
733                         _fmpz_poly_reduce(s, d + len - 1, a, j, lena);
734                         _fmpz_poly_sqr(t, s, d);
735                         _fmpz_poly_reduce(t, 2 * d - 1, a, j, lena);
736                         _fmpz_poly_sub(t, u + 0 * len, len, t, d);
737                         _fmpz_vec_scalar_fdiv_q_2exp(t, t, d, 1);
738                         _fmpz_poly_mul(r, rop, d, t, d);
739                         _fmpz_poly_reduce(r, 2 * d - 1, a, j, lena);
740                         _fmpz_poly_add(rop, r, d, s, d);
741                         _fmpz_vec_scalar_fdiv_r_2exp(rop, rop, d, e[0]);
742                     }
743 
744                     _fmpz_vec_clear(u, len * n);
745                     flint_free(e);
746                 }
747             }
748         }
749     }
750 
751     _fmpz_vec_clear(g, 2 * d - 1);
752     _fmpz_vec_clear(r, 2 * d - 1);
753     _fmpz_vec_clear(s, 2 * d - 1);
754     _fmpz_vec_clear(t, 2 * d - 1);
755 
756     return ans;
757 }
758 
759 /*
760     Returns whether \code{(op, len)} is a square, and if so
761     sets \code{(rop, 2 * d - 1)} to a square root, reduced
762     modulo $2^N$.
763 
764     Assumes that \code{(op, len)} has valuation $0$.
765  */
_qadic_sqrt(fmpz * rop,const fmpz * op,slong len,const fmpz * a,const slong * j,slong lena,const fmpz_t p,slong N)766 int _qadic_sqrt(fmpz *rop, const fmpz *op, slong len,
767                 const fmpz *a, const slong *j, slong lena,
768                 const fmpz_t p, slong N)
769 {
770     if (*p == WORD(2))
771     {
772         return _qadic_sqrt_2(rop, op, len, a, j, lena, N);
773     }
774     else
775     {
776         return _qadic_sqrt_p(rop, op, len, a, j, lena, p, N);
777     }
778 }
779 
qadic_sqrt(qadic_t rop,const qadic_t op,const qadic_ctx_t ctx)780 int qadic_sqrt(qadic_t rop, const qadic_t op, const qadic_ctx_t ctx)
781 {
782     const fmpz *p = (&ctx->pctx)->p;
783     const slong d  = qadic_ctx_degree(ctx);
784     const slong N  = qadic_prec(rop);
785 
786     fmpz *t;
787     int ans;
788 
789     if (qadic_is_zero(op))
790     {
791         qadic_zero(rop);
792         return 1;
793     }
794     if (op->val & WORD(1))
795     {
796         return 0;
797     }
798 
799     rop->val = op->val / 2;
800 
801     /*
802         FIXME:  In this case, we don't actually
803         check whether the element is a square!
804      */
805     if (rop->val >= N)
806     {
807         qadic_zero(rop);
808         return 1;
809     }
810 
811     if (rop == op)
812     {
813         t = _fmpz_vec_init(2 * d - 1);
814     }
815     else
816     {
817         padic_poly_fit_length(rop, 2 * d - 1);
818         t = rop->coeffs;
819     }
820 
821     ans = _qadic_sqrt(t, op->coeffs, op->length, ctx->a, ctx->j, ctx->len, p, N - rop->val);
822 
823     if (rop == op)
824     {
825         _fmpz_vec_clear(rop->coeffs, rop->alloc);
826         rop->coeffs = t;
827         rop->alloc  = 2 * d - 1;
828         rop->length = d;
829     }
830     _padic_poly_set_length(rop, d);
831     _padic_poly_normalise(rop);
832     if (padic_poly_length(rop) == 0)
833         padic_poly_val(rop) = 0;
834 
835     return ans;
836 }
837