-- | Prime numbers and related number theoretical stuff.

{-# LANGUAGE BangPatterns #-}

module Math.FiniteField.Primes 
  ( -- * Integer logarithm
    integerLog2
  , ceilingLog2
    -- * Integer square root
  , isSquare
  , integerSquareRoot
  , ceilingSquareRoot
  , integerSquareRoot' 
  , integerSquareRootNewton'
    -- * Elementary number theory
  , divides
  , divisors, squareFreeDivisors, squareFreeDivisors_  
  , divisorSum , divisorSum'
  , moebiusMu , eulerTotient , liouvilleLambda
    -- * List of prime numbers
  , primes
  , primesSimple
  , primesTMWE
    -- * Prime factorization
  , factorize, factorizeNaive
  , productOfFactors
  , integerFactorsTrialDivision
  , groupIntegerFactors
    -- * Modulo @m@ arithmetic
  , powerMod
    -- * Prime testing
  , isPrimeTrialDivision
  , millerRabinPrimalityTest
  , isProbablyPrime
  , isVeryProbablyPrime
  )
  where

--------------------------------------------------------------------------------

import Data.List ( group , sort , foldl' )
import Data.Bits

import System.Random

import Math.FiniteField.Sign
import Math.FiniteField.Misc


--------------------------------------------------------------------------------
-- Integer logarithm

-- | Largest integer @k@ such that @2^k@ is smaller or equal to @n@
integerLog2 :: Integer -> Integer
integerLog2 :: Integer -> Integer
integerLog2 Integer
n = forall {t} {a}. (Num t, Num a, Bits t) => t -> a
go Integer
n where
  go :: t -> a
go t
0 = -a
1
  go t
k = a
1 forall a. Num a => a -> a -> a
+ t -> a
go (forall a. Bits a => a -> Int -> a
shiftR t
k Int
1)

-- | Smallest integer @k@ such that @2^k@ is larger or equal to @n@
ceilingLog2 :: Integer -> Integer
ceilingLog2 :: Integer -> Integer
ceilingLog2 Integer
0 = Integer
0
ceilingLog2 Integer
n = Integer
1 forall a. Num a => a -> a -> a
+ forall {t} {a}. (Num t, Num a, Bits t) => t -> a
go (Integer
nforall a. Num a => a -> a -> a
-Integer
1) where
  go :: t -> a
go t
0 = -a
1
  go t
k = a
1 forall a. Num a => a -> a -> a
+ t -> a
go (forall a. Bits a => a -> Int -> a
shiftR t
k Int
1)
  
--------------------------------------------------------------------------------
-- Integer square root

isSquare :: Integer -> Bool
isSquare :: Integer -> Bool
isSquare Integer
n = 
  if (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> a -> a
mod Integer
n Integer
32) forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int]
rs 
    then forall a b. (a, b) -> b
snd (Integer -> (Integer, Integer)
integerSquareRoot' Integer
n) forall a. Eq a => a -> a -> Bool
== Integer
0
    else Bool
False
  where
    rs :: [Int]
rs = [Int
0,Int
1,Int
4,Int
9,Int
16,Int
17,Int
25] :: [Int]
    
-- | Integer square root (largest integer whose square is smaller or equal to the input)
-- using Newton's method, with a faster (for large numbers) inital guess based on bit shifts.
integerSquareRoot :: Integer -> Integer
integerSquareRoot :: Integer -> Integer
integerSquareRoot = forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> (Integer, Integer)
integerSquareRoot'

-- | Smallest integer whose square is larger or equal to the input
ceilingSquareRoot :: Integer -> Integer
ceilingSquareRoot :: Integer -> Integer
ceilingSquareRoot Integer
n = (if Integer
rforall a. Ord a => a -> a -> Bool
>Integer
0 then Integer
uforall a. Num a => a -> a -> a
+Integer
1 else Integer
u) where (Integer
u,Integer
r) = Integer -> (Integer, Integer)
integerSquareRoot' Integer
n 

-- | We also return the excess residue; that is
--
-- > (a,r) = integerSquareRoot' n
-- 
-- means that
--
-- > a*a + r = n
-- > a*a <= n < (a+1)*(a+1)
integerSquareRoot' :: Integer -> (Integer,Integer)
integerSquareRoot' :: Integer -> (Integer, Integer)
integerSquareRoot' Integer
n
  | Integer
nforall a. Ord a => a -> a -> Bool
<Integer
0 = forall a. HasCallStack => [Char] -> a
error [Char]
"integerSquareRoot: negative input"
  | Integer
nforall a. Ord a => a -> a -> Bool
<Integer
2 = (Integer
n,Integer
0)
  | Bool
otherwise = Integer -> (Integer, Integer)
go Integer
firstGuess 
  where
    k :: Integer
k = Integer -> Integer
integerLog2 Integer
n
    firstGuess :: Integer
firstGuess = Integer
2forall a b. (Num a, Integral b) => a -> b -> a
^(forall a. Integral a => a -> a -> a
div (Integer
kforall a. Num a => a -> a -> a
+Integer
2) Integer
2) -- !! note that (div (k+1) 2) is NOT enough !!
    go :: Integer -> (Integer, Integer)
go Integer
a = 
      if Integer
m forall a. Ord a => a -> a -> Bool
< Integer
a
        then Integer -> (Integer, Integer)
go Integer
a' 
        else (Integer
a, Integer
r forall a. Num a => a -> a -> a
+ Integer
aforall a. Num a => a -> a -> a
*(Integer
mforall a. Num a => a -> a -> a
-Integer
a))
      where
        (Integer
m,Integer
r) = forall a. Integral a => a -> a -> (a, a)
divMod Integer
n Integer
a
        a' :: Integer
a' = forall a. Integral a => a -> a -> a
div (Integer
m forall a. Num a => a -> a -> a
+ Integer
a) Integer
2

-- | Newton's method without an initial guess. For very small numbers (<10^10) it
-- is somewhat faster than the above version.
integerSquareRootNewton' :: Integer -> (Integer,Integer)
integerSquareRootNewton' :: Integer -> (Integer, Integer)
integerSquareRootNewton' Integer
n
  | Integer
nforall a. Ord a => a -> a -> Bool
<Integer
0 = forall a. HasCallStack => [Char] -> a
error [Char]
"integerSquareRootNewton: negative input"
  | Integer
nforall a. Ord a => a -> a -> Bool
<Integer
2 = (Integer
n,Integer
0)
  | Bool
otherwise = Integer -> (Integer, Integer)
go (forall a. Integral a => a -> a -> a
div Integer
n Integer
2) 
  where
    go :: Integer -> (Integer, Integer)
go Integer
a = 
      if Integer
m forall a. Ord a => a -> a -> Bool
< Integer
a
        then Integer -> (Integer, Integer)
go Integer
a' 
        else (Integer
a, Integer
r forall a. Num a => a -> a -> a
+ Integer
aforall a. Num a => a -> a -> a
*(Integer
mforall a. Num a => a -> a -> a
-Integer
a))
      where
        (Integer
m,Integer
r) = forall a. Integral a => a -> a -> (a, a)
divMod Integer
n Integer
a
        a' :: Integer
a' = forall a. Integral a => a -> a -> a
div (Integer
m forall a. Num a => a -> a -> a
+ Integer
a) Integer
2

{-
-- brute force test of integer square root
isqrt_test n1 n2 = 
  [ k 
  | k<-[n1..n2] 
  , let (a,r) = integerSquareRoot' k
  , (a*a+r/=k) || (a*a>k) || (a+1)*(a+1)<=k 
  ]
-}

--------------------------------------------------------------------------------

-- | @d `divides` n@
divides :: Integer -> Integer -> Bool
divides :: Integer -> Integer -> Bool
divides !Integer
d !Integer
n = (forall a. Integral a => a -> a -> a
mod Integer
n Integer
d forall a. Eq a => a -> a -> Bool
== Integer
0)

{-# SPECIALIZE moebiusMu :: Int     -> Int     #-}
{-# SPECIALIZE moebiusMu :: Integer -> Integer #-}
-- | The Moebius mu function (naive implementation)
moebiusMu :: (Integral a, Num b) => a -> b
moebiusMu :: forall a b. (Integral a, Num b) => a -> b
moebiusMu a
n 
  | forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (forall a. Ord a => a -> a -> Bool
>Int
1) [Int]
expos       =  b
0
  | forall a. Integral a => a -> Bool
even (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
primes) =  b
1
  | Bool
otherwise            = -b
1
  where
    factors :: [(Integer, Int)]
factors = [Integer] -> [(Integer, Int)]
groupIntegerFactors forall a b. (a -> b) -> a -> b
$ Integer -> [Integer]
integerFactorsTrialDivision forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n
    ([Integer]
primes,[Int]
expos) = forall a b. [(a, b)] -> ([a], [b])
unzip [(Integer, Int)]
factors

{-# SPECIALIZE liouvilleLambda :: Int     -> Int     #-}
{-# SPECIALIZE liouvilleLambda :: Integer -> Integer #-}
-- | The Liouville lambda function (naive implementation)
liouvilleLambda :: (Integral a, Num b) => a -> b
liouvilleLambda :: forall a b. (Integral a, Num b) => a -> b
liouvilleLambda a
n = 
  if forall a. Integral a => a -> Bool
odd (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a. Num a => a -> a -> a
(+) Int
0 forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(Integer, Int)]
grps)
    then -b
1
    else  b
1
  where
    grps :: [(Integer, Int)]
grps = [Integer] -> [(Integer, Int)]
groupIntegerFactors forall a b. (a -> b) -> a -> b
$ Integer -> [Integer]
integerFactorsTrialDivision forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n

-- | Sum ofthe of the divisors
divisorSum :: Integer -> Integer
divisorSum :: Integer -> Integer
divisorSum Integer
n = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a. Num a => a -> a -> a
(+) Integer
0 [ Integer
d | Integer
d <- Integer -> [Integer]
divisors Integer
n ]

-- | Sum of @k@-th powers of the divisors
divisorSum' :: Int -> Integer -> Integer
divisorSum' :: Int -> Integer -> Integer
divisorSum' Int
k Integer
n = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a. Num a => a -> a -> a
(+) Integer
0 [ Integer
dforall a b. (Num a, Integral b) => a -> b -> a
^Int
k | Integer
d <- Integer -> [Integer]
divisors Integer
n ]

-- | Euler's totient function
eulerTotient :: Integer -> Integer
eulerTotient :: Integer -> Integer
eulerTotient Integer
n = forall a. Integral a => a -> a -> a
div Integer
n Integer
prodp forall a. Num a => a -> a -> a
* Integer
prodp1 where
  grps :: [(Integer, Int)]
grps   = [Integer] -> [(Integer, Int)]
groupIntegerFactors forall a b. (a -> b) -> a -> b
$ Integer -> [Integer]
integerFactorsTrialDivision Integer
n
  ps :: [Integer]
ps     = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Integer, Int)]
grps
  prodp :: Integer
prodp  = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a. Num a => a -> a -> a
(*) Integer
1 [ Integer
p   | Integer
p <- [Integer]
ps ] 
  prodp1 :: Integer
prodp1 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a. Num a => a -> a -> a
(*) Integer
1 [ Integer
pforall a. Num a => a -> a -> a
-Integer
1 | Integer
p <- [Integer]
ps ] 

-- | Divisors of @n@ (note: the result is /not/ ordered!)
divisors :: Integer -> [Integer]
divisors :: Integer -> [Integer]
divisors Integer
n = [ forall {b}. Integral b => [b] -> Integer
f [Int]
tup | [Int]
tup <- [Int] -> [[Int]]
tuples' [Int]
expos ] where
  grps :: [(Integer, Int)]
grps = [Integer] -> [(Integer, Int)]
groupIntegerFactors forall a b. (a -> b) -> a -> b
$ Integer -> [Integer]
integerFactorsTrialDivision Integer
n
  ([Integer]
ps,[Int]
expos) = forall a b. [(a, b)] -> ([a], [b])
unzip [(Integer, Int)]
grps
  f :: [b] -> Integer
f [b]
es = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a. Num a => a -> a -> a
(*) Integer
1 forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall a b. (Num a, Integral b) => a -> b -> a
(^) [Integer]
ps [b]
es

-- | List of square-free divisors together with their Mobius mu value
-- (note: the result is /not/ ordered!)
squareFreeDivisors :: Integer -> [(Integer,Sign)]
squareFreeDivisors :: Integer -> [(Integer, Sign)]
squareFreeDivisors Integer
n = forall a b. (a -> b) -> [a] -> [b]
map forall {t :: * -> *} {a}. (Foldable t, Num a) => t a -> (a, Sign)
f (forall a. [a] -> [[a]]
sublists [Integer]
primes) where
  grps :: [(Integer, Int)]
grps = [Integer] -> [(Integer, Int)]
groupIntegerFactors forall a b. (a -> b) -> a -> b
$ Integer -> [Integer]
integerFactorsTrialDivision Integer
n
  primes :: [Integer]
primes = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Integer, Int)]
grps
  f :: t a -> (a, Sign)
f t a
ps = ( forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a. Num a => a -> a -> a
(*) a
1 t a
ps , if forall a. Integral a => a -> Bool
even (forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
ps) then Sign
Plus else Sign
Minus)

-- | List of square-free divisors 
-- (note: the result is /not/ ordered!)
squareFreeDivisors_ :: Integer -> [Integer]
squareFreeDivisors_ :: Integer -> [Integer]
squareFreeDivisors_ Integer
n = forall a b. (a -> b) -> [a] -> [b]
map forall {t :: * -> *} {a}. (Foldable t, Num a) => t a -> a
f (forall a. [a] -> [[a]]
sublists [Integer]
primes) where
  grps :: [(Integer, Int)]
grps = [Integer] -> [(Integer, Int)]
groupIntegerFactors forall a b. (a -> b) -> a -> b
$ Integer -> [Integer]
integerFactorsTrialDivision Integer
n
  primes :: [Integer]
primes = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Integer, Int)]
grps
  f :: t a -> a
f t a
ps = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a. Num a => a -> a -> a
(*) a
1 t a
ps

-- | To avoid cyclic dependencies, I made a local copy of this...
sublists :: [a] -> [[a]]
sublists :: forall a. [a] -> [[a]]
sublists [] = [[]]
sublists (a
x:[a]
xs) = forall a. [a] -> [[a]]
sublists [a]
xs forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map (a
xforall a. a -> [a] -> [a]
:) (forall a. [a] -> [[a]]
sublists [a]
xs)  

--------------------------------------------------------------------------------
-- List of prime numbers 

-- | Infinite list of primes, using the TMWE algorithm.
primes :: [Integer]
primes :: [Integer]
primes = [Integer]
primesTMWE

-- | A relatively simple but still quite fast implementation of list of primes.
-- By Will Ness <http://www.haskell.org/pipermail/haskell-cafe/2009-November/068441.html>
primesSimple :: [Integer]
primesSimple :: [Integer]
primesSimple = Integer
2 forall a. a -> [a] -> [a]
: Integer
3 forall a. a -> [a] -> [a]
: Int -> [Integer] -> Integer -> [Integer]
sieve Int
0 [Integer]
primes' Integer
5 where
  primes' :: [Integer]
primes' = forall a. [a] -> [a]
tail [Integer]
primesSimple
  sieve :: Int -> [Integer] -> Integer -> [Integer]
sieve Int
k (Integer
p:[Integer]
ps) Integer
x = Int -> [Integer] -> [Integer]
noDivs Int
k [Integer]
h forall a. [a] -> [a] -> [a]
++ Int -> [Integer] -> Integer -> [Integer]
sieve (Int
kforall a. Num a => a -> a -> a
+Int
1) [Integer]
ps (Integer
tforall a. Num a => a -> a -> a
+Integer
2) where
    t :: Integer
t = Integer
pforall a. Num a => a -> a -> a
*Integer
p 
    h :: [Integer]
h = [Integer
x,Integer
xforall a. Num a => a -> a -> a
+Integer
2..Integer
tforall a. Num a => a -> a -> a
-Integer
2]
  noDivs :: Int -> [Integer] -> [Integer]
noDivs Int
k = forall a. (a -> Bool) -> [a] -> [a]
filter (\Integer
x -> forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Integer
y -> forall a. Integral a => a -> a -> a
rem Integer
x Integer
y forall a. Eq a => a -> a -> Bool
/= Integer
0) (forall a. Int -> [a] -> [a]
take Int
k [Integer]
primes'))
  
-- | List of primes, using tree merge with wheel. Code by Will Ness.
primesTMWE :: [Integer]
primesTMWE :: [Integer]
primesTMWE = Integer
2forall a. a -> [a] -> [a]
:Integer
3forall a. a -> [a] -> [a]
:Integer
5forall a. a -> [a] -> [a]
:Integer
7forall a. a -> [a] -> [a]
: forall {a}. (Eq a, Num a) => a -> [a] -> [a] -> [a]
gaps Integer
11 [Integer]
wheel (forall {a}. Ord a => [[a]] -> [a]
fold3t forall a b. (a -> b) -> a -> b
$ forall {t}. (Eq t, Num t) => t -> [t] -> [t] -> [[t]]
roll Integer
11 [Integer]
wheel [Integer]
primes') where                                                             

  primes' :: [Integer]
primes' = Integer
11forall a. a -> [a] -> [a]
: forall {a}. (Eq a, Num a) => a -> [a] -> [a] -> [a]
gaps Integer
13 (forall a. [a] -> [a]
tail [Integer]
wheel) (forall {a}. Ord a => [[a]] -> [a]
fold3t forall a b. (a -> b) -> a -> b
$ forall {t}. (Eq t, Num t) => t -> [t] -> [t] -> [[t]]
roll Integer
11 [Integer]
wheel [Integer]
primes')
  fold3t :: [[a]] -> [a]
fold3t ((a
x:[a]
xs): ~([a]
ys:[a]
zs:[[a]]
t)) 
    = a
x forall a. a -> [a] -> [a]
: forall {a}. Ord a => [a] -> [a] -> [a]
union [a]
xs (forall {a}. Ord a => [a] -> [a] -> [a]
union [a]
ys [a]
zs) forall {a}. Ord a => [a] -> [a] -> [a]
`union` [[a]] -> [a]
fold3t (forall {a}. Ord a => [[a]] -> [[a]]
pairs [[a]]
t)            
  pairs :: [[a]] -> [[a]]
pairs ((a
x:[a]
xs):[a]
ys:[[a]]
t) = (a
x forall a. a -> [a] -> [a]
: forall {a}. Ord a => [a] -> [a] -> [a]
union [a]
xs [a]
ys) forall a. a -> [a] -> [a]
: [[a]] -> [[a]]
pairs [[a]]
t 
  wheel :: [Integer]
wheel = Integer
2forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
8forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:  
          Integer
4forall a. a -> [a] -> [a]
:Integer
8forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:Integer
6forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:Integer
4forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:Integer
10forall a. a -> [a] -> [a]
:Integer
2forall a. a -> [a] -> [a]
:Integer
10forall a. a -> [a] -> [a]
:[Integer]
wheel 
  gaps :: a -> [a] -> [a] -> [a]
gaps a
k ws :: [a]
ws@(a
w:[a]
t) cs :: [a]
cs@(~(a
c:[a]
u)) 
    | a
kforall a. Eq a => a -> a -> Bool
==a
c  = a -> [a] -> [a] -> [a]
gaps (a
kforall a. Num a => a -> a -> a
+a
w) [a]
t [a]
u              
    | Bool
True  = a
k forall a. a -> [a] -> [a]
: a -> [a] -> [a] -> [a]
gaps (a
kforall a. Num a => a -> a -> a
+a
w) [a]
t [a]
cs  
  roll :: t -> [t] -> [t] -> [[t]]
roll t
k ws :: [t]
ws@(t
w:[t]
t) ps :: [t]
ps@(~(t
p:[t]
u))
    | t
kforall a. Eq a => a -> a -> Bool
==t
p  = forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl (\t
c t
d->t
cforall a. Num a => a -> a -> a
+t
pforall a. Num a => a -> a -> a
*t
d) (t
pforall a. Num a => a -> a -> a
*t
p) [t]
ws forall a. a -> [a] -> [a]
: t -> [t] -> [t] -> [[t]]
roll (t
kforall a. Num a => a -> a -> a
+t
w) [t]
t [t]
u              
    | Bool
True  = t -> [t] -> [t] -> [[t]]
roll (t
kforall a. Num a => a -> a -> a
+t
w) [t]
t [t]
ps   

  minus :: [a] -> [a] -> [a]
minus xxs :: [a]
xxs@(a
x:[a]
xs) yys :: [a]
yys@(a
y:[a]
ys) = case forall a. Ord a => a -> a -> Ordering
compare a
x a
y of 
    Ordering
LT -> a
x forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
minus [a]
xs  [a]
yys
    Ordering
EQ ->     [a] -> [a] -> [a]
minus [a]
xs  [a]
ys 
    Ordering
GT ->     [a] -> [a] -> [a]
minus [a]
xxs [a]
ys
  minus [a]
xs [] = [a]
xs
  minus [] [a]
_  = []
  
  union :: [a] -> [a] -> [a]
union xxs :: [a]
xxs@(a
x:[a]
xs) yys :: [a]
yys@(a
y:[a]
ys) = case forall a. Ord a => a -> a -> Ordering
compare a
x a
y of 
    Ordering
LT -> a
x forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
union [a]
xs  [a]
yys
    Ordering
EQ -> a
x forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
union [a]
xs  [a]
ys 
    Ordering
GT -> a
y forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
union [a]
xxs [a]
ys
  union [a]
xs [] = [a]
xs
  union [] [a]
ys =[a]
ys

--------------------------------------------------------------------------------
-- Prime factorization

factorize :: Integer -> [(Integer,Int)]
factorize :: Integer -> [(Integer, Int)]
factorize = Integer -> [(Integer, Int)]
factorizeNaive

factorizeNaive :: Integer -> [(Integer,Int)]
factorizeNaive :: Integer -> [(Integer, Int)]
factorizeNaive = [Integer] -> [(Integer, Int)]
groupIntegerFactors forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> [Integer]
integerFactorsTrialDivision

productOfFactors :: [(Integer,Int)] -> Integer
productOfFactors :: [(Integer, Int)] -> Integer
productOfFactors = [Integer] -> Integer
productInterleaved forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall {t}. (Bits t, Num t) => t -> Int -> t
pow) where
  pow :: t -> Int -> t
pow t
_ Int
0 = t
1
  pow t
p Int
1 = t
p
  pow t
2 Int
n = forall a. Bits a => a -> Int -> a
shiftL t
1 Int
n
  pow t
p Int
2 = t
pforall a. Num a => a -> a -> a
*t
p
  pow t
p Int
n = if forall a. Integral a => a -> Bool
even Int
n
              then     (t -> Int -> t
pow t
p (forall a. Bits a => a -> Int -> a
shiftR Int
n Int
1))forall a b. (Num a, Integral b) => a -> b -> a
^Integer
2
              else t
p forall a. Num a => a -> a -> a
* (t -> Int -> t
pow t
p (forall a. Bits a => a -> Int -> a
shiftR Int
n Int
1))forall a b. (Num a, Integral b) => a -> b -> a
^Integer
2 

-- | Groups integer factors. Example: from [2,2,2,3,3,5] we produce [(2,3),(3,2),(5,1)]  
groupIntegerFactors :: [Integer] -> [(Integer,Int)]
groupIntegerFactors :: [Integer] -> [(Integer, Int)]
groupIntegerFactors = forall a b. (a -> b) -> [a] -> [b]
map forall {a}. [a] -> (a, Int)
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Eq a => [a] -> [[a]]
group forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => [a] -> [a]
sort where
  f :: [a] -> (a, Int)
f [a]
xs = (forall a. [a] -> a
head [a]
xs, forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)

-- | The naive trial division algorithm.
integerFactorsTrialDivision :: Integer -> [Integer]
integerFactorsTrialDivision :: Integer -> [Integer]
integerFactorsTrialDivision Integer
n 
  | Integer
nforall a. Ord a => a -> a -> Bool
<Integer
1 = forall a. HasCallStack => [Char] -> a
error [Char]
"integerFactorsTrialDivision: n should be at least 1"
  | Bool
otherwise = forall {a}. Integral a => [a] -> a -> [a]
go [Integer]
primes Integer
n 
  where
    go :: [a] -> a -> [a]
go [a]
_  a
1 = []
    go [a]
rs a
k = [a] -> a -> [a]
sub [a]
ps a
k where
      sub :: [a] -> a -> [a]
sub [] a
k = [a
k]
      sub qqs :: [a]
qqs@(a
q:[a]
qs) a
k = case forall a. Integral a => a -> a -> a
mod a
k a
q of
        a
0 -> a
q forall a. a -> [a] -> [a]
: [a] -> a -> [a]
go [a]
qqs (forall a. Integral a => a -> a -> a
div a
k a
q)
        a
_ -> [a] -> a -> [a]
sub [a]
qs a
k
      ps :: [a]
ps = forall a. (a -> Bool) -> [a] -> [a]
takeWhile (\a
p -> a
pforall a. Num a => a -> a -> a
*a
p forall a. Ord a => a -> a -> Bool
<= a
k) [a]
rs  
{-
    go 1 = []
    go k = sub ps k where
      sub [] k = [k]
      sub (q:qs) k = case mod k q of
        0 -> q : go (div k q)
        _ -> sub qs k
      ps = takeWhile (\p -> p*p <= k) primes
-}

{-    
-- brute force testing of factors
ifactorsTest :: (Integer -> [Integer]) -> Integer -> Bool
ifactorsTest alg n = and [ product (alg k) == k | k<-[1..n] ]   
-}

--------------------------------------------------------------------------------
-- Modulo @m@ arithmetic

-- | Efficient powers modulo m.
-- 
-- > powerMod a k m == (a^k) `mod` m
powerMod :: Integer -> Integer -> Integer -> Integer
powerMod :: Integer -> Integer -> Integer -> Integer
powerMod Integer
a' Integer
k Integer
m = {- debug bs $ -} Integer -> [Bool] -> Integer
go Integer
a [Bool]
bs where

  bs :: [Bool]
bs = forall {t}. (Num t, Bits t) => t -> [Bool]
bin Integer
k

  bin :: t -> [Bool]
bin t
0 = []
  bin t
x = (t
x forall a. Bits a => a -> a -> a
.&. t
1 forall a. Eq a => a -> a -> Bool
/= t
0) forall a. a -> [a] -> [a]
: t -> [Bool]
bin (forall a. Bits a => a -> Int -> a
shiftR t
x Int
1)

  a :: Integer
a = forall a. Integral a => a -> a -> a
mod Integer
a' Integer
m

  go :: Integer -> [Bool] -> Integer
go Integer
_ [] = Integer
1
  go Integer
x (Bool
b:[Bool]
bs) = -- debug (x,b) $ 
    if Bool
b 
      then forall a. Integral a => a -> a -> a
mod (Integer
xforall a. Num a => a -> a -> a
*Integer
rest) Integer
m
      else Integer
rest
    where 
      rest :: Integer
rest = Integer -> [Bool] -> Integer
go (forall a. Integral a => a -> a -> a
mod (Integer
xforall a. Num a => a -> a -> a
*Integer
x) Integer
m) [Bool]
bs 
      
--------------------------------------------------------------------------------
-- * Prime testing

-- | Prime testing using trial division 
isPrimeTrialDivision :: Integer -> Bool
isPrimeTrialDivision :: Integer -> Bool
isPrimeTrialDivision Integer
n = forall (t :: * -> *). Foldable t => t Bool -> Bool
and [ Bool -> Bool
not (Integer -> Integer -> Bool
divides Integer
p Integer
n) | Integer
p <- [Integer]
ps ] where
  ps :: [Integer]
ps = forall a. (a -> Bool) -> [a] -> [a]
takeWhile (forall a. Ord a => a -> a -> Bool
<= Integer
nsqrt) [Integer]
primes 
  nsqrt :: Integer
nsqrt = Integer -> Integer
integerSquareRoot Integer
n
  
-- | Miller-Rabin Primality Test (taken from Haskell wiki). 
-- We test the primality of the first argument @n@ by using the second argument @a@ as a candidate witness.
-- If it returs @False@, then @n@ is composite. If it returns @True@, then @n@ is either prime or composite.
--
-- A random choice between @2@ and @(n-2)@ is a good choice for @a@.
millerRabinPrimalityTest :: Integer -> Integer -> Bool
millerRabinPrimalityTest :: Integer -> Integer -> Bool
millerRabinPrimalityTest Integer
n Integer
a
  | Integer
a forall a. Ord a => a -> a -> Bool
<= Integer
1 Bool -> Bool -> Bool
|| Integer
a forall a. Ord a => a -> a -> Bool
>= Integer
nforall a. Num a => a -> a -> a
-Integer
1 = 
      forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"millerRabinPrimalityTest: a out of range (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Integer
a forall a. [a] -> [a] -> [a]
++ [Char]
" for "forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Integer
n forall a. [a] -> [a] -> [a]
++ [Char]
")" 
  | Integer
n forall a. Ord a => a -> a -> Bool
< Integer
2 = Bool
False
  | forall a. Integral a => a -> Bool
even Integer
n = Bool
False
  | Integer
b0 forall a. Eq a => a -> a -> Bool
== Integer
1 Bool -> Bool -> Bool
|| Integer
b0 forall a. Eq a => a -> a -> Bool
== Integer
n' = Bool
True
  | Bool
otherwise = [Integer] -> Bool
iter (forall a. [a] -> [a]
tail [Integer]
b)
  where
    n' :: Integer
n' = Integer
nforall a. Num a => a -> a -> a
-Integer
1
    (Integer
k,Integer
m) = forall a. Integral a => a -> (a, a)
find2km Integer
n'
    b0 :: Integer
b0 = forall a. Integral a => a -> a -> a -> a
powMod Integer
n Integer
a Integer
m
    b :: [Integer]
b = forall a. Int -> [a] -> [a]
take (forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
k) forall a b. (a -> b) -> a -> b
$ forall a. (a -> a) -> a -> [a]
iterate (forall a. Integral a => a -> a -> a
squareMod Integer
n) Integer
b0
    iter :: [Integer] -> Bool
iter [] = Bool
False
    iter (Integer
x:[Integer]
xs)
      | Integer
x forall a. Eq a => a -> a -> Bool
== Integer
1  = Bool
False
      | Integer
x forall a. Eq a => a -> a -> Bool
== Integer
n' = Bool
True
      | Bool
otherwise = [Integer] -> Bool
iter [Integer]
xs


{-# SPECIALIZE find2km :: Integer -> (Integer,Integer) #-}
find2km :: Integral a => a -> (a,a)
find2km :: forall a. Integral a => a -> (a, a)
find2km a
n = forall {t} {t}. (Num t, Integral t) => t -> t -> (t, t)
f a
0 a
n where 
  f :: t -> t -> (t, t)
f t
k t
m
    | t
r forall a. Eq a => a -> a -> Bool
== t
1 = (t
k,t
m)
    | Bool
otherwise = t -> t -> (t, t)
f (t
kforall a. Num a => a -> a -> a
+t
1) t
q
    where (t
q,t
r) = forall a. Integral a => a -> a -> (a, a)
quotRem t
m t
2        
 
{-# SPECIALIZE pow' :: (Integer -> Integer -> Integer) -> (Integer -> Integer) -> Integer -> Integer -> Integer #-}
pow' :: (Num a, Integral b) => (a -> a -> a) -> (a -> a) -> a -> b -> a
pow' :: forall a b.
(Num a, Integral b) =>
(a -> a -> a) -> (a -> a) -> a -> b -> a
pow' a -> a -> a
_ a -> a
_ a
_ b
0 = a
1
pow' a -> a -> a
mul a -> a
sq a
x' b
n' = forall {t}. Integral t => a -> t -> a -> a
f a
x' b
n' a
1 where 
  f :: a -> t -> a -> a
f !a
x !t
n !a
y
    | t
n forall a. Eq a => a -> a -> Bool
== t
1 = a
x a -> a -> a
`mul` a
y
    | t
r forall a. Eq a => a -> a -> Bool
== t
0 = a -> t -> a -> a
f a
x2 t
q a
y
    | Bool
otherwise = a -> t -> a -> a
f a
x2 t
q (a
x a -> a -> a
`mul` a
y)
    where
      (t
q,t
r) = forall a. Integral a => a -> a -> (a, a)
quotRem t
n t
2
      x2 :: a
x2    = a -> a
sq a
x
 
{-# SPECIALIZE mulMod :: Integer -> Integer -> Integer -> Integer #-}
mulMod :: Integral a => a -> a -> a -> a
mulMod :: forall a. Integral a => a -> a -> a -> a
mulMod !a
a !a
b !a
c = (a
b forall a. Num a => a -> a -> a
* a
c) forall a. Integral a => a -> a -> a
`mod` a
a

{-# SPECIALIZE squareMod :: Integer -> Integer -> Integer #-}
squareMod :: Integral a => a -> a -> a
squareMod :: forall a. Integral a => a -> a -> a
squareMod !a
a !a
b = (a
b forall a. Num a => a -> a -> a
* a
b) forall a. Integral a => a -> a -> a
`rem` a
a

{-# SPECIALIZE powMod :: Integer -> Integer -> Integer -> Integer #-}
powMod :: Integral a => a -> a -> a -> a
powMod :: forall a. Integral a => a -> a -> a -> a
powMod a
m = forall a b.
(Num a, Integral b) =>
(a -> a -> a) -> (a -> a) -> a -> b -> a
pow' (forall a. Integral a => a -> a -> a -> a
mulMod a
m) (forall a. Integral a => a -> a -> a
squareMod a
m)

--------------------------------------------------------------------------------

-- | For very small numbers, we use trial division, for larger numbers, we apply the 
-- Miller-Rabin primality test @log4(n)@ times, with candidate witnesses derived 
-- deterministically from @n@ using a pseudo-random sequence 
-- (which /should be/ based on a cryptographic hash function, but isn\'t, yet). 
--
-- Thus the candidate witnesses should behave essentially like random, but the 
-- resulting function is still a deterministic, pure function.
--
-- TODO: implement the hash sequence, at the moment we use 'System.Random' instead...
--
isProbablyPrime :: Integer -> Bool
isProbablyPrime :: Integer -> Bool
isProbablyPrime Integer
n 
  | Integer
n forall a. Ord a => a -> a -> Bool
< Integer
2      = Bool
False
  | forall a. Integral a => a -> Bool
even Integer
n     = (Integer
nforall a. Eq a => a -> a -> Bool
==Integer
2)
  | Integer
n forall a. Ord a => a -> a -> Bool
< Integer
1000   = forall (t :: * -> *) a. Foldable t => t a -> Int
length (Integer -> [Integer]
integerFactorsTrialDivision Integer
n) forall a. Eq a => a -> a -> Bool
== Int
1
  | Bool
otherwise  = forall (t :: * -> *). Foldable t => t Bool -> Bool
and [ Integer -> Integer -> Bool
millerRabinPrimalityTest Integer
n Integer
a | Integer
a <- [Integer]
witnessList ]
  where
    log2n :: Integer
log2n       = Integer -> Integer
integerLog2 Integer
n 
    nchecks :: Int
nchecks     = Int
1 forall a. Num a => a -> a -> a
+ forall a. Num a => Integer -> a
fromInteger (forall a. Integral a => a -> a -> a
div Integer
log2n Integer
2) :: Int
    witnessList :: [Integer]
witnessList = forall a. Int -> [a] -> [a]
take Int
nchecks [Integer]
pseudoRnds
    pseudoRnds :: [Integer]
pseudoRnds  = Integer
2 forall a. a -> [a] -> [a]
: [ Integer
a | Integer
a <- Integer -> [Integer]
integerRndSequence Integer
n , Integer
a forall a. Ord a => a -> a -> Bool
> Integer
1 Bool -> Bool -> Bool
&& Integer
a forall a. Ord a => a -> a -> Bool
< (Integer
nforall a. Num a => a -> a -> a
-Integer
1) ]

-- | A more exhaustive version of 'isProbablyPrime', this one tests candidate
-- witnesses both the first log4(n) prime numbers and then log4(n) pseudo-random
-- numbers
isVeryProbablyPrime :: Integer -> Bool
isVeryProbablyPrime :: Integer -> Bool
isVeryProbablyPrime Integer
n
  | Integer
n forall a. Ord a => a -> a -> Bool
< Integer
2      = Bool
False
  | forall a. Integral a => a -> Bool
even Integer
n     = (Integer
nforall a. Eq a => a -> a -> Bool
==Integer
2)
  | Integer
n forall a. Ord a => a -> a -> Bool
< Integer
1000   = forall (t :: * -> *) a. Foldable t => t a -> Int
length (Integer -> [Integer]
integerFactorsTrialDivision Integer
n) forall a. Eq a => a -> a -> Bool
== Int
1
  | Bool
otherwise  = forall (t :: * -> *). Foldable t => t Bool -> Bool
and [ Integer -> Integer -> Bool
millerRabinPrimalityTest Integer
n Integer
a | Integer
a <- [Integer]
witnessList ]
  where
    log2n :: Integer
log2n       = Integer -> Integer
integerLog2 Integer
n 
    nchecks :: Int
nchecks     = Int
1 forall a. Num a => a -> a -> a
+ forall a. Num a => Integer -> a
fromInteger (forall a. Integral a => a -> a -> a
div Integer
log2n Integer
2) :: Int
    witnessList :: [Integer]
witnessList = forall a. Int -> [a] -> [a]
take Int
nchecks [Integer]
primes forall a. [a] -> [a] -> [a]
++ forall a. Int -> [a] -> [a]
take Int
nchecks [Integer]
pseudoRnds
    pseudoRnds :: [Integer]
pseudoRnds  = [ Integer
a | Integer
a <- Integer -> [Integer]
integerRndSequence (Integer
nforall a. Num a => a -> a -> a
+Integer
3) , Integer
a forall a. Ord a => a -> a -> Bool
> Integer
1 Bool -> Bool -> Bool
&& Integer
a forall a. Ord a => a -> a -> Bool
< (Integer
nforall a. Num a => a -> a -> a
-Integer
1) ]

--------------------------------------------------------------------------------

-- | Given an integer @n@, we initialize a system random generator with using a 
-- seed derived from @n@ (note that this uses at most 32 or 64 bits), and generate 
-- an infinite sequence of pseudo-random integers between @0..n-1@, generated by 
-- that random generator. 
--
-- Note that this is not really a preferred way of generating such sequences!
-- 
integerRndSequence :: Integer -> [Integer]
integerRndSequence :: Integer -> [Integer]
integerRndSequence Integer
n = forall a g. (Random a, RandomGen g) => (a, a) -> g -> [a]
randomRs (Integer
0,Integer
nforall a. Num a => a -> a -> a
-Integer
1) StdGen
gen where
  gen :: StdGen
gen = Int -> StdGen
mkStdGen forall a b. (a -> b) -> a -> b
$ forall a. Num a => Integer -> a
fromInteger (Integer
n forall a. Num a => a -> a -> a
+ Integer
17 forall a. Num a => a -> a -> a
* Integer -> Integer
integerLog2 Integer
n)

--------------------------------------------------------------------------------