Modular square root function with Hensel's Stronger Lemma; algorithm for $\bmod 2^x$? Also $x^2 = a\pmod{p^n}$ when $\gcd(a,p) \neq $1

Question

I'm programming a solution to HackerRank's Project Euler 100, which is a much stronger variant than the original question, so I have to essentially make a quadratic solver, which involves a lot of quadratic modular equations. I'm trying to implement a full "modular square root" function which can take a number, a prime, and the power of that prime, i.e. find all x such that: $$ x^2 \equiv \text{num} \pmod{\text{prime}^{\text{power}}} $$ It currently works well enough for odd primes, but I have some edge cases that I'm still fleshing out:

I don't understand mod $2^n$ quite yet, but I feel like Hensel's stronger Lemma will help me? Likewise with num $= 0$. Also situations with num having prime as a factor. For example, $x^2 \equiv \text{num} \pmod{{2}^{{8}}}$ has many solutions for a variety of numbers, as does $x^2 \equiv \text{num} \pmod{\text{3}^{\text{power}}}$...

My guess is that Hensel's Stronger Lemma is what will let me lift, and even then, I'm not quite sure algorithmically what this implies. Hensel's Stronger Lemma says that you can lift lower roots to several higher powers, but not all of them, so does that mean I'd need to exhaustively search all qualifying lower roots of a prime to lift to all higher roots?

I am not a number theorist, just teaching myself coding, so a lot of this has been a huge crash course in number theory concepts!

Answer 1

After crawling through many PDFs and wikis, I finally grok'd how to generally apply Hensel's Stronger Lemma in an algorithmic fashion; I'm not sure why none of the other papers at least described a way to do it; it's all obfuscated by the jargon. As someone who doesn't know rings, p-adics, or anything else, here is what I found; I understand it may not be general enough for all polynomials, but it worked in my randomized testing suite for quadratic congruences of form $Ax^2 +Bx + C \equiv 0 \mod p^k$ with prime mods and powers.

1. Find the roots of $f(x) \equiv 0 \mod p^i$ using your choice of algorithm. For odd primes I used Tonelli-Shanks for 4k+1 and the (p+1)/4 rule for the others, then completed the square. For 2, I simply tested 0 and 1.
2. If $f'(x) \not\equiv 0 \mod p^i$, you can use simple Hensel's Lemma
3. If not, then you must use stronger Hensel's Lemma
4. For each root $r$, check if $f(r) \equiv 0 \mod p^{i+1}$. If so, $r$ is a root of $\mod p^{i+1}$, but also so is every congruent number to $r \mod p^{i+1}$ of form $r + k*p^i$ i.e. if 1 is a root mod $3^2$, then so is 4 and 7. This is how the number of roots can grow (or not) at every layer. For $3^3$, you would then check all the roots you just found; 1,4,7 could each sprout 3 new roots, etc.

Rinse, repeat step 4 until you've gotten as many powers as you want to go!