-- | Prime numbers and related number theoretical stuff.

module Math.Combinat.Numbers.Primes 
  ( -- * 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
  , millerRabinPrimalityTest
  , isProbablyPrime
  , isVeryProbablyPrime
  )
  where

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

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

import Math.Combinat.Sign
import Math.Combinat.Helper
import Math.Combinat.Numbers.Integers

-- import Math.Combinat.Sets   ( sublists )       -- cyclic dependency...
import Math.Combinat.Tuples ( tuples'  )

import Data.Bits

import System.Random

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

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

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

{-# SPECIALIZE liouvilleLambda :: Int     -> Int     #-}
{-# SPECIALIZE liouvilleLambda :: Integer -> Integer #-}
-- | The Liouville lambda function
liouvilleLambda :: (Integral a, Num b) => a -> b
liouvilleLambda :: a -> b
liouvilleLambda a
n = 
  if Int -> Bool
forall a. Integral a => a -> Bool
odd ((Int -> Int -> Int) -> Int -> [Int] -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ ((Integer, Int) -> Int) -> [(Integer, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Integer, Int) -> Int
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 ([Integer] -> [(Integer, Int)]) -> [Integer] -> [(Integer, Int)]
forall a b. (a -> b) -> a -> b
$ Integer -> [Integer]
integerFactorsTrialDivision (Integer -> [Integer]) -> Integer -> [Integer]
forall a b. (a -> b) -> a -> b
$ a -> Integer
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 = (Integer -> Integer -> Integer) -> Integer -> [Integer] -> Integer
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Integer -> Integer -> Integer
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 = (Integer -> Integer -> Integer) -> Integer -> [Integer] -> Integer
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+) Integer
0 [ Integer
dInteger -> Int -> Integer
forall 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 = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div Integer
n Integer
prodp Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
prodp1 where
  grps :: [(Integer, Int)]
grps   = [Integer] -> [(Integer, Int)]
groupIntegerFactors ([Integer] -> [(Integer, Int)]) -> [Integer] -> [(Integer, Int)]
forall a b. (a -> b) -> a -> b
$ Integer -> [Integer]
integerFactorsTrialDivision Integer
n
  ps :: [Integer]
ps     = ((Integer, Int) -> Integer) -> [(Integer, Int)] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer, Int) -> Integer
forall a b. (a, b) -> a
fst [(Integer, Int)]
grps
  prodp :: Integer
prodp  = (Integer -> Integer -> Integer) -> Integer -> [Integer] -> Integer
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(*) Integer
1 [ Integer
p   | Integer
p <- [Integer]
ps ] 
  prodp1 :: Integer
prodp1 = (Integer -> Integer -> Integer) -> Integer -> [Integer] -> Integer
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(*) Integer
1 [ Integer
pInteger -> Integer -> Integer
forall 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 = [ [Int] -> Integer
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 ([Integer] -> [(Integer, Int)]) -> [Integer] -> [(Integer, Int)]
forall a b. (a -> b) -> a -> b
$ Integer -> [Integer]
integerFactorsTrialDivision Integer
n
  ([Integer]
ps,[Int]
expos) = [(Integer, Int)] -> ([Integer], [Int])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Integer, Int)]
grps
  f :: [b] -> Integer
f [b]
es = (Integer -> Integer -> Integer) -> Integer -> [Integer] -> Integer
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(*) Integer
1 ([Integer] -> Integer) -> [Integer] -> Integer
forall a b. (a -> b) -> a -> b
$ (Integer -> b -> Integer) -> [Integer] -> [b] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Integer -> b -> Integer
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 = ([Integer] -> (Integer, Sign)) -> [[Integer]] -> [(Integer, Sign)]
forall a b. (a -> b) -> [a] -> [b]
map [Integer] -> (Integer, Sign)
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> (a, Sign)
f ([Integer] -> [[Integer]]
forall a. [a] -> [[a]]
sublists [Integer]
primes) where
  grps :: [(Integer, Int)]
grps = [Integer] -> [(Integer, Int)]
groupIntegerFactors ([Integer] -> [(Integer, Int)]) -> [Integer] -> [(Integer, Int)]
forall a b. (a -> b) -> a -> b
$ Integer -> [Integer]
integerFactorsTrialDivision Integer
n
  primes :: [Integer]
primes = ((Integer, Int) -> Integer) -> [(Integer, Int)] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer, Int) -> Integer
forall a b. (a, b) -> a
fst [(Integer, Int)]
grps
  f :: t a -> (a, Sign)
f t a
ps = ( (a -> a -> a) -> a -> t a -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' a -> a -> a
forall a. Num a => a -> a -> a
(*) a
1 t a
ps , if Int -> Bool
forall a. Integral a => a -> Bool
even (t a -> Int
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 = ([Integer] -> Integer) -> [[Integer]] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map [Integer] -> Integer
forall (t :: * -> *) b. (Foldable t, Num b) => t b -> b
f ([Integer] -> [[Integer]]
forall a. [a] -> [[a]]
sublists [Integer]
primes) where
  grps :: [(Integer, Int)]
grps = [Integer] -> [(Integer, Int)]
groupIntegerFactors ([Integer] -> [(Integer, Int)]) -> [Integer] -> [(Integer, Int)]
forall a b. (a -> b) -> a -> b
$ Integer -> [Integer]
integerFactorsTrialDivision Integer
n
  primes :: [Integer]
primes = ((Integer, Int) -> Integer) -> [(Integer, Int)] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer, Int) -> Integer
forall a b. (a, b) -> a
fst [(Integer, Int)]
grps
  f :: t b -> b
f t b
ps = (b -> b -> b) -> b -> t b -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' b -> b -> b
forall a. Num a => a -> a -> a
(*) b
1 t b
ps

-- | To avoid cyclic dependencies, I made a local copy of this...
sublists :: [a] -> [[a]]
sublists :: [a] -> [[a]]
sublists [] = [[]]
sublists (a
x:[a]
xs) = [a] -> [[a]]
forall a. [a] -> [[a]]
sublists [a]
xs [[a]] -> [[a]] -> [[a]]
forall a. [a] -> [a] -> [a]
++ ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (a
xa -> [a] -> [a]
forall a. a -> [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 Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: Integer
3 Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: Int -> [Integer] -> Integer -> [Integer]
sieve Int
0 [Integer]
primes' Integer
5 where
  primes' :: [Integer]
primes' = [Integer] -> [Integer]
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 [Integer] -> [Integer] -> [Integer]
forall a. [a] -> [a] -> [a]
++ Int -> [Integer] -> Integer -> [Integer]
sieve (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [Integer]
ps (Integer
tInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
2) where
    t :: Integer
t = Integer
pInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
p 
    h :: [Integer]
h = [Integer
x,Integer
xInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
2..Integer
tInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
2]
  noDivs :: Int -> [Integer] -> [Integer]
noDivs Int
k = (Integer -> Bool) -> [Integer] -> [Integer]
forall a. (a -> Bool) -> [a] -> [a]
filter (\Integer
x -> (Integer -> Bool) -> [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Integer
y -> Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
rem Integer
x Integer
y Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
0) (Int -> [Integer] -> [Integer]
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
2Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
:Integer
3Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
:Integer
5Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
:Integer
7Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: Integer -> [Integer] -> [Integer] -> [Integer]
forall a. (Eq a, Num a) => a -> [a] -> [a] -> [a]
gaps Integer
11 [Integer]
wheel ([[Integer]] -> [Integer]
forall a. Ord a => [[a]] -> [a]
fold3t ([[Integer]] -> [Integer]) -> [[Integer]] -> [Integer]
forall a b. (a -> b) -> a -> b
$ Integer -> [Integer] -> [Integer] -> [[Integer]]
forall a. (Eq a, Num a) => a -> [a] -> [a] -> [[a]]
roll Integer
11 [Integer]
wheel [Integer]
primes') where                                                             

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

  minus :: [a] -> [a] -> [a]
minus xxs :: [a]
xxs@(a
x:[a]
xs) yys :: [a]
yys@(a
y:[a]
ys) = case a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
x a
y of 
    Ordering
LT -> a
x a -> [a] -> [a]
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 a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
x a
y of 
    Ordering
LT -> a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
union [a]
xs  [a]
yys
    Ordering
EQ -> a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
union [a]
xs  [a]
ys 
    Ordering
GT -> a
y a -> [a] -> [a]
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 ([Integer] -> [(Integer, Int)])
-> (Integer -> [Integer]) -> Integer -> [(Integer, Int)]
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 ([Integer] -> Integer)
-> ([(Integer, Int)] -> [Integer]) -> [(Integer, Int)] -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Integer, Int) -> Integer) -> [(Integer, Int)] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map ((Integer -> Int -> Integer) -> (Integer, Int) -> Integer
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Integer -> Int -> Integer
forall a. (Bits a, Num a) => a -> Int -> a
pow) where
  pow :: a -> Int -> a
pow a
_ Int
0 = a
1
  pow a
p Int
1 = a
p
  pow a
2 Int
n = a -> Int -> a
forall a. Bits a => a -> Int -> a
shiftL a
1 Int
n
  pow a
p Int
2 = a
pa -> a -> a
forall a. Num a => a -> a -> a
*a
p
  pow a
p Int
n = if Int -> Bool
forall a. Integral a => a -> Bool
even Int
n
              then     (a -> Int -> a
pow a
p (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
shiftR Int
n Int
1))a -> Integer -> a
forall a b. (Num a, Integral b) => a -> b -> a
^Integer
2
              else a
p a -> a -> a
forall a. Num a => a -> a -> a
* (a -> Int -> a
pow a
p (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
shiftR Int
n Int
1))a -> Integer -> a
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 = ([Integer] -> (Integer, Int)) -> [[Integer]] -> [(Integer, Int)]
forall a b. (a -> b) -> [a] -> [b]
map [Integer] -> (Integer, Int)
forall a. [a] -> (a, Int)
f ([[Integer]] -> [(Integer, Int)])
-> ([Integer] -> [[Integer]]) -> [Integer] -> [(Integer, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Integer] -> [[Integer]]
forall a. Eq a => [a] -> [[a]]
group ([Integer] -> [[Integer]])
-> ([Integer] -> [Integer]) -> [Integer] -> [[Integer]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Integer] -> [Integer]
forall a. Ord a => [a] -> [a]
sort where
  f :: [a] -> (a, Int)
f [a]
xs = ([a] -> a
forall a. [a] -> a
head [a]
xs, [a] -> Int
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
nInteger -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<Integer
1 = [Char] -> [Integer]
forall a. HasCallStack => [Char] -> a
error [Char]
"integerFactorsTrialDivision: n should be at least 1"
  | Bool
otherwise = [Integer] -> Integer -> [Integer]
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 a -> a -> a
forall a. Integral a => a -> a -> a
mod a
k a
q of
        a
0 -> a
q a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> a -> [a]
go [a]
qqs (a -> a -> a
forall a. Integral a => a -> a -> a
div a
k a
q)
        a
_ -> [a] -> a -> [a]
sub [a]
qs a
k
      ps :: [a]
ps = (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (\a
p -> a
pa -> a -> a
forall a. Num a => a -> a -> a
*a
p a -> a -> Bool
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 = Integer -> [Bool]
forall t. (Num t, Bits t) => t -> [Bool]
bin Integer
k

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

  a :: Integer
a = Integer -> Integer -> Integer
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 Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
mod (Integer
xInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
rest) Integer
m
      else Integer
rest
    where 
      rest :: Integer
rest = Integer -> [Bool] -> Integer
go (Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
mod (Integer
xInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
x) Integer
m) [Bool]
bs 
      
--------------------------------------------------------------------------------
-- Prime testing
 
-- | 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 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
1 Bool -> Bool -> Bool
|| Integer
a Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1 = 
      [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error ([Char] -> Bool) -> [Char] -> Bool
forall a b. (a -> b) -> a -> b
$ [Char]
"millerRabinPrimalityTest: a out of range (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
a [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" for "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
n [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")" 
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2 = Bool
False
  | Integer -> Bool
forall a. Integral a => a -> Bool
even Integer
n = Bool
False
  | Integer
b0 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 Bool -> Bool -> Bool
|| Integer
b0 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
n' = Bool
True
  | Bool
otherwise = [Integer] -> Bool
iter ([Integer] -> [Integer]
forall a. [a] -> [a]
tail [Integer]
b)
  where
    n' :: Integer
n' = Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1
    (Integer
k,Integer
m) = Integer -> (Integer, Integer)
forall a. Integral a => a -> (a, a)
find2km Integer
n'
    b0 :: Integer
b0 = Integer -> Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a -> a
powMod Integer
n Integer
a Integer
m
    b :: [Integer]
b = Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take (Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
k) ([Integer] -> [Integer]) -> [Integer] -> [Integer]
forall a b. (a -> b) -> a -> b
$ (Integer -> Integer) -> Integer -> [Integer]
forall a. (a -> a) -> a -> [a]
iterate (Integer -> Integer -> Integer
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 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 = Bool
False
      | Integer
x Integer -> Integer -> Bool
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 :: a -> (a, a)
find2km a
n = a -> a -> (a, a)
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 t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
1 = (t
k,t
m)
    | Bool
otherwise = t -> t -> (t, t)
f (t
kt -> t -> t
forall a. Num a => a -> a -> a
+t
1) t
q
    where (t
q,t
r) = t -> t -> (t, t)
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' :: (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' = a -> b -> a -> a
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 t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
1 = a
x a -> a -> a
`mul` a
y
    | t
r t -> t -> Bool
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) = t -> t -> (t, t)
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 :: a -> a -> a -> a
mulMod a
a a
b a
c = (a
b a -> a -> a
forall a. Num a => a -> a -> a
* a
c) a -> a -> a
forall a. Integral a => a -> a -> a
`mod` a
a

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

{-# SPECIALIZE powMod :: Integer -> Integer -> Integer -> Integer #-}
powMod :: Integral a => a -> a -> a -> a
powMod :: a -> a -> a -> a
powMod a
m = (a -> a -> a) -> (a -> a) -> a -> a -> a
forall a b.
(Num a, Integral b) =>
(a -> a -> a) -> (a -> a) -> a -> b -> a
pow' (a -> a -> a -> a
forall a. Integral a => a -> a -> a -> a
mulMod a
m) (a -> a -> a
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 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2      = Bool
False
  | Integer -> Bool
forall a. Integral a => a -> Bool
even Integer
n     = (Integer
nInteger -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
==Integer
2)
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
1000   = [Integer] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Integer -> [Integer]
integerFactorsTrialDivision Integer
n) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
  | Bool
otherwise  = [Bool] -> Bool
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div Integer
log2n Integer
2) :: Int
    witnessList :: [Integer]
witnessList = Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take Int
nchecks [Integer]
pseudoRnds
    pseudoRnds :: [Integer]
pseudoRnds  = Integer
2 Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: [ Integer
a | Integer
a <- Integer -> [Integer]
integerRndSequence Integer
n , Integer
a Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
1 Bool -> Bool -> Bool
&& Integer
a Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< (Integer
nInteger -> Integer -> Integer
forall 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 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2      = Bool
False
  | Integer -> Bool
forall a. Integral a => a -> Bool
even Integer
n     = (Integer
nInteger -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
==Integer
2)
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
1000   = [Integer] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Integer -> [Integer]
integerFactorsTrialDivision Integer
n) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
  | Bool
otherwise  = [Bool] -> Bool
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div Integer
log2n Integer
2) :: Int
    witnessList :: [Integer]
witnessList = Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take Int
nchecks [Integer]
primes [Integer] -> [Integer] -> [Integer]
forall a. [a] -> [a] -> [a]
++ Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take Int
nchecks [Integer]
pseudoRnds
    pseudoRnds :: [Integer]
pseudoRnds  = [ Integer
a | Integer
a <- Integer -> [Integer]
integerRndSequence (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
3) , Integer
a Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
1 Bool -> Bool -> Bool
&& Integer
a Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) ]

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

{-
-- | Given an integer @n@, we return an infinite sequence of pseudo-random integers 
-- between @0..n-1@, generated using a crypographic hash function.
--
integerHashSequence :: Integer -> [Integer]
integerHashSequence = error "integerHashSequence: not implemented yet"
-}

-- | 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 = (Integer, Integer) -> StdGen -> [Integer]
forall a g. (Random a, RandomGen g) => (a, a) -> g -> [a]
randomRs (Integer
0,Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) StdGen
gen where
  gen :: StdGen
gen = Int -> StdGen
mkStdGen (Int -> StdGen) -> Int -> StdGen
forall a b. (a -> b) -> a -> b
$ Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
17 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer -> Integer
integerLog2 Integer
n)

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