feat: add sqrtmod_pk (#2155)

- only for units for now

Author: thofma

src/flint/fmpz.jl

@@ -1078,6 +1078,41 @@ function sqrtmod(x::ZZRingElem, m::ZZRingElem)
   return z
 end
 
+function sqrtmod_pk(a::IntegerUnion, p::IntegerUnion, k::Int)
+  @req a != 0 && valuation(a, p) == 0 "Element must be a unit modulo p^k"
+  return _sqrtmod_pk_big(ZZ(a), ZZ(p), k)
+end
+
+function _sqrtmod_pk_big(a::ZZRingElem, p::ZZRingElem, k::Int)
+  # must be unit modulo p
+  r = ZZ()
+  fl = @ccall libflint._padic_sqrt(r::Ref{ZZRingElem}, a::Ref{ZZRingElem}, p::Ref{ZZRingElem}, k::Int)::Bool
+  if !fl
+    error("No square root exists")
+  end
+  return r
+end
+
+function _sqrtmod_pk_small(a::Int, p::Int, k::Int)
+  # this is too expensive to check
+  #@req k < @ccall Nemo.libflint.n_flog((UInt(1) << 63)::UInt, p::Int)::Int
+  if a < 0
+    a = mod(a, p)
+  end
+  # must be unit modulo p
+  res = Ref(Ptr{UInt}())
+  n = @ccall libflint.n_sqrtmod_primepow(res::Ref{Ptr{UInt}}, a::Int, p::Int, k::Int)::Int
+
+  if n == 0
+    error("No square root exists")
+  end
+
+  ptr = res[]
+  b = unsafe_load(ptr, 1)
+  @ccall libflint.flint_free(ptr::Ptr{UInt})::Nothing
+  return b % Int
+end
+
 function _normalize_crt(r::ZZRingElem, m::ZZRingElem, signed)
   s = sign(Int, m)
   if s > 0

test/flint/fmpz-test.jl

@@ -958,6 +958,21 @@ end
   @test mod_sym(ZZ(4), ZZ(-5)) == ZZ(-1)
   @test mod_sym(ZZ(-4), ZZ(5)) == ZZ(1)
   @test mod_sym(ZZ(-4), ZZ(-5)) == ZZ(1)
+
+  for (p, k) in [(2, 4), (2, 5), (5, 1), (5, 5)]
+    b = rand(1:p^k - 1) % p^k
+    while gcd(b, p) != 1
+      b = rand(1:p^k - 1) % p^k
+    end
+    a = powermod(b, 2, p^k)
+    c = Nemo.sqrtmod_pk(a, p, k)
+    @test c^2 % p^k == a
+    c = Nemo._sqrtmod_pk_small(a, p, k)
+    @test c^2 % p^k == a
+  end
+
+  @test_throws ArgumentError Nemo.sqrtmod_pk(2, 2, 5)
+  @test_throws ErrorException Nemo.sqrtmod_pk(3, 5, 10)
 end
 
 @testset "ZZRingElem.crt" begin