1 /*
2     Copyright (C) 2020 Daniel Schultz
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 "fq_nmod_mpoly.h"
13 
14 
_rbnode_clear_sp(mpoly_rbtree_t tree,mpoly_rbnode_t node,slong s,fq_nmod_poly_t l,const fq_nmod_poly_t x,const fq_nmod_mpoly_ctx_t ctx)15 static void _rbnode_clear_sp(mpoly_rbtree_t tree, mpoly_rbnode_t node,
16                    slong s, fq_nmod_poly_t l, const fq_nmod_poly_t x,
17                                                  const fq_nmod_mpoly_ctx_t ctx)
18 {
19     fq_nmod_poly_t r, xp;
20     slong e = node->key;
21     FLINT_ASSERT(e >= s);
22 
23     fq_nmod_poly_init(r, ctx->fqctx);
24     fq_nmod_poly_zero(r, ctx->fqctx);
25     if (node->right != tree->null)
26         _rbnode_clear_sp(tree, node->right, e, r, x, ctx);
27 
28     fq_nmod_poly_zero(l, ctx->fqctx);
29     if (node->left != tree->null)
30         _rbnode_clear_sp(tree, node->left, s, l, x, ctx);
31 
32     fq_nmod_poly_init(xp, ctx->fqctx);
33     fq_nmod_poly_pow(xp, x, e - s, ctx->fqctx);
34 
35     fq_nmod_poly_add(r, r, node->data, ctx->fqctx);
36     fq_nmod_poly_mul(r, xp, r, ctx->fqctx);
37     fq_nmod_poly_add(l, l, r, ctx->fqctx);
38 
39     fq_nmod_poly_clear(r, ctx->fqctx);
40     fq_nmod_poly_clear(xp, ctx->fqctx);
41     fq_nmod_poly_clear(node->data, ctx->fqctx);
42     flint_free(node->data);
43     flint_free(node);
44 }
45 
46 
_fq_nmod_mpoly_compose_fq_nmod_poly_sp(fq_nmod_poly_t A,const fq_nmod_mpoly_t B,fq_nmod_poly_struct * const * C,const fq_nmod_mpoly_ctx_t ctx)47 int _fq_nmod_mpoly_compose_fq_nmod_poly_sp(fq_nmod_poly_t A, const fq_nmod_mpoly_t B,
48                 fq_nmod_poly_struct * const * C, const fq_nmod_mpoly_ctx_t ctx)
49 {
50     int success = 1;
51     int new;
52     slong i, j, k, N, bits, nvars = ctx->minfo->nvars;
53     slong main_exp, main_var, main_shift, main_off, shift, off;
54     ulong mask;
55     slong entries, k_len;
56     slong Blen;
57     const fq_nmod_struct * Bcoeff;
58     ulong * Bexp;
59     slong * degrees;
60     slong * offs;
61     ulong * masks;
62     fq_nmod_poly_struct * powers;
63     fq_nmod_poly_t t, t2;
64     mpoly_rbtree_t tree;
65     mpoly_rbnode_struct * node;
66     TMP_INIT;
67 
68     Blen = B->length;
69     Bcoeff = B->coeffs;
70     Bexp = B->exps;
71     bits = B->bits;
72 
73     FLINT_ASSERT(Blen != 0);
74 
75     TMP_START;
76 
77     degrees = (slong *) TMP_ALLOC(nvars*sizeof(slong));
78     mpoly_degrees_si(degrees, Bexp, Blen, bits, ctx->minfo);
79 
80     /* pick main variable with highest degree */
81     main_var = 0;
82     for (i = 1; i < nvars; i++)
83     {
84         if (degrees[i] > degrees[main_var])
85             main_var = i;
86     }
87 
88     /* compute how many masks are needed */
89     entries = 0;
90     for (i = 0; i < nvars; i++)
91     {
92         if (_ff_poly_pow_ui_is_not_feasible(C[i]->length, degrees[i]))
93         {
94             success = 0;
95             goto cleanup_degrees;
96         }
97 
98         if (i == main_var)
99             continue;
100 
101         entries += FLINT_BIT_COUNT(degrees[i]);
102     }
103     offs = (slong *) TMP_ALLOC(entries*sizeof(slong));
104     masks = (ulong *) TMP_ALLOC(entries*sizeof(slong));
105     powers = (fq_nmod_poly_struct *) TMP_ALLOC(entries*
106                                                   sizeof(fq_nmod_poly_struct));
107 
108     N = mpoly_words_per_exp(bits, ctx->minfo);
109 
110     /* store bit masks for each power of two of the non-main variables */
111     k = 0;
112     for (i = 0; i < nvars; i++)
113     {
114         flint_bitcnt_t varibits;
115 
116         if (i == main_var)
117             continue;
118 
119         mpoly_gen_offset_shift_sp(&off, &shift, i, bits, ctx->minfo);
120         varibits = FLINT_BIT_COUNT(degrees[i]);
121         for (j = 0; j < varibits; j++)
122         {
123             offs[k] = off;
124             masks[k] = UWORD(1) << (shift + j);
125             fq_nmod_poly_init(powers + k, ctx->fqctx);
126             if (j == 0)
127                 fq_nmod_poly_set(powers + k, C[i], ctx->fqctx);
128             else
129                 fq_nmod_poly_mul(powers + k, powers + k - 1, powers + k - 1,
130                                                                    ctx->fqctx);
131             k++;
132         }
133     }
134     k_len = k;
135     FLINT_ASSERT(k_len == entries);
136 
137     /* accumulate coefficients of the main variable */
138     mpoly_gen_offset_shift_sp(&main_off, &main_shift, main_var, bits, ctx->minfo);
139     mpoly_rbtree_init(tree);
140     fq_nmod_poly_init(t, ctx->fqctx);
141     fq_nmod_poly_init(t2, ctx->fqctx);
142     mask = (-UWORD(1)) >> (FLINT_BITS - bits);
143     for (i = 0; i < Blen; i++)
144     {
145         main_exp = (Bexp[N*i + main_off] >> main_shift) & mask;
146         node = mpoly_rbtree_get(&new, tree, main_exp);
147         if (new)
148         {
149             node->data = flint_malloc(sizeof(fq_nmod_poly_struct));
150             fq_nmod_poly_init(node->data, ctx->fqctx);
151             fq_nmod_poly_zero(node->data, ctx->fqctx);
152         }
153 
154         fq_nmod_poly_zero(t, ctx->fqctx);
155         fq_nmod_poly_set_coeff(t, 0, Bcoeff + i, ctx->fqctx);
156         for (k = 0; k < k_len; k++)
157         {
158             if ((Bexp[N*i + offs[k]] & masks[k]) != WORD(0))
159             {
160                 fq_nmod_poly_mul(t2, t, powers + k, ctx->fqctx);
161                 fq_nmod_poly_swap(t, t2, ctx->fqctx);
162             }
163         }
164         fq_nmod_poly_add(t2, t, node->data, ctx->fqctx);
165         fq_nmod_poly_swap(t2, node->data, ctx->fqctx);
166     }
167     fq_nmod_poly_clear(t, ctx->fqctx);
168     fq_nmod_poly_clear(t2, ctx->fqctx);
169 
170     for (k = 0; k < k_len; k++)
171         fq_nmod_poly_clear(powers + k, ctx->fqctx);
172 
173     /* use tree method to evaluate in the main variable */
174     _rbnode_clear_sp(tree, tree->head->left, WORD(0), A, C[main_var], ctx);
175 
176 cleanup_degrees:
177 
178     TMP_END;
179 
180     return success;
181 }
182 
183 
_rbnode_clear_mp(mpoly_rbtree_t tree,mpoly_rbnode_t node,const fmpz_t s,fq_nmod_poly_t l,const fq_nmod_poly_t x,const fq_nmod_mpoly_ctx_t ctx)184 static int _rbnode_clear_mp(mpoly_rbtree_t tree, mpoly_rbnode_t node,
185                      const fmpz_t s, fq_nmod_poly_t l, const fq_nmod_poly_t x,
186                                                  const fq_nmod_mpoly_ctx_t ctx)
187 {
188     int success = 1;
189     fq_nmod_poly_t r, xp;
190     FLINT_ASSERT(fmpz_cmp(&node->key, s) >= 0);
191 
192     fq_nmod_poly_init(r, ctx->fqctx);
193     fq_nmod_poly_zero(r, ctx->fqctx);
194     if (node->right != tree->null)
195     {
196         if (!_rbnode_clear_mp(tree, node->right, &node->key, r, x, ctx))
197             success = 0;
198     }
199 
200     fq_nmod_poly_zero(l, ctx->fqctx);
201     if (node->left != tree->null)
202     {
203         if (!_rbnode_clear_mp(tree, node->left, s, l, x, ctx))
204             success = 0;
205     }
206 
207     fq_nmod_poly_init(xp, ctx->fqctx);
208     fmpz_sub(&node->key, &node->key, s);
209     FLINT_ASSERT(fmpz_sgn(&node->key) >= 0);
210     if (fmpz_fits_si(&node->key))
211     {
212         fq_nmod_poly_pow(xp, x, fmpz_get_si(&node->key), ctx->fqctx);
213     }
214     else
215     {
216         slong degree = fq_nmod_poly_degree(x, ctx->fqctx);
217         fq_nmod_poly_zero(xp, ctx->fqctx);
218         if (degree == 0)
219         {
220             fq_nmod_t t;
221             fq_nmod_init(t, ctx->fqctx);
222             fq_nmod_pow(t, x->coeffs + 0, &node->key, ctx->fqctx);
223             fq_nmod_poly_set_coeff(xp, 0, t, ctx->fqctx);
224             fq_nmod_clear(t, ctx->fqctx);
225         }
226         else if (degree > 0)
227         {
228             success = 0;
229         }
230     }
231     fq_nmod_poly_add(r, r, node->data, ctx->fqctx);
232     fq_nmod_poly_mul(r, xp, r, ctx->fqctx);
233     fq_nmod_poly_add(l, l, r, ctx->fqctx);
234 
235     fmpz_clear(&node->key);
236     fq_nmod_poly_clear(r, ctx->fqctx);
237     fq_nmod_poly_clear(xp, ctx->fqctx);
238     fq_nmod_poly_clear(node->data, ctx->fqctx);
239     flint_free(node->data);
240     flint_free(node);
241 
242     return success;
243 }
244 
245 
_fq_nmod_mpoly_compose_fq_nmod_poly_mp(fq_nmod_poly_t A,const fq_nmod_mpoly_t B,fq_nmod_poly_struct * const * C,const fq_nmod_mpoly_ctx_t ctx)246 int _fq_nmod_mpoly_compose_fq_nmod_poly_mp(fq_nmod_poly_t A, const fq_nmod_mpoly_t B,
247                 fq_nmod_poly_struct * const * C, const fq_nmod_mpoly_ctx_t ctx)
248 {
249     int success = 1;
250     int new;
251     ulong l;
252     slong i, k, N, bits, nvars = ctx->minfo->nvars;
253     fmpz_t main_exp;
254     slong main_var, main_off, off;
255     slong entries, k_len;
256     slong Blen;
257     const fq_nmod_struct * Bcoeff;
258     ulong * Bexp;
259     fmpz * degrees;
260     slong * offs;
261     ulong * masks;
262     flint_bitcnt_t * bitcounts;
263     fmpz_t s;
264     fq_nmod_poly_struct * powers;
265     fq_nmod_poly_t t, t2;
266     mpoly_rbtree_t tree;
267     mpoly_rbnode_struct * node;
268     TMP_INIT;
269 
270     Blen = B->length;
271     Bcoeff = B->coeffs;
272     Bexp = B->exps;
273     bits = B->bits;
274 
275     FLINT_ASSERT(Blen != 0);
276 
277     TMP_START;
278 
279     bitcounts = (flint_bitcnt_t * ) TMP_ALLOC(nvars*sizeof(flint_bitcnt_t));
280     degrees = (fmpz *) TMP_ALLOC(nvars*sizeof(fmpz));
281     for (i = 0; i < nvars; i++)
282     {
283         fmpz_init(degrees + i);
284     }
285     mpoly_degrees_ffmpz(degrees, Bexp, Blen, bits, ctx->minfo);
286 
287     /* pick main variable with highest degree */
288     main_var = 0;
289     for (i = 1; i < nvars; i++)
290     {
291         if (fmpz_cmp(degrees + i, degrees + main_var) > 0)
292             main_var = i;
293     }
294 
295     /* compute how many masks are needed */
296     entries = 0;
297     for (i = 0; i < nvars; i++)
298     {
299         if (_ff_poly_pow_fmpz_is_not_feasible(C[i]->length, degrees + i))
300         {
301             success = 0;
302             goto cleanup_degrees;
303         }
304 
305         bitcounts[i] = fmpz_bits(degrees + i);
306 
307         if (i == main_var)
308             continue;
309 
310         entries += bitcounts[i];
311     }
312     offs = (slong *) TMP_ALLOC(entries*sizeof(slong));
313     masks = (ulong *) TMP_ALLOC(entries*sizeof(slong));
314     powers = (fq_nmod_poly_struct *) TMP_ALLOC(entries*
315                                                   sizeof(fq_nmod_poly_struct));
316 
317     N = mpoly_words_per_exp(bits, ctx->minfo);
318 
319     /* store bit masks for each power of two of the non-main variables */
320     k = 0;
321     for (i = 0; i < nvars; i++)
322     {
323         if (i == main_var)
324             continue;
325 
326         off = mpoly_gen_offset_mp(i, bits, ctx->minfo);
327 
328         for (l = 0; l < bitcounts[i]; l++)
329         {
330             offs[k] = off + (l/FLINT_BITS);
331             masks[k] = UWORD(1) << (l%FLINT_BITS);
332             fq_nmod_poly_init(powers + k, ctx->fqctx);
333             if (l == 0)
334                 fq_nmod_poly_set(powers + k, C[i], ctx->fqctx);
335             else
336                 fq_nmod_poly_mul(powers + k, powers + k - 1, powers + k - 1,
337                                                                    ctx->fqctx);
338             k++;
339         }
340     }
341     k_len = k;
342     FLINT_ASSERT(k_len == entries);
343 
344     /* accumulate coefficients of the main variable */
345     main_off = mpoly_gen_offset_mp(main_var, bits, ctx->minfo);
346     mpoly_rbtree_init(tree);
347     fq_nmod_poly_init(t, ctx->fqctx);
348     fq_nmod_poly_init(t2, ctx->fqctx);
349     fmpz_init(main_exp);
350     for (i = 0; i < Blen; i++)
351     {
352         fmpz_set_ui_array(main_exp, Bexp + N*i + main_off, bits/FLINT_BITS);
353         node = mpoly_rbtree_get_fmpz(&new, tree, main_exp);
354         if (new)
355         {
356             node->data = flint_malloc(sizeof(fq_nmod_poly_struct));
357             fq_nmod_poly_init(node->data, ctx->fqctx);
358             fq_nmod_poly_zero(node->data, ctx->fqctx);
359         }
360 
361         fq_nmod_poly_zero(t, ctx->fqctx);
362         fq_nmod_poly_set_coeff(t, 0, Bcoeff + i, ctx->fqctx);
363         for (k = 0; k < k_len; k++)
364         {
365             if ((Bexp[N*i + offs[k]] & masks[k]) != WORD(0))
366             {
367                 fq_nmod_poly_mul(t2, t, powers + k, ctx->fqctx);
368                 fq_nmod_poly_swap(t, t2, ctx->fqctx);
369             }
370         }
371         fq_nmod_poly_add(t2, t, node->data, ctx->fqctx);
372         fq_nmod_poly_swap(t2, node->data, ctx->fqctx);
373     }
374     fmpz_clear(main_exp);
375     fq_nmod_poly_clear(t, ctx->fqctx);
376     fq_nmod_poly_clear(t2, ctx->fqctx);
377 
378     for (k = 0; k < k_len; k++)
379         fq_nmod_poly_clear(powers + k, ctx->fqctx);
380 
381     /* use tree method to evaluate in the main variable */
382     fmpz_init(s);
383     if (!_rbnode_clear_mp(tree, tree->head->left, s, A, C[main_var], ctx))
384         success = 0;
385     fmpz_clear(s);
386 
387 cleanup_degrees:
388 
389     for (i = 0; i < nvars; i++)
390         fmpz_clear(degrees + i);
391 
392     TMP_END;
393 
394     return success;
395 }
396 
397 
fq_nmod_mpoly_compose_fq_nmod_poly(fq_nmod_poly_t A,const fq_nmod_mpoly_t B,fq_nmod_poly_struct * const * C,const fq_nmod_mpoly_ctx_t ctx)398 int fq_nmod_mpoly_compose_fq_nmod_poly(fq_nmod_poly_t A,
399                     const fq_nmod_mpoly_t B, fq_nmod_poly_struct * const * C,
400                                                  const fq_nmod_mpoly_ctx_t ctx)
401 {
402     if (B->length == 0)
403     {
404         fq_nmod_poly_zero(A, ctx->fqctx);
405         return 1;
406     }
407 
408     if (B->bits <= FLINT_BITS)
409     {
410         return _fq_nmod_mpoly_compose_fq_nmod_poly_sp(A, B, C, ctx);
411     }
412     else
413     {
414         return _fq_nmod_mpoly_compose_fq_nmod_poly_mp(A, B, C, ctx);
415     }
416 }
417 
418