Computing a square root modulo a power of two

Computing a square root modulo a power of two turns out to be remarkably simple. E.g., for \( 2^{64} \) the simple function shown below written in Sage does the job (and obviously, this function can be used for any power).

def sqrt_pow_2(v):
    assert v % 8 == 1
    e = 1
    v_ = (v - 1) >> 2
    for i in range(1, 64):
        e = e - (e^2 - e - v_)
    return 2*e - 1

Z64 = IntegerModRing(2^64)
v = Z64(17)
sqrt_pow_2(v)^2 == v  # returns True

If you don't care about how the above function works (just that it does), then congratulations, you've now reached the end of this post.

Otherwise, keep reading. The rest of this post is my best effort at explaining how sqrt_pow_2 works.

Let \( v\in\mathbb{Z}_{2^k} \) for some \( k \). The goal is to find some \( e\in\mathbb{Z}_{2^k} \) such that \( e^2 \equiv v \pmod{2^k} \). The value \( e \) can be viewed as the square root of \( v \).

A generic way to compute a square root modulo some \( p^k \) is to find a root for the polynomial \( f(X) = X^2 - v \). This root, let's call it \( e \), will clearly satisfy \( e^2 \equiv v \pmod{p^k} \).

One way to find such a root is to (1) find a root \( e' \) modulo \( p^{i} \), then (2) use a Hensel lift to find a root modulo \( p^{i+1} \). This we continue to do until \( i = k \).

The problem is just that, in order to use a Hensel lift, we must have \( f'(X) \not \equiv 0 \pmod{p} \), which is not the case when \( p=2 \).

The trick is to consider a different polynomial \( g \), which also has a root from which a square root can be computed, but for which \( g'(X) \not \equiv 0 \pmod{2} \) is the case. One such polynomial is

\[ g(X) = X^2 - X - \frac{v - 1}{2}, \]

which has a root at \( e = (\sqrt{v} + 1)/2 \).

Consider the polynomial \( g \) and let \( e \) be a root modulo \( 2^{i} \) for some \( i \). I.e., \( g(e) \equiv 0 \pmod{2^{i}} \).

Then \( e' = e + t\cdot 2^{i} \) for some \( t\in\mathbb{Z}_2 \) is a root modulo \( 2^{i+1} \). We can calculate \( t \) by writing everything out:

and so

The above gives us a straightforward way to compute a square root of \( v\in\mathbb{Z}_{2^{k}} \):

1. Set \( i \leftarrow 1 \) and find an \( e \) such that \( g(e) \equiv 0 \pmod{2} \). This shouldn't be too hard to find.¹
2. Compute \( e' \leftarrow e + t\cdot 2^{i} \) where \( t = (-g(e) / 2^{i})/g'(e) \).
3. Set \( i \leftarrow i + 1 \) and \( e \leftarrow e' \).
4. Repeat step 2 and 3 until \( i = k \).
5. Return \( 2\cdot e - 1 \).

That is almost what the code at the top of this post does. What remains is to realize that, first

Second, \( g'(e) \) only has to be computed modulo 2, which honestly doesn't leave a lot of options. That is, we can simplify the above to

Which gives the final expression.²

The sqrt_pow_2 function gives us one square root. Can we find them all?

Yes, and it is in fact easy to do so: If \( e \) is a square root modulo \( 2^{k} \), then \( -e \), \( e + 2^{k-1} \) and \( -e + 2^{k-1} \) are square roots as well.