Clean up and include discrete logarithm functions into fmpz_mod

src/flint/flintlib/fmpz_mod.pxd

@@ -1,10 +1,11 @@
 from flint.flintlib.flint cimport ulong, slong
-from flint.flintlib.fmpz cimport fmpz_t, fmpz_preinvn_struct
+from flint.flintlib.fmpz cimport fmpz_t, fmpz_struct, fmpz_preinvn_struct
 from flint.flintlib.nmod cimport nmod_t
 
-# unimported types  {'fmpz_mod_discrete_log_pohlig_hellman_t'}
-
 cdef extern from "flint/fmpz_mod.h":
+    #
+    # fmpz_mod structs, a la Pohlig - Hellman
+    #
     ctypedef struct fmpz_mod_ctx_struct:
         fmpz_t n
         nmod_t mod
@@ -13,6 +14,36 @@ cdef extern from "flint/fmpz_mod.h":
         fmpz_preinvn_struct * ninv_huge
     ctypedef fmpz_mod_ctx_struct fmpz_mod_ctx_t[1]
 
+    #
+    # discrete logs structs, a la Pohlig - Hellman
+    #
+
+    ctypedef struct fmpz_mod_discrete_log_pohlig_hellman_table_entry_struct:
+        fmpz_t gammapow
+        ulong cm
+
+    ctypedef struct fmpz_mod_discrete_log_pohlig_hellman_entry_struct:
+        slong exp
+        ulong prime
+        fmpz_t gamma
+        fmpz_t gammainv
+        fmpz_t startingbeta
+        fmpz_t co
+        fmpz_t startinge
+        fmpz_t idem
+        ulong cbound
+        ulong dbound
+        fmpz_mod_discrete_log_pohlig_hellman_table_entry_struct * table # length cbound */
+
+    ctypedef struct fmpz_mod_discrete_log_pohlig_hellman_struct:
+        fmpz_mod_ctx_t fpctx
+        fmpz_t pm1         # p - 1 */
+        fmpz_t alpha       # p.r. of p */
+        fmpz_t alphainv
+        slong num_factors  # factors of p - 1
+        fmpz_mod_discrete_log_pohlig_hellman_entry_struct * entries
+    ctypedef fmpz_mod_discrete_log_pohlig_hellman_struct fmpz_mod_discrete_log_pohlig_hellman_t[1]
+
     # Parsed from here
     void fmpz_mod_ctx_init(fmpz_mod_ctx_t ctx, const fmpz_t n)
     void fmpz_mod_ctx_clear(fmpz_mod_ctx_t ctx)
@@ -37,9 +68,9 @@ cdef extern from "flint/fmpz_mod.h":
     int fmpz_mod_divides(fmpz_t a, const fmpz_t b, const fmpz_t c, const fmpz_mod_ctx_t ctx)
     void fmpz_mod_pow_ui(fmpz_t a, const fmpz_t b, ulong e, const fmpz_mod_ctx_t ctx)
     int fmpz_mod_pow_fmpz(fmpz_t a, const fmpz_t b, const fmpz_t e, const fmpz_mod_ctx_t ctx)
-    # void fmpz_mod_discrete_log_pohlig_hellman_init(fmpz_mod_discrete_log_pohlig_hellman_t L)
-    # void fmpz_mod_discrete_log_pohlig_hellman_clear(fmpz_mod_discrete_log_pohlig_hellman_t L)
-    # double fmpz_mod_discrete_log_pohlig_hellman_precompute_prime(fmpz_mod_discrete_log_pohlig_hellman_t L, const fmpz_t p)
-    # const fmpz_struct * fmpz_mod_discrete_log_pohlig_hellman_primitive_root(const fmpz_mod_discrete_log_pohlig_hellman_t L)
-    # void fmpz_mod_discrete_log_pohlig_hellman_run(fmpz_t x, const fmpz_mod_discrete_log_pohlig_hellman_t L, const fmpz_t y)
+    void fmpz_mod_discrete_log_pohlig_hellman_init(fmpz_mod_discrete_log_pohlig_hellman_t L)
+    void fmpz_mod_discrete_log_pohlig_hellman_clear(fmpz_mod_discrete_log_pohlig_hellman_t L)
+    double fmpz_mod_discrete_log_pohlig_hellman_precompute_prime(fmpz_mod_discrete_log_pohlig_hellman_t L, const fmpz_t p)
+    const fmpz_struct * fmpz_mod_discrete_log_pohlig_hellman_primitive_root(const fmpz_mod_discrete_log_pohlig_hellman_t L)
+    void fmpz_mod_discrete_log_pohlig_hellman_run(fmpz_t x, const fmpz_mod_discrete_log_pohlig_hellman_t L, const fmpz_t y)
     int fmpz_next_smooth_prime(fmpz_t a, const fmpz_t b)

src/flint/test/test.py

@@ -1770,6 +1770,50 @@ def test_fmpz_mod():
         assert fmpz(test_y) / F_test(test_x) == (test_y * pow(test_x, -1, test_mod)) % test_mod
         assert test_y / F_test(test_x) == (test_y * pow(test_x, -1, test_mod)) % test_mod
 
+def test_fmpz_mod_dlog():
+    from flint import fmpz, fmpz_mod_ctx
+
+    # Input modulus must be prime
+    F = fmpz_mod_ctx(4)
+    g, a = F(1), F(2)
+    assert raises(lambda: g.discrete_log(a, check=True), ValueError)
+
+    # Moduli must match
+    F1, F2 = fmpz_mod_ctx(2), fmpz_mod_ctx(3)
+    g = F1(2)
+    a = F2(4)
+    assert raises(lambda: g.discrete_log(a, check=True), ValueError)
+
+    # Need to use either fmpz_mod or something which can be case to
+    # fmpz
+    assert raises(lambda: g.discrete_log("A", check=True), TypeError)
+
+    F = fmpz_mod_ctx(163)
+    g = F(2)
+    a = g**123
+
+    assert 123 == g.discrete_log(a)
+
+    a_int = pow(2, 123, 163)
+    a_fmpz = fmpz(a_int)
+    assert 123 == g.discrete_log(a_int)
+    assert 123 == g.discrete_log(a_fmpz)
+
+    # Randomised testing with smooth large modulus
+    e2, e3 = 92, 79
+    p = 2**e2 * 3**e3 + 1
+    F = fmpz_mod_ctx(p)
+
+    import random
+    for _ in range(10):
+        g = F(random.randint(0,p))
+        for _ in range(10):
+            i = random.randint(0,p)
+            a = g**i
+            x = g.discrete_log(a)
+            assert g**x == a
+
+
 
 all_tests = [
     test_pyflint,
@@ -1790,4 +1834,5 @@ def test_fmpz_mod():
     test_nmod_mat,
     test_arb,
     test_fmpz_mod,
+    test_fmpz_mod_dlog
 ]

src/flint/types/fmpz_mod.pxd

@@ -1,14 +1,20 @@
 from flint.flint_base.flint_base cimport flint_scalar
 from flint.flintlib.fmpz cimport fmpz_t
-from flint.flintlib.fmpz_mod cimport fmpz_mod_ctx_t
+from flint.flintlib.fmpz_mod cimport (
+    fmpz_mod_ctx_t,
+    fmpz_mod_discrete_log_pohlig_hellman_t
+)
 
 
 cdef class fmpz_mod_ctx:
     cdef fmpz_mod_ctx_t val
-
+    cdef fmpz_mod_discrete_log_pohlig_hellman_t L
+    cdef bint _dlog_precomputed
+    cdef any_as_fmpz_mod(self, obj)
+    cdef _precompute_dlog_prime(self)
+
 cdef class fmpz_mod(flint_scalar):
     cdef fmpz_mod_ctx ctx
     cdef fmpz_t val
-
-    cdef any_as_fmpz_mod(self, obj)
-
+    cdef bint base_dlog_precomputed
+    cdef fmpz_t x_g