Hensel lifting in Python
Lifts a factorisation modulo p to one modulo p^k with quadratic convergence in the exponent.
- File
hensel_lifting.py, 203 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/hensel-lifting/. The authors are listed in prov/ZKSCMK.jsonl.
"""Quadratic Hensel lifting of a factorization f = g h modulo p to modulo p^k.
A polynomial a_0 + a_1 x + ... + a_n x^n is the list [a_0, ..., a_n] of
Python ints, lowest degree first; the zero polynomial is []. Residues modulo
m are kept in the range 0 .. m - 1.
The input is f in Z[x], a prime p, factors g, h with f = g h mod p, h monic,
g mod p and h mod p coprime, and the leading coefficient of f prime to p.
Each Hensel step goes from modulus m to modulus m^2, so ceil(log2 k) steps
reach p^k. The step itself (hensel_step) is valid for any integer modulus
m >= 2, prime or not. All loops are iterative.
"""
from __future__ import annotations
from typing import List, Sequence, Tuple
Poly = List[int]
def trim(a: Sequence[int]) -> Poly:
"""Return a copy of a without trailing zero coefficients."""
out = list(a)
while out and out[-1] == 0:
out.pop()
return out
def reduce_mod(a: Sequence[int], m: int) -> Poly:
"""Reduce every coefficient into 0 .. m - 1 and strip leading zeros."""
return trim([c % m for c in a])
def symmetric(a: Sequence[int], m: int) -> Poly:
"""Map every coefficient into the symmetric range -m/2 < c <= m/2."""
out = []
for c in a:
c %= m
out.append(c - m if 2 * c > m else c)
return trim(out)
def mul(a: Sequence[int], b: Sequence[int]) -> Poly:
"""Schoolbook product over Z."""
if not a or not b:
return []
c = [0] * (len(a) + len(b) - 1)
for i, x in enumerate(a):
if x:
for j, y in enumerate(b):
c[i + j] += x * y
return trim(c)
def add(a: Sequence[int], b: Sequence[int]) -> Poly:
"""Coefficientwise sum."""
n = max(len(a), len(b))
return trim([(a[i] if i < len(a) else 0) + (b[i] if i < len(b) else 0) for i in range(n)])
def sub(a: Sequence[int], b: Sequence[int]) -> Poly:
"""Coefficientwise difference a - b."""
n = max(len(a), len(b))
return trim([(a[i] if i < len(a) else 0) - (b[i] if i < len(b) else 0) for i in range(n)])
def divmod_monic(a: Sequence[int], h: Sequence[int], m: int) -> Tuple[Poly, Poly]:
"""Divide a by the monic h with remainder, coefficients modulo m."""
h = reduce_mod(h, m)
if not h or h[-1] != 1:
raise ValueError("the divisor must be monic modulo m")
r = reduce_mod(a, m)
e = len(h) - 1
if len(r) <= e:
return [], r
q = [0] * (len(r) - e)
for i in range(len(r) - 1, e - 1, -1):
c = r[i] % m
if c:
q[i - e] = c
for j in range(e + 1):
r[i - e + j] = (r[i - e + j] - c * h[j]) % m
return trim(q), reduce_mod(r[:e], m)
def cofactors_mod_p(g: Sequence[int], h: Sequence[int], p: int) -> Tuple[Poly, Poly]:
"""Return s, t with s g + t h = 1 mod p, deg s < deg h and deg t < deg g.
Uses the extended Euclidean algorithm in (Z/pZ)[x]; p must be prime and
g, h coprime modulo p.
"""
r0, r1 = reduce_mod(g, p), reduce_mod(h, p)
s0, s1 = [1], []
t0, t1 = [], [1]
while r1:
inv = pow(r1[-1], -1, p)
q = [0] * max(0, len(r0) - len(r1) + 1)
r = list(r0)
for i in range(len(r) - 1, len(r1) - 2, -1):
c = r[i] * inv % p
if c:
q[i - len(r1) + 1] = c
for j in range(len(r1)):
r[i - len(r1) + 1 + j] = (r[i - len(r1) + 1 + j] - c * r1[j]) % p
q = trim(q)
r0, r1 = r1, reduce_mod(r, p)
s0, s1 = s1, reduce_mod(sub(s0, mul(q, s1)), p)
t0, t1 = t1, reduce_mod(sub(t0, mul(q, t1)), p)
if len(r0) != 1:
raise ValueError("g and h are not coprime modulo p")
inv = pow(r0[0], -1, p)
return reduce_mod([c * inv for c in s0], p), reduce_mod([c * inv for c in t0], p)
def hensel_step(f, g, h, s, t, m):
"""One Hensel step: from f = g h, s g + t h = 1 mod m to the same mod m^2."""
m2 = m * m
e = reduce_mod(sub(f, mul(g, h)), m2) # @line L4
q, r = divmod_monic(mul(s, e), h, m2) # @line L5
g_new = reduce_mod(add(g, add(mul(t, e), mul(q, g))), m2) # @line L6
h_new = reduce_mod(add(h, r), m2) # @line L6
b = reduce_mod(sub(add(mul(s, g_new), mul(t, h_new)), [1]), m2) # @line L7
c, d = divmod_monic(mul(s, b), h_new, m2) # @line L8
s_new = reduce_mod(sub(s, d), m2) # @line L9
t_new = reduce_mod(sub(t, add(mul(t, b), mul(c, g_new))), m2) # @line L9
return g_new, h_new, s_new, t_new
def hensel_lift(f: Sequence[int], g: Sequence[int], h: Sequence[int], p: int, k: int,
steps: list | None = None) -> Tuple[Poly, Poly]:
"""Lift f = g h mod p, with h monic and gcd(g, h) = 1 mod p, to f = g h mod p^k.
Returns (g_k, h_k) with coefficients in 0 .. p^k - 1, h_k monic,
g_k = g and h_k = h mod p. If steps is a list, the moduli and factors of
every step are appended to it.
"""
if k < 1:
raise ValueError("k must be at least 1")
f, g, h = trim(f), reduce_mod(g, p), reduce_mod(h, p)
if not h or h[-1] != 1:
raise ValueError("h must be monic modulo p")
if reduce_mod(sub(f, mul(g, h)), p):
raise ValueError("f is not g h modulo p")
s, t = cofactors_mod_p(g, h, p) # @line L1
m = p # @line L2
pk = p ** k
while m < pk: # @line L3
g, h, s, t = hensel_step(f, g, h, s, t, m)
m = m * m # @line L10
if steps is not None:
steps.append((m, list(g), list(h)))
return reduce_mod(g, pk), reduce_mod(h, pk) # @line L11
def lift_factors(f: Sequence[int], factors: Sequence[Sequence[int]], p: int, k: int) -> List[Poly]:
"""Lift f = lc(f) f_1 ... f_r mod p, all f_i monic and pairwise coprime mod p.
Splits off one factor at a time (a linear factor tree). Returns monic
f_1*, ..., f_r* with f = lc(f) f_1* ... f_r* mod p^k.
"""
f = trim(f)
pk = p ** k
rest = reduce_mod(f, pk)
lifted: List[Poly] = []
for i, fi in enumerate(factors[:-1]):
others = [f[-1] % p]
for fj in factors[i + 1:]:
others = reduce_mod(mul(others, fj), p)
g_k, h_k = hensel_lift(rest, others, fi, p, k)
lifted.append(h_k)
rest = g_k
if factors:
inv = pow(rest[-1], -1, pk)
lifted.append(reduce_mod([c * inv for c in rest], pk))
return lifted
def lift_root(f: Sequence[int], a: int, p: int, k: int) -> int:
"""Lift a simple root a of f modulo p to the unique root modulo p^k congruent to a.
This is the linear factor case h = x - a, done by the Newton step
b <- b - f(b) / f'(b) modulo p^(2j).
"""
f = trim(f)
df = trim([i * c for i, c in enumerate(f)][1:])
ev = lambda poly, x, m: _horner(poly, x, m)
if ev(f, a, p) != 0 or ev(df, a, p) == 0:
raise ValueError("a must be a simple root modulo p")
b, m, pk = a % p, p, p ** k
while m < pk:
m = m * m
b = (b - ev(f, b, m) * pow(ev(df, b, m), -1, m)) % m
return b % pk
def _horner(a: Sequence[int], x: int, m: int) -> int:
v = 0
for c in reversed(a):
v = (v * x + c) % m
return v
test_hensel_lifting.py 283 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/hensel-lifting/. The authors are listed in prov/ZKSCMK.jsonl.
import itertools
import random
import pytest
from hensel_lifting import (
cofactors_mod_p,
hensel_lift,
hensel_step,
lift_factors,
lift_root,
symmetric,
)
PRIMES = [2, 3, 7, 998244353]
# Independent oracle: plain integer polynomial arithmetic written separately.
def ref_trim(a):
a = list(a)
while a and a[-1] == 0:
a.pop()
return a
def ref_mul(a, b):
if not a or not b:
return []
out = [0] * (len(a) + len(b) - 1)
for i in range(len(a)):
for j in range(len(b)):
out[i + j] += a[i] * b[j]
return ref_trim(out)
def ref_mod(a, m):
return ref_trim([c % m for c in a])
def ref_rem_monic(a, h, m):
"""Remainder of a by the monic h, modulo m, by repeated subtraction of shifts."""
r = ref_mod(a, m)
e = len(h) - 1
while len(r) > e:
c = r[-1]
shift = len(r) - 1 - e
r = ref_mod([r[i] - (c * h[i - shift] if 0 <= i - shift <= e else 0) for i in range(len(r))], m)
return r
def congruent(a, b, m):
n = max(len(a), len(b))
a = list(a) + [0] * (n - len(a))
b = list(b) + [0] * (n - len(b))
return all((x - y) % m == 0 for x, y in zip(a, b))
def gcd_is_one_mod_p(a, b, p):
a, b = ref_mod(a, p), ref_mod(b, p)
while b:
inv = pow(b[-1], -1, p)
while len(a) >= len(b):
c = a[-1] * inv % p
s = len(a) - len(b)
a = ref_mod([a[i] - (c * b[i - s] if i >= s else 0) for i in range(len(a))], p)
a, b = b, a
return len(a) == 1
def random_monic(rng, deg, bound):
return [rng.randint(-bound, bound) for _ in range(deg)] + [1]
def check_lift(f, g0, h0, p, k):
g, h = hensel_lift(f, g0, h0, p, k)
pk = p ** k
assert congruent(f, ref_mul(g, h), pk)
assert congruent(g, g0, p) and congruent(h, h0, p)
assert h and h[-1] == 1 and len(h) == len(ref_mod(h0, p))
assert all(0 <= c < pk for c in g + h)
return g, h
def test_worked_example():
f = [-1, 0, 0, 0, 1]
g0, h0 = [-2, -1, 2, 1], [-2, 1]
steps = []
g, h = hensel_lift(f, g0, h0, 5, 4, steps)
assert cofactors_mod_p(g0, h0, 5) == ([3], [4, 3, 2])
assert [m for m, _, _ in steps] == [25, 625]
assert symmetric(steps[0][2], 25) == [-7, 1]
assert symmetric(steps[0][1], 25) == [-7, -1, 7, 1]
assert symmetric(h, 625) == [-182, 1]
assert symmetric(g, 625) == [-182, -1, 182, 1]
assert (182 * 182 + 1) % 625 == 0
@pytest.mark.parametrize("p", PRIMES)
def test_recovers_true_factors(p):
rng = random.Random(1000 + p % 1000)
done = 0
while done < 25:
a = random_monic(rng, rng.randint(0, 4), 9)
b = random_monic(rng, rng.randint(1, 4), 9)
lc = rng.choice([1, -1, 2, 3, 5, 7, 11])
if lc % p == 0 or not gcd_is_one_mod_p(a, b, p):
continue
f = ref_mul([lc], ref_mul(a, b))
bound = 2 * max(abs(c) for c in b) + 1
k = 1
while p ** k <= bound:
k += 1
k += rng.randint(0, 2)
g, h = check_lift(f, ref_mul([lc], a), b, p, k)
assert symmetric(h, p ** k) == b
assert congruent(g, ref_mul([lc], a), p ** k)
done += 1
@pytest.mark.parametrize("p", [2, 3, 7, 998244353])
def test_random_factorizations_mod_p(p):
"""f is only known to factor modulo p; check the lifted congruence."""
rng = random.Random(7 * (p % 10007))
done = 0
while done < 40:
g0 = [rng.randrange(p) for _ in range(rng.randint(1, 5))]
h0 = [rng.randrange(p) for _ in range(rng.randint(0, 4))] + [1]
g0 = ref_mod(g0, p)
if not g0 or not gcd_is_one_mod_p(g0, h0, p):
continue
noise = [p * rng.randint(-50, 50) for _ in range(len(g0) + len(h0) - 2)]
f = ref_trim([x + (noise[i] if i < len(noise) else 0) for i, x in enumerate(ref_mul(g0, h0))])
if not f or f[-1] % p == 0:
continue
k = rng.randint(1, 6) if p < 100 else rng.randint(1, 3)
check_lift(f, g0, h0, p, k)
done += 1
@pytest.mark.parametrize("p,k", [(2, 3), (3, 2), (3, 3), (7, 2)])
def test_uniqueness_by_brute_force(p, k):
rng = random.Random(p * 100 + k)
pk = p ** k
done = 0
while done < 8:
a = random_monic(rng, rng.randint(1, 3), 5)
b = random_monic(rng, rng.randint(1, 2), 5)
if not gcd_is_one_mod_p(a, b, p):
continue
f = ref_mul(a, b)
g, h = hensel_lift(f, a, b, p, k)
e = len(b) - 1
found = []
base = [c % p for c in b[:e]]
for lifts in itertools.product(range(p ** (k - 1)), repeat=e):
cand = [base[i] + p * lifts[i] for i in range(e)] + [1]
if not ref_rem_monic(f, cand, pk):
found.append(ref_mod(cand, pk))
assert found == [h]
done += 1
def test_composite_modulus_step():
"""The Hensel step is valid for any modulus, here m = 6 and then m = 36."""
rng = random.Random(6)
done = 0
while done < 30:
a = random_monic(rng, rng.randint(1, 3), 9)
b = random_monic(rng, rng.randint(1, 3), 9)
if not (gcd_is_one_mod_p(a, b, 2) and gcd_is_one_mod_p(a, b, 3)):
continue
s2, t2 = cofactors_mod_p(a, b, 2)
s3, t3 = cofactors_mod_p(a, b, 3)
def crt(x, y):
n = max(len(x), len(y))
x = x + [0] * (n - len(x))
y = y + [0] * (n - len(y))
return [(3 * u + 4 * v) % 6 for u, v in zip(x, y)]
s, t = crt(s2, s3), crt(t2, t3)
f = [c + 6 * rng.randint(-20, 20) for c in ref_mul(a, b)]
f[-1] = 1
g, h, m = ref_mod(a, 6), ref_mod(b, 6), 6
for _ in range(2):
g, h, s, t = hensel_step(f, g, h, s, t, m)
m = m * m
assert congruent(f, ref_mul(g, h), m)
assert congruent(ref_trim([x + y for x, y in itertools.zip_longest(ref_mul(s, g), ref_mul(t, h), fillvalue=0)]), [1], m)
assert h[-1] == 1 and len(h) == len(b)
assert len(s) < len(h) and len(t) < len(g)
done += 1
def test_edge_cases():
# leading zero coefficients in the input are stripped
g, h = hensel_lift([-1, 0, 0, 0, 1, 0, 0], [-2, -1, 2, 1, 0], [-2, 1, 0], 5, 2)
assert symmetric(h, 25) == [-7, 1]
# constant cofactor g: f = 3 (x^2 + x + 1), p = 2
g, h = hensel_lift([3, 3, 3], [1], [1, 1, 1], 2, 5)
assert g == [3] and h == [1, 1, 1]
# monic h of degree 0: everything goes to g
g, h = hensel_lift([4, 1, 9], [4, 1, 2], [1], 7, 3)
assert h == [1] and g == [4, 1, 9]
# k = 1 returns the input reduced modulo p
assert hensel_lift([-1, 0, 1], [-1, 1], [1, 1], 3, 1) == ([2, 1], [1, 1])
# a constant f with h = 1
assert hensel_lift([5], [5], [1], 3, 2) == ([5], [1])
def test_rejects_bad_input():
with pytest.raises(ValueError):
hensel_lift([0, 0, 1], [0, 1], [0, 1], 5, 3) # x and x are not coprime
with pytest.raises(ValueError):
hensel_lift([1, 0, 2], [1, 0, 2], [0, 2], 5, 3) # h not monic
with pytest.raises(ValueError):
hensel_lift([1, 0, 1], [1, 1], [1, 1], 5, 3) # f is not g h mod p
with pytest.raises(ValueError):
lift_root([-7, 0, 1], 0, 7, 2) # not a simple root
@pytest.mark.parametrize("p", [3, 7, 998244353])
def test_lift_root_square_roots(p):
rng = random.Random(p % 991)
k = 6 if p < 100 else 2
pk = p ** k
for _ in range(30):
c = rng.randrange(1, p)
if pow(c, (p - 1) // 2, p) != 1:
continue
r = next(x for x in range(p) if x * x % p == c) if p < 100 else _sqrt_big(c, p)
b = lift_root([-c, 0, 1], r, p, k)
assert (b * b - c) % pk == 0 and b % p == r
def _sqrt_big(c, p):
# p = 998244353 = 119 * 2^23 + 1; Tonelli-Shanks
q, s = p - 1, 0
while q % 2 == 0:
q, s = q // 2, s + 1
z = 3
m, cc, t, r = s, pow(z, q, p), pow(c, q, p), pow(c, (q + 1) // 2, p)
while t != 1:
i, tt = 0, t
while tt != 1:
tt, i = tt * tt % p, i + 1
b = pow(cc, 1 << (m - i - 1), p)
m, cc, t, r = i, b * b % p, t * b * b % p, r * b % p
return r
def test_lift_root_matches_factor_lift():
assert lift_root([-2, 0, 1], 3, 7, 4) == 2166
g, h = hensel_lift([-2, 0, 1], [3, 1], [-3, 1], 7, 4)
assert h == [(-2166) % 7 ** 4, 1]
@pytest.mark.parametrize("p", [2, 3, 7])
def test_lift_factors(p):
rng = random.Random(55 + p)
done = 0
while done < 12:
r = rng.randint(2, 4)
facs = [random_monic(rng, rng.randint(1, 2), 6) for _ in range(r)]
if any(not gcd_is_one_mod_p(facs[i], facs[j], p) for i in range(r) for j in range(i + 1, r)):
continue
lc = rng.choice([1, -1, 5, 11])
if lc % p == 0:
continue
f = [lc]
for x in facs:
f = ref_mul(f, x)
k = 8
lifted = lift_factors(f, [ref_mod(x, p) for x in facs], p, k)
prod = [lc]
for x in lifted:
prod = ref_mul(prod, x)
assert congruent(prod, f, p ** k)
assert [symmetric(x, p ** k) for x in lifted] == facs
done += 1
What the next level needs
- the tests did not run in this build