Data commit

Task/P-Adic-square-roots/Haskell/p-adic-square-roots.hs

@@ -0,0 +1,33 @@
+{-# LANGUAGE KindSignatures, DataKinds  #-}
+
+import Data.Ratio
+import Data.List (find)
+import GHC.TypeLits
+import Padic
+
+pSqrt :: KnownNat p => Rational -> Padic p
+pSqrt r = res
+  where
+    res = maybe Null mkUnit series
+    (a, b) = (numerator r, denominator r)
+    series = case modulo res of
+
+      2 | eqMod 4 a 3 -> Nothing
+        | not (eqMod 8 a 1) -> Nothing
+        | otherwise -> Just $ 1 : 0 : go 8 1
+        where
+          go pk x =
+            let q = ((b*x*x - a) `div` pk) `mod` 2
+            in q : go (2*pk) (x + q * (pk `div` 2))
+
+      p -> do
+        y <- find (\x -> eqMod p (b*x*x) a) [1..p-1]
+        df <- recipMod p (2*b*y)
+        let go pk x =
+              let f = (b*x*x - a) `div` pk
+                  d = (f * (p - df)) `mod` p
+              in x `div` (pk `div` p) : go (p*pk) (x + d*pk)
+        Just $ go p y
+
+eqMod :: Integral a => a -> a -> a -> Bool
+eqMod p a b = a `mod` p == b `mod` p