add transformation interface to bkz_reduction

fplll · Feb 4, 2024 · 821e975 · 821e975
1 parent 50c3040
commit 821e975
Showing 1 changed file with 16 additions and 5 deletions.
diff --git a/src/fpylll/fplll/bkz.pyx b/src/fpylll/fplll/bkz.pyx
@@ -1106,12 +1106,13 @@ cdef class BKZReduction:
 
 
 
-def bkz_reduction(IntegerMatrix B, BKZParam o, float_type=None, int precision=0):
+def bkz_reduction(IntegerMatrix B, BKZParam o, U=None, float_type=None, int precision=0):
     """
     Run BKZ reduction.
 
     :param IntegerMatrix B: Integer matrix, modified in place.
     :param BKZParam o: BKZ parameters
+    :param U: Transformation matrix or ``None``
     :param float_type: either ``None``: for automatic choice or an entry of `fpylll.config.float_types`
     :param precision: bit precision to use if ``float_type`` is ``'mpfr'``
 
@@ -1125,10 +1126,20 @@ def bkz_reduction(IntegerMatrix B, BKZParam o, float_type=None, int precision=0)
     if B._type != ZT_MPZ:
         raise NotImplementedError("C++ BKZ is not implemented over longs, try the Python version.")
 
-    with nogil:
-        sig_on()
-        r = bkz_reduction_c(B._core.mpz, NULL, o.o[0], float_type_, precision)
-        sig_off()
+    if U is not None and isinstance(U, IntegerMatrix):
+        with nogil:
+            sig_on()
+            r = bkz_reduction_c((<IntegerMatrix>B)._core.mpz,
+                                (<IntegerMatrix>U)._core.mpz,
+                                o.o[0], float_type_, precision)
+            sig_off()
+    else:
+        with nogil:
+            sig_on()
+            r = bkz_reduction_c((<IntegerMatrix>B)._core.mpz,
+                                NULL, o.o[0], float_type_, precision)
+            sig_off()
+
 
     if r and r not in (RED_BKZ_LOOPS_LIMIT, RED_BKZ_TIME_LIMIT):
         raise ReductionError( str(get_red_status_str(r)) )