{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

module Algebra.Arithmetic
  ( repeatedSquare,
    modPow,
    fermatTest,
    isPseudoPrime,
  )
where

import AlgebraicPrelude hiding (div, mod)
import Control.Lens ((&), (+~), _1)
import Control.Monad.Random (MonadRandom, uniform)
import Numeric.Domain.Euclidean ()
import Prelude (div, mod)
import qualified Prelude as P

data PrimeResult = Composite | ProbablyPrime | Prime
  deriving (ReadPrec [PrimeResult]
ReadPrec PrimeResult
Int -> ReadS PrimeResult
ReadS [PrimeResult]
(Int -> ReadS PrimeResult)
-> ReadS [PrimeResult]
-> ReadPrec PrimeResult
-> ReadPrec [PrimeResult]
-> Read PrimeResult
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [PrimeResult]
$creadListPrec :: ReadPrec [PrimeResult]
readPrec :: ReadPrec PrimeResult
$creadPrec :: ReadPrec PrimeResult
readList :: ReadS [PrimeResult]
$creadList :: ReadS [PrimeResult]
readsPrec :: Int -> ReadS PrimeResult
$creadsPrec :: Int -> ReadS PrimeResult
Read, Int -> PrimeResult -> ShowS
[PrimeResult] -> ShowS
PrimeResult -> String
(Int -> PrimeResult -> ShowS)
-> (PrimeResult -> String)
-> ([PrimeResult] -> ShowS)
-> Show PrimeResult
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PrimeResult] -> ShowS
$cshowList :: [PrimeResult] -> ShowS
show :: PrimeResult -> String
$cshow :: PrimeResult -> String
showsPrec :: Int -> PrimeResult -> ShowS
$cshowsPrec :: Int -> PrimeResult -> ShowS
Show, PrimeResult -> PrimeResult -> Bool
(PrimeResult -> PrimeResult -> Bool)
-> (PrimeResult -> PrimeResult -> Bool) -> Eq PrimeResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PrimeResult -> PrimeResult -> Bool
$c/= :: PrimeResult -> PrimeResult -> Bool
== :: PrimeResult -> PrimeResult -> Bool
$c== :: PrimeResult -> PrimeResult -> Bool
Eq, Eq PrimeResult
Eq PrimeResult
-> (PrimeResult -> PrimeResult -> Ordering)
-> (PrimeResult -> PrimeResult -> Bool)
-> (PrimeResult -> PrimeResult -> Bool)
-> (PrimeResult -> PrimeResult -> Bool)
-> (PrimeResult -> PrimeResult -> Bool)
-> (PrimeResult -> PrimeResult -> PrimeResult)
-> (PrimeResult -> PrimeResult -> PrimeResult)
-> Ord PrimeResult
PrimeResult -> PrimeResult -> Bool
PrimeResult -> PrimeResult -> Ordering
PrimeResult -> PrimeResult -> PrimeResult
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: PrimeResult -> PrimeResult -> PrimeResult
$cmin :: PrimeResult -> PrimeResult -> PrimeResult
max :: PrimeResult -> PrimeResult -> PrimeResult
$cmax :: PrimeResult -> PrimeResult -> PrimeResult
>= :: PrimeResult -> PrimeResult -> Bool
$c>= :: PrimeResult -> PrimeResult -> Bool
> :: PrimeResult -> PrimeResult -> Bool
$c> :: PrimeResult -> PrimeResult -> Bool
<= :: PrimeResult -> PrimeResult -> Bool
$c<= :: PrimeResult -> PrimeResult -> Bool
< :: PrimeResult -> PrimeResult -> Bool
$c< :: PrimeResult -> PrimeResult -> Bool
compare :: PrimeResult -> PrimeResult -> Ordering
$ccompare :: PrimeResult -> PrimeResult -> Ordering
$cp1Ord :: Eq PrimeResult
Ord)

-- | Calculates @n@-th power efficiently, using repeated square method.
repeatedSquare :: Multiplicative r => r -> Natural -> r
repeatedSquare :: r -> Natural -> r
repeatedSquare r
a Natural
n =
  let bits :: [Natural]
bits = [Natural] -> [Natural]
forall a. [a] -> [a]
tail ([Natural] -> [Natural]) -> [Natural] -> [Natural]
forall a b. (a -> b) -> a -> b
$ Natural -> [Natural]
binRep Natural
n
   in (r -> Natural -> r) -> r -> [Natural] -> r
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\r
b Natural
nk -> if Natural
nk Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Natural
1 then r
b r -> r -> r
forall r. Multiplicative r => r -> r -> r
* r
b r -> r -> r
forall r. Multiplicative r => r -> r -> r
* r
a else r
b r -> r -> r
forall r. Multiplicative r => r -> r -> r
* r
b) r
a [Natural]
bits

binRep :: Natural -> [Natural]
binRep :: Natural -> [Natural]
binRep = (Natural -> [Natural] -> [Natural])
-> [Natural] -> Natural -> [Natural]
forall a b c. (a -> b -> c) -> b -> a -> c
flip Natural -> [Natural] -> [Natural]
forall a. Integral a => a -> [a] -> [a]
go []
  where
    go :: a -> [a] -> [a]
go a
0 = [a] -> [a]
forall a. a -> a
id
    go a
k = a -> [a] -> [a]
go (a
k a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2) ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a
k a -> a -> a
forall a. Integral a => a -> a -> a
`mod` a
2) a -> [a] -> [a]
forall a. a -> [a] -> [a]
:)

-- | Fermat-test for pseudo-primeness.
fermatTest :: MonadRandom m => Integer -> m PrimeResult
fermatTest :: Integer -> m PrimeResult
fermatTest Integer
2 = PrimeResult -> m PrimeResult
forall (m :: * -> *) a. Monad m => a -> m a
return PrimeResult
Prime
fermatTest Integer
n = do
  Integer
a <- [Integer] -> m Integer
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadRandom m) =>
t a -> m a
uniform [Integer
2 .. Integer
n Integer -> Integer -> Integer
forall r. Group r => r -> r -> r
- Integer
2]
  let b :: Integer
b = Integer -> Integer -> Natural -> Integer
forall a r. (Integral a, Euclidean r) => r -> r -> a -> r
modPow Integer
n (Integer -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
a) (Integer -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Natural) -> Integer -> Natural
forall a b. (a -> b) -> a -> b
$ Integer
n Integer -> Integer -> Integer
forall r. Group r => r -> r -> r
- Integer
1 :: Natural)
  if Integer
b Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
1
    then PrimeResult -> m PrimeResult
forall (m :: * -> *) a. Monad m => a -> m a
return PrimeResult
Composite
    else PrimeResult -> m PrimeResult
forall (m :: * -> *) a. Monad m => a -> m a
return PrimeResult
ProbablyPrime

-- | @'modPow' x m p@ efficiently calculates @x ^ p `'mod'` m@.
modPow :: (P.Integral a, Euclidean r) => r -> r -> a -> r
modPow :: r -> r -> a -> r
modPow r
i r
p = r -> r -> a -> r
go r
i r
forall r. Unital r => r
one
  where
    go :: r -> r -> a -> r
go r
_ r
acc a
0 = r
acc
    go r
b r
acc a
e = r -> r -> a -> r
go ((r
b r -> r -> r
forall r. Multiplicative r => r -> r -> r
* r
b) r -> r -> r
forall d. Euclidean d => d -> d -> d
`rem` r
p) (if a
e a -> a -> a
forall a. Integral a => a -> a -> a
`mod` a
2 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1 then (r
acc r -> r -> r
forall r. Multiplicative r => r -> r -> r
* r
b) r -> r -> r
forall d. Euclidean d => d -> d -> d
`rem` r
p else r
acc) (a
e a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2)

splitFactor :: Euclidean r => r -> r -> (Int, r)
splitFactor :: r -> r -> (Int, r)
splitFactor r
d r
n =
  let (r
q, r
r) = r
n r -> r -> (r, r)
forall d. Euclidean d => d -> d -> (d, d)
`divide` r
d
   in if r -> Bool
forall r. DecidableZero r => r -> Bool
isZero r
q
        then (Int
0, r
n)
        else r -> r -> (Int, r)
forall r. Euclidean r => r -> r -> (Int, r)
splitFactor r
d r
r (Int, r) -> ((Int, r) -> (Int, r)) -> (Int, r)
forall a b. a -> (a -> b) -> b
& (Int -> Identity Int) -> (Int, r) -> Identity (Int, r)
forall s t a b. Field1 s t a b => Lens s t a b
_1 ((Int -> Identity Int) -> (Int, r) -> Identity (Int, r))
-> Int -> (Int, r) -> (Int, r)
forall a s t. Num a => ASetter s t a a -> a -> s -> t
+~ Int
1

{- | @'isPseudoPrime' n@ tests if the given integer @n@ is pseudo prime.
   It returns @'Left' p@ if @p < n@ divides @n@,
   @'Right' 'True'@ if @n@ is pseudo-prime,
   @'Right' 'False'@ if it is not pseudo-prime but no clue can be found.
-}
isPseudoPrime ::
  MonadRandom m =>
  Integer ->
  m (Either Integer Bool)
isPseudoPrime :: Integer -> m (Either Integer Bool)
isPseudoPrime Integer
2 = Either Integer Bool -> m (Either Integer Bool)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Integer Bool -> m (Either Integer Bool))
-> Either Integer Bool -> m (Either Integer Bool)
forall a b. (a -> b) -> a -> b
$ Bool -> Either Integer Bool
forall a b. b -> Either a b
Right Bool
True
isPseudoPrime Integer
3 = Either Integer Bool -> m (Either Integer Bool)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Integer Bool -> m (Either Integer Bool))
-> Either Integer Bool -> m (Either Integer Bool)
forall a b. (a -> b) -> a -> b
$ Bool -> Either Integer Bool
forall a b. b -> Either a b
Right Bool
True
isPseudoPrime Integer
n = do
  Integer
a <- [Integer] -> m Integer
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadRandom m) =>
t a -> m a
uniform [Integer
2 .. Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
P.- Integer
2]
  let d :: Integer
d = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
P.gcd Integer
a Integer
n
  Either Integer Bool -> m (Either Integer Bool)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Integer Bool -> m (Either Integer Bool))
-> Either Integer Bool -> m (Either Integer Bool)
forall a b. (a -> b) -> a -> b
$
    if Integer
d Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
1
      then Integer -> Either Integer Bool
forall a b. a -> Either a b
Left Integer
d
      else
        let (Int
v, Integer
m) = Integer -> Integer -> (Int, Integer)
forall r. Euclidean r => r -> r -> (Int, r)
splitFactor Integer
2 (Integer
n Integer -> Integer -> Integer
forall r. Group r => r -> r -> r
-Integer
1)
            b0 :: Integer
b0 = Integer -> Integer -> Integer -> Integer
forall a r. (Integral a, Euclidean r) => r -> r -> a -> r
modPow Integer
n Integer
a Integer
m
            bs :: [Integer]
bs = Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take (Int
v Int -> Int -> Int
forall r. Additive r => r -> r -> r
+ Int
1) ([Integer] -> [Integer]) -> [Integer] -> [Integer]
forall a b. (a -> b) -> a -> b
$ (Integer -> Integer) -> Integer -> [Integer]
forall a. (a -> a) -> a -> [a]
iterate (\Integer
b -> Integer
b Integer -> Integer -> Integer
forall r. Multiplicative r => r -> r -> r
* Integer
b Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n) Integer
b0
         in if Integer
b0 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1
              then Bool -> Either Integer Bool
forall a b. b -> Either a b
Right Bool
True
              else case Integer -> [Integer] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex Integer
1 [Integer]
bs of
                Maybe Int
Nothing -> Bool -> Either Integer Bool
forall a b. b -> Either a b
Right Bool
False
                Just Int
j ->
                  let g :: Integer
g = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
P.gcd ([Integer]
bs [Integer] -> Int -> Integer
forall a. [a] -> Int -> a
!! Int
j Integer -> Integer -> Integer
forall r. Group r => r -> r -> r
- Integer
1) Integer
n
                   in if Integer
g Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 Bool -> Bool -> Bool
|| Integer
g Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
n then Bool -> Either Integer Bool
forall a b. b -> Either a b
Right Bool
True else Integer -> Either Integer Bool
forall a b. a -> Either a b
Left Integer
g