Ticket #9051: 9051-FpT_1.patch

File 9051-FpT_1.patch, 25.9 KB (added by robertwb, 12 years ago)
  • module_list.py

    # HG changeset patch
    # User Robert Bradshaw <robertwb@math.washington.edu>
    # Date 1260035621 28800
    # Node ID a2a1e2eafbcaf6fdb624db1cc2b6c3f2bffc1796
    # Parent  2ef9d956a42f1bb80136513dbc8ee3210d7e0735
    #7585 Fast arithmatic in F_p(t).
    * * *
    #7585 More work on F_p(t)
    
    Fix memory leak, misc speedups.
    
    diff -r 2ef9d956a42f -r a2a1e2eafbca module_list.py
    a b  
    10361036    Extension('sage.rings.fraction_field_element',
    10371037              sources = ['sage/rings/fraction_field_element.pyx']),
    10381038   
     1039    Extension('sage.rings.fraction_field_FpT',
     1040              sources = ['sage/rings/fraction_field_FpT.pyx'],
     1041              libraries = ["csage", "flint", "gmp", "gmpxx", "ntl", "zn_poly"],
     1042              extra_compile_args=["-std=c99", "-D_XPG6"],
     1043              include_dirs = [SAGE_ROOT+'/local/include/FLINT/'],
     1044              depends = [SAGE_ROOT + "/local/include/FLINT/flint.h"]),
     1045   
    10391046    Extension('sage.rings.laurent_series_ring_element',
    10401047              sources = ['sage/rings/laurent_series_ring_element.pyx']),
    10411048
  • sage/rings/finite_rings/integer_mod.pxd

    diff -r 2ef9d956a42f -r a2a1e2eafbca sage/rings/finite_rings/integer_mod.pxd
    a b  
    4747    cdef void set_from_int(IntegerMod_int64 self, int_fast64_t value)
    4848    cdef int_fast64_t get_int_value(IntegerMod_int64 self)
    4949    cdef IntegerMod_int64 _new_c(self, int_fast64_t value)
     50
     51cdef int jacobi_int(int_fast32_t a, int_fast32_t m) except -2
     52cdef int_fast32_t mod_inverse_int(int_fast32_t x, int_fast32_t n) except 0
     53cdef int_fast32_t mod_pow_int(int_fast32_t base, int_fast32_t exp, int_fast32_t n)
    5054    cdef shift(IntegerMod_int64 self, int k)
    5155
  • new file sage/rings/fraction_field_FpT.pyx

    diff -r 2ef9d956a42f -r a2a1e2eafbca sage/rings/fraction_field_FpT.pyx
    - +  
     1import sys
     2
     3from sage.rings.all import GF
     4from sage.rings.ring cimport Field
     5from sage.libs.flint.zmod_poly cimport *
     6from sage.structure.element cimport Element, ModuleElement, RingElement
     7import sage.algebras.algebra
     8
     9from sage.rings.integer_mod cimport jacobi_int, mod_inverse_int, mod_pow_int
     10
     11class Fpt(Field):
     12    def __init__(self, long p, names='t'):
     13        assert 2 < p < 2**16
     14        if isinstance(names, tuple):
     15            names, = names
     16        self.p = p
     17        base_ring = GF(p)
     18        self.poly_ring = base_ring[names]
     19        sage.algebras.algebra.Algebra.__init__(self, base_ring, names=names, normalize=True)
     20        self._populate_coercion_lists_(element_constructor=FptElement)
     21   
     22#    def _element_constructor_(self, *args):
     23#        return FptElement(self, *args)
     24   
     25    def ngens(self):
     26        return 1
     27   
     28    def gen(self, ix):
     29        assert ix == 0
     30        return FptElement(self, self.poly_ring.gen())
     31   
     32    def characteristic(self):
     33        return self.p
     34   
     35    def _coerce_map_from_(self, R):
     36        return self.poly_ring.has_coerce_map_from(R)
     37   
     38    def __iter__(self):
     39        return self.iter()
     40   
     41    def iter(self, bound=None, start=None):
     42        """
     43            sage: from sage.rings.fraction_field_FpT import *
     44            sage: R.<t> = Fpt(5)
     45            sage: list(R.iter(2))
     46        """
     47        return Fpt_iter(self, bound, start)
     48
     49cdef class FptElement(RingElement):
     50
     51    cdef long p
     52    cdef zmod_poly_t _numer, _denom
     53    cdef bint initalized
     54
     55    def __init__(self, parent, numer, denom=1):
     56        """
     57        EXAMPLES::
     58       
     59            sage: from sage.rings.fraction_field_FpT import *
     60            sage: R.<t> = Fpt(5)
     61            sage: R(7)
     62            2
     63           
     64        """
     65        RingElement.__init__(self, parent)
     66        numer = parent.poly_ring(numer)
     67        denom = parent.poly_ring(denom)
     68        self.p = parent.p
     69        zmod_poly_init(self._numer, self.p)
     70        zmod_poly_init(self._denom, self.p)
     71        self.initalized = True
     72        cdef long n
     73        for n, a in enumerate(numer):
     74            zmod_poly_set_coeff_ui(self._numer, n, a)
     75        for n, a in enumerate(denom):
     76            zmod_poly_set_coeff_ui(self._denom, n, a)
     77        normalize(self._numer, self._denom, self.p)
     78
     79    def __dealloc__(self):
     80        if self.initalized:
     81            zmod_poly_clear(self._numer)
     82            zmod_poly_clear(self._denom)
     83   
     84    cdef FptElement _new_c(self):
     85        cdef FptElement x = <FptElement>PY_NEW(FptElement)
     86        x._parent = self._parent
     87        x.p = self.p
     88        zmod_poly_init_precomp(x._numer, x.p, self._numer.p_inv)
     89        zmod_poly_init_precomp(x._denom, x.p, self._numer.p_inv)
     90        x.initalized = True
     91        return x
     92   
     93    cdef FptElement _copy_c(self):
     94        cdef FptElement x = <FptElement>PY_NEW(FptElement)
     95        x._parent = self._parent
     96        x.p = self.p
     97        zmod_poly_init2_precomp(x._numer, x.p, self._numer.p_inv, self._numer.length)
     98        zmod_poly_init2_precomp(x._denom, x.p, self._denom.p_inv, self._denom.length)
     99        zmod_poly_set(x._numer, self._numer)
     100        zmod_poly_set(x._denom, self._denom)
     101        x.initalized = True
     102        return x
     103   
     104    def numer(self):
     105        cdef long n
     106        return self._parent.poly_ring(
     107            [zmod_poly_get_coeff_ui (self._numer, n) for n in range(0, zmod_poly_degree(self._numer)+1)])
     108
     109    def denom(self):
     110        cdef long n
     111        return self._parent.poly_ring(
     112            [zmod_poly_get_coeff_ui (self._denom, n) for n in range(0, zmod_poly_degree(self._denom)+1)])
     113
     114    def _repr_(self):
     115        """
     116            sage: from sage.rings.fraction_field_FpT import *
     117            sage: R.<t> = Fpt(17)
     118            sage: -t
     119            16*t
     120            sage: 1/t
     121            1/t
     122            sage: 1/(t+1)
     123            1/(t + 1)
     124            sage: 1-t/t
     125            0
     126            sage: (1-t)/t
     127            (16*t + 1)/t
     128        """
     129        if zmod_poly_degree(self._denom) == 0 and zmod_poly_get_coeff_ui(self._denom, 0) == 1:
     130            return repr(self.numer())
     131        else:
     132            numer_s = repr(self.numer())
     133            denom_s = repr(self.denom())
     134            if '-' in numer_s or '+' in numer_s:
     135                numer_s = "(%s)" % numer_s
     136            if '-' in denom_s or '+' in denom_s:
     137                denom_s = "(%s)" % denom_s
     138            return "%s/%s" % (numer_s, denom_s)
     139   
     140    def _latex_(self):
     141        if zmod_poly_degree(self._denom) == 0 and zmod_poly_get_coeff_ui(self._denom, 0) == 1:
     142            return self.numer()._latex_()
     143        else:
     144            return "{%s}/{%s}" % (self.numer()._latex_(), self.denom()._latex_())
     145        return
     146   
     147    def __cmp__(self, other):
     148        """
     149        TESTS::
     150       
     151            sage: from sage.rings.fraction_field_FpT import *
     152            sage: R.<t> = Fpt(7)
     153            sage: t == t
     154            True
     155            sage: t == -t
     156            False
     157            sage: -t == 6*t
     158            True
     159            sage: 1/t == 1/t
     160            True
     161            sage: 1/t == 1/(t+1)
     162            False
     163        """
     164        if isinstance(other, FptElement):
     165            # They are normalized.
     166            if (zmod_poly_equal(self._numer, (<FptElement>other)._numer) and
     167                    zmod_poly_equal(self._denom, (<FptElement>other)._denom)):
     168                return 0
     169            else:
     170                return -1
     171        else:
     172            return cmp(type(self), type(other))
     173   
     174    def __hash__(self):
     175        """
     176        TESTS::
     177           
     178            sage: from sage.rings.fraction_field_FpT import *
     179            sage: K.<t> = Fpt(7)
     180            sage: hash(K(0))
     181            0
     182            sage: hash(K(5))
     183            5
     184            sage: set([1, t, 1/t, t, t, 1/t, 1+1/t, t/t])
     185            set([1, (t + 1)/t, t, 1/t])
     186        """
     187        if self.denom() == 1:
     188            return hash(self.numer())
     189        return hash(str(self))
     190       
     191    def __neg__(self):
     192        cdef FptElement x = self._new_c()
     193        zmod_poly_neg(x._numer, self._numer)
     194        zmod_poly_set(x._denom, self._denom)
     195        return x
     196       
     197    def __invert__(self):
     198        if zmod_poly_degree(self._numer) == -1:
     199            raise ZeroDivisionError
     200        cdef FptElement x = self._new_c()
     201        zmod_poly_set(x._denom, self._numer)
     202        zmod_poly_set(x._numer, self._denom)
     203        return x
     204
     205    cpdef ModuleElement _add_(self, ModuleElement _other):
     206        """
     207        EXAMPLES::
     208       
     209            sage: from sage.rings.fraction_field_FpT import *
     210            sage: R.<t> = Fpt(7)
     211            sage: t + t
     212            2*t
     213            sage: (t + 3) + (t + 10)
     214            2*t + 6
     215            sage: sum([t] * 7)
     216            0
     217            sage: 1/t + t
     218            (t^2 + 1)/(t)
     219            sage: 1/t + 1/t^2
     220            (t + 1)/(t^2)
     221        """
     222        cdef FptElement other = <FptElement>_other
     223        cdef FptElement x = self._new_c()
     224        zmod_poly_mul(x._numer, self._numer, other._denom)
     225        zmod_poly_mul(x._denom, self._denom, other._numer) # use x._denom as a temp
     226        zmod_poly_add(x._numer, x._numer, x._denom)
     227        zmod_poly_mul(x._denom, self._denom, other._denom)
     228        normalize(x._numer, x._denom, self.p)
     229        return x
     230
     231    cpdef ModuleElement _sub_(self, ModuleElement _other):
     232        """
     233        EXAMPLES::
     234       
     235            sage: from sage.rings.fraction_field_FpT import *
     236            sage: R.<t> = Fpt(7)
     237            sage: t - t
     238            0
     239            sage: (t + 3) - (t + 11)
     240            sage: 6
     241        """
     242        cdef FptElement other = <FptElement>_other
     243        cdef FptElement x = self._new_c()
     244        zmod_poly_mul(x._numer, self._numer, other._denom)
     245        zmod_poly_mul(x._denom, self._denom, other._numer) # use x._denom as a temp
     246        zmod_poly_sub(x._numer, x._numer, x._denom)
     247        zmod_poly_mul(x._denom, self._denom, other._denom)
     248        normalize(x._numer, x._denom, self.p)
     249        return x
     250
     251    cpdef RingElement _mul_(self, RingElement _other):
     252        """
     253        EXAMPLES::
     254       
     255            sage: from sage.rings.fraction_field_FpT import *
     256            sage: R.<t> = Fpt(7)
     257            sage: t * t
     258            t^2
     259            sage: (t + 3) * (t + 10)
     260            t^2 + 6*t + 2
     261        """
     262        cdef FptElement other = <FptElement>_other
     263        cdef FptElement x = self._new_c()
     264        zmod_poly_mul(x._numer, self._numer, other._numer)
     265        zmod_poly_mul(x._denom, self._denom, other._denom)
     266        normalize(x._numer, x._denom, self.p)
     267        return x
     268
     269    cpdef RingElement _div_(self, RingElement _other):
     270        """
     271        EXAMPLES::
     272       
     273            sage: from sage.rings.fraction_field_FpT import *
     274            sage: R.<t> = Fpt(5)
     275            sage: t / t
     276            1
     277            sage: (t + 3) / (t + 11)
     278            (t + 3)/(t + 1)
     279            sage: (t^2 + 2*t + 1) / (t + 1)
     280            t + 1
     281        """
     282        cdef FptElement other = <FptElement>_other
     283        if zmod_poly_degree(other._numer) == -1:
     284            raise ZeroDivisionError
     285        cdef FptElement x = self._new_c()
     286        zmod_poly_mul(x._numer, self._numer, other._denom)
     287        zmod_poly_mul(x._denom, self._denom, other._numer)
     288        normalize(x._numer, x._denom, self.p)
     289        return x
     290   
     291    cpdef FptElement next(self):
     292        # TODO: This misses denom > numer...
     293        """
     294        EXAMPLES::
     295       
     296            sage: from sage.rings.fraction_field_FpT import *
     297            sage: R.<t> = Fpt(3)
     298            sage: a = R(0)
     299            sage: for _ in range(15):
     300            ...       a = a.next()
     301            ...       print a
     302            ...   
     303            1
     304            2
     305            t
     306            t + 1
     307            (t + 1)/(t)
     308            t + 2
     309            (t + 2)/(t)
     310            (t + 2)/(t + 1)
     311            2*t
     312            (2*t)/(t + 1)
     313            (2*t)/(t + 2)
     314            2*t + 1
     315            (2*t + 1)/(t)
     316            (2*t + 1)/(t + 1)
     317            2*t + 2
     318        """
     319        cdef FptElement x = self._new_c()
     320        if zmod_poly_degree(self._numer) == -1:
     321            zmod_poly_set_coeff_ui(x._numer, 0, 1)
     322            zmod_poly_set_coeff_ui(x._denom, 0, 1)
     323            return x
     324        elif zmod_poly_degree(self._numer) == 0:
     325            zmod_poly_set_coeff_ui(x._numer, 0, zmod_poly_get_coeff_ui(self._numer, 0) + 1)
     326            zmod_poly_set_coeff_ui(x._denom, 0, 1)
     327            if zmod_poly_get_coeff_ui(x._numer, 0) == 0:
     328                zmod_poly_set_coeff_ui(x._numer, 1, 1)
     329            return x
     330        cdef zmod_poly_t g
     331        zmod_poly_init(g, self.p)
     332        zmod_poly_set(x._numer, self._numer)
     333        zmod_poly_set(x._denom, self._denom)
     334        while True:
     335            zmod_poly_inc(x._denom, True)
     336            if zmod_poly_degree(x._numer) < zmod_poly_degree(x._denom):
     337                zmod_poly_inc(x._numer, False)
     338                zmod_poly_zero(x._denom)
     339                zmod_poly_set_coeff_ui(x._denom, 0, 1)
     340            zmod_poly_gcd(g, x._numer, x._denom)
     341            if zmod_poly_is_one(g):
     342                zmod_poly_clear(g)
     343                return x
     344   
     345    cpdef _sqrt_or_None(self):
     346        """
     347        Returns the squre root of self, or None. Differs from sqrt() by not raising an exception.
     348       
     349        TESTS::
     350       
     351            sage: from sage.rings.fraction_field_FpT import *
     352            sage: R.<t> = Fpt(17)
     353            sage: sqrt(t^2)
     354            t
     355            sage: sqrt(1/t^2)
     356            1/t
     357            sage: sqrt((1+t)^2)
     358            t + 1
     359            sage: sqrt((1+t)^2 / t^2)
     360            (t + 1)/t
     361
     362            sage: sqrt(4 * (1+t)^2 / t^2)
     363            (2*t + 2)/t
     364           
     365            sage: sqrt(R(0))
     366            0
     367            sage: sqrt(R(-1))
     368            4
     369           
     370            sage: sqrt(t^4)
     371            t^2
     372            sage: sqrt(4*t^4/(1+t)^8)
     373            4*t^2/(t^4 + 4*t^3 + 6*t^2 + 4*t + 1)
     374           
     375            sage: R.<t> = Fpt(5)
     376            sage: [a for a in R.iter(2) if (a^2).sqrt() not in (a,-a)]
     377            []
     378            sage: [a for a in R.iter(2) if a.is_square() and a.sqrt()^2 != a]
     379            []
     380
     381        """
     382        if zmod_poly_degree(self._numer) == -1:
     383            return self
     384        if not zmod_poly_sqrt_check(self._numer) or not zmod_poly_sqrt_check(self._denom):
     385            return None
     386        cdef zmod_poly_t numer
     387        cdef zmod_poly_t denom
     388        cdef FptElement res
     389        try:
     390            zmod_poly_init(denom, self.p)
     391            zmod_poly_init(numer, self.p)
     392            if not zmod_poly_sqrt0(numer, self._numer):
     393                return None
     394            if not zmod_poly_sqrt0(denom, self._denom):
     395                return None
     396            res = self._new_c()
     397            zmod_poly_swap(numer, res._numer)
     398            zmod_poly_swap(denom, res._denom)
     399            return res
     400        finally:
     401            zmod_poly_clear(numer)
     402            zmod_poly_clear(denom)
     403   
     404    cpdef bint is_square(self):
     405        return self._sqrt_or_None() is not None
     406   
     407    def sqrt(self):
     408        """
     409        EXAMPLES::
     410       
     411            sage: from sage.rings.fraction_field_FpT import *
     412            sage: R.<t> = Fpt(7)
     413        """
     414        s = self._sqrt_or_None()
     415        if s is None:
     416            raise ValueError, "not a perfect square"
     417        else:
     418            return s
     419   
     420    def __pow__(FptElement self, Py_ssize_t e, dummy):
     421        """
     422        EXAMPLES::
     423       
     424            sage: from sage.rings.fraction_field_FpT import *
     425            sage: R.<t> = Fpt(7)
     426            sage: t^5
     427            t^5
     428            sage: t^-5
     429            1/t^5
     430           
     431            sage: a = (t+1)/(t-1); a
     432            (t + 1)/(t + 6)
     433            sage: a^5
     434            (t^5 + 5*t^4 + 3*t^3 + 3*t^2 + 5*t + 1)/(t^5 + 2*t^4 + 3*t^3 + 4*t^2 + 5*t + 6)
     435            sage: a^7
     436            (t^7 + 1)/(t^7 + 6)
     437            sage: a^14
     438            (t^14 + 2*t^7 + 1)/(t^14 + 5*t^7 + 1)
     439           
     440            sage: (a^2)^2 == a^4
     441            True
     442            sage: a^3 * a^2 == a^5
     443            True
     444            sage: a^47 * a^92 == a^(47+92)
     445            True
     446        """
     447        cdef long a
     448        assert dummy is None
     449        cdef FptElement x = self._new_c()
     450        if e >= 0:
     451            zmod_poly_pow(x._numer, self._numer, e)
     452            zmod_poly_pow(x._denom, self._denom, e)
     453        else:
     454            zmod_poly_pow(x._denom, self._numer, -e)
     455            zmod_poly_pow(x._numer, self._denom, -e)
     456            if zmod_poly_leading(x._denom) != 1:
     457                a = mod_inverse_int(zmod_poly_leading(x._denom), self.p)
     458                zmod_poly_scalar_mul(x._numer, x._numer, a)
     459                zmod_poly_scalar_mul(x._denom, x._denom, a)
     460        return x
     461
     462       
     463cdef class Fpt_iter:
     464   
     465    cdef parent
     466    cdef long degree
     467    cdef FptElement cur
     468    cdef zmod_poly_t g
     469   
     470    def __init__(self, parent, degree=None, FptElement start=None):
     471        if degree is None:
     472            raise NotImplementedError
     473        self.parent = parent
     474        self.cur = start
     475        self.degree = sys.maxint if degree is None else degree
     476   
     477    def __cinit__(self, parent, *args, **kwds):
     478        zmod_poly_init(self.g, parent.characteristic())
     479   
     480    def __dealloc__(self):
     481        zmod_poly_clear(self.g)
     482   
     483    def __iter__(self):
     484        return self
     485   
     486    def __next__(self):
     487        """
     488        EXAMPLES::
     489       
     490            sage: from sage.rings.fraction_field_FpT import *
     491            sage: K.<t> = Fpt(3)
     492            sage: list(K.iter(1))
     493            [0,
     494             1,
     495             2,
     496             t,
     497             t + 1,
     498             t + 2,
     499             2*t,
     500             2*t + 1,
     501             2*t + 2,
     502             1/t,
     503             2/t,
     504             (t + 1)/t,
     505             (t + 2)/t,
     506             (2*t + 1)/t,
     507             (2*t + 2)/t,
     508             1/(t + 1),
     509             2/(t + 1),
     510             t/(t + 1),
     511             (t + 2)/(t + 1),
     512             2*t/(t + 1),
     513             (2*t + 1)/(t + 1),
     514             1/(t + 2),
     515             2/(t + 2),
     516             t/(t + 2),
     517             (t + 1)/(t + 2),
     518             2*t/(t + 2),
     519             (2*t + 2)/(t + 2)]
     520             
     521            sage: len(list(K.iter(3)))
     522            2187
     523
     524            sage: K.<t> = Fpt(5)
     525            sage: L = list(K.iter(3))
     526            78125
     527            sage: L[:10]
     528            [0, 1, 2, 3, 4, t, t + 1, t + 2, t + 3, t + 4]
     529            sage: L[2000]
     530            (3*t^3 + 3*t^2 + 3*t + 4)/(t + 2)
     531            sage: L[-1]
     532            (4*t^3 + 4*t^2 + 4*t + 4)/(t^3 + 4*t^2 + 4*t + 4)
     533        """
     534        cdef FptElement next
     535        if self.cur is None:
     536            self.cur = self.parent(0)
     537        else:
     538            next = self.cur._copy_c()
     539            while True:
     540                zmod_poly_inc(next._numer, False)
     541                if zmod_poly_degree(next._numer) > self.degree:
     542                    zmod_poly_inc(next._denom, True)
     543                    if zmod_poly_degree(next._denom) > self.degree:
     544                        raise StopIteration
     545                    zmod_poly_zero(next._numer)
     546                    zmod_poly_set_coeff_ui(next._numer, 0, 1)
     547                zmod_poly_gcd(self.g, next._numer, next._denom)
     548                if zmod_poly_is_one(self.g):
     549                    break
     550            self.cur = next
     551#            self.cur = self.cur.next()
     552#            if zmod_poly_degree(self.cur._numer) > self.degree:
     553#                raise StopIteration
     554        return self.cur
     555
     556
     557
     558cdef inline bint normalize(zmod_poly_t numer, zmod_poly_t denom, long p):
     559    cdef long a
     560    if zmod_poly_degree(numer) == -1:
     561        zmod_poly_truncate(denom, 0)
     562        zmod_poly_set_coeff_ui(denom, 0, 1)
     563    elif zmod_poly_degree(numer) == 0 or zmod_poly_degree(denom) == 0:
     564        if zmod_poly_leading(denom) != 1:
     565            a = mod_inverse_int(zmod_poly_leading(denom), p)
     566            zmod_poly_scalar_mul(numer, numer, a)
     567            zmod_poly_scalar_mul(denom, denom, a)
     568            return True
     569        return False
     570    cdef zmod_poly_t g
     571    cdef bint changed = False
     572    try:
     573        zmod_poly_init_precomp(g, p, numer.p_inv)
     574        zmod_poly_gcd(g, numer, denom)
     575        if zmod_poly_degree(g) != 0:
     576            # Divide knowing divisible by? Can we get these quotients as a byproduct of the gcd?
     577            zmod_poly_div(numer, numer, g)
     578            zmod_poly_div(denom, denom, g)
     579            changed = True
     580        if zmod_poly_leading(denom) != 1:
     581            a = mod_inverse_int(zmod_poly_leading(denom), p)
     582            zmod_poly_scalar_mul(numer, numer, a)
     583            zmod_poly_scalar_mul(denom, denom, a)
     584            changed = True
     585        return changed
     586    finally:
     587        zmod_poly_clear(g)
     588
     589cdef inline unsigned long zmod_poly_leading(zmod_poly_t poly):
     590    return zmod_poly_get_coeff_ui(poly, zmod_poly_degree(poly))
     591
     592cdef inline void zmod_poly_inc(zmod_poly_t poly, bint monic):
     593    cdef long n
     594    cdef long a
     595    cdef long p = poly.p
     596    for n from 0 <= n <= zmod_poly_degree(poly) + 1:
     597        a = zmod_poly_get_coeff_ui(poly, n) + 1
     598        if a == poly.p:
     599            zmod_poly_set_coeff_ui(poly, n, 0)
     600        else:
     601            zmod_poly_set_coeff_ui(poly, n, a)
     602            break
     603    if monic and a == 2 and n == zmod_poly_degree(poly):
     604        zmod_poly_set_coeff_ui(poly, n, 0)
     605        zmod_poly_set_coeff_ui(poly, n+1, 1)
     606
     607cdef void zmod_poly_pow(zmod_poly_t res, zmod_poly_t poly, unsigned long e):
     608    if zmod_poly_degree(poly) < 2:
     609        if zmod_poly_degree(poly) == -1:
     610            zmod_poly_zero(res)
     611            return
     612        elif zmod_poly_is_one(poly):
     613            zmod_poly_set(res, poly)
     614            return
     615        elif e > 1 and zmod_poly_degree(poly) == 1 and zmod_poly_get_coeff_ui(poly, 0) == 0 and zmod_poly_get_coeff_ui(poly, 1) == 1:
     616            zmod_poly_left_shift(res, poly, e-1)
     617            return
     618   
     619    # TODO: Could use the fact that (a+b)^p = a^p + b^p
     620    # Only seems to be a big gain for large p, large exponents...
     621    cdef zmod_poly_t pow2
     622   
     623    if e < 5:
     624        if e == 0:
     625            zmod_poly_zero(res)
     626            zmod_poly_set_coeff_ui(res, 0, 1)
     627        elif e == 1:
     628            zmod_poly_set(res, poly)
     629        elif e == 2:
     630            zmod_poly_sqr(res, poly)
     631        elif e == 3:
     632            zmod_poly_sqr(res, poly)
     633            zmod_poly_mul(res, res, poly)
     634        elif e == 4:
     635            zmod_poly_sqr(res, poly)
     636            zmod_poly_sqr(res, res)
     637    else:
     638        zmod_poly_init(pow2, poly.p)
     639        zmod_poly_zero(res)
     640        zmod_poly_set_coeff_ui(res, 0, 1)
     641        zmod_poly_set(pow2, poly)
     642        while True:
     643            if e & 1:
     644                zmod_poly_mul(res, res, pow2)
     645            e >>= 1
     646            if e == 0:
     647                break
     648            zmod_poly_sqr(pow2, pow2)
     649        zmod_poly_clear(pow2)
     650
     651
     652cdef long sqrt_mod_int(long a, long p) except -1:
     653    return GF(p)(a).sqrt()
     654
     655cdef bint zmod_poly_sqrt_check(zmod_poly_t poly):
     656    """
     657    Quick check to see if poly could possibly be a square.
     658    """
     659    return (zmod_poly_degree(poly) % 2 == 0
     660        and jacobi_int(zmod_poly_leading(poly), poly.p) == 1
     661        and jacobi_int(zmod_poly_get_coeff_ui(poly, 0), poly.p) != -1)
     662 
     663cdef bint zmod_poly_sqrt(zmod_poly_t res, zmod_poly_t poly):
     664    """
     665    Compute the square root of poly as res if res is a perfect square.
     666   
     667    Returns True on success, False on failure.
     668    """
     669    if not zmod_poly_sqrt_check(poly):
     670        return False
     671    return zmod_poly_sqrt0(res, poly)
     672
     673cdef bint zmod_poly_sqrt0(zmod_poly_t res, zmod_poly_t poly):
     674    """
     675    Compute the square root of poly as res if res is a perfect square, assuming that poly passes zmod_poly_sqrt_check.
     676   
     677    Returns True on success, False on failure.
     678    """
     679    if zmod_poly_degree(poly) == -1:
     680        zmod_poly_zero(res)
     681        return True
     682    if zmod_poly_degree(poly) == 0:
     683        zmod_poly_zero(res)
     684        zmod_poly_set_coeff_ui(res, 0, sqrt_mod_int(zmod_poly_get_coeff_ui(poly, 0), poly.p))
     685        return True
     686    cdef zmod_poly_t g
     687    cdef long p = poly.p
     688    cdef long n, leading
     689    cdef zmod_poly_factor_t factors
     690    try:
     691        zmod_poly_init(g, p)
     692        zmod_poly_derivative(res, poly)
     693        zmod_poly_gcd(g, res, poly)
     694        if 2*zmod_poly_degree(g) < zmod_poly_degree(poly):
     695            return False
     696           
     697        elif 2*zmod_poly_degree(g) > zmod_poly_degree(poly):
     698            try:
     699                zmod_poly_factor_init(factors)
     700                leading = zmod_poly_leading(poly)
     701                if leading == 1:
     702                    zmod_poly_factor_square_free(factors, poly)
     703                else:
     704                    zmod_poly_scalar_mul(res, poly, mod_inverse_int(leading, p))
     705                    zmod_poly_factor_square_free(factors, res)
     706                for n in range(factors.num_factors):
     707                    if factors.exponents[n] % 2 != 0:
     708                        return False
     709                zmod_poly_pow(res, factors.factors[0], factors.exponents[0]//2)
     710                for n in range(1, factors.num_factors):
     711                    zmod_poly_pow(factors.factors[n], factors.factors[n], factors.exponents[n]//2)
     712                    zmod_poly_mul(res, res, factors.factors[n])
     713                if leading != 1:
     714                    zmod_poly_scalar_mul(res, res, sqrt_mod_int(zmod_poly_leading(poly), p))
     715                return True
     716            finally:
     717                zmod_poly_factor_clear(factors)
     718           
     719        else: # deg(g) == deg(poly)/2
     720            zmod_poly_sqr(res, g)
     721            leading = zmod_poly_leading(poly) * mod_inverse_int(zmod_poly_leading(res), p)
     722            zmod_poly_scalar_mul(res, res, leading)
     723            if zmod_poly_equal(res, poly):
     724                zmod_poly_scalar_mul(res, g, sqrt_mod_int(leading, p))
     725                return True
     726            else:
     727                return False
     728    finally:
     729        zmod_poly_clear(g)