-- |
-- Module:      Math.NumberTheory.ArithmeticFunctions.Standard
-- Copyright:   (c) 2016 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Textbook arithmetic functions.
--

{-# LANGUAGE ScopedTypeVariables #-}

module Math.NumberTheory.ArithmeticFunctions.Standard
  ( -- * List divisors
    divisors, divisorsA
  , divisorsList, divisorsListA
  , divisorsSmall, divisorsSmallA
  , divisorsTo, divisorsToA
    -- * Multiplicative functions
  , multiplicative
  , divisorCount, tau, tauA
  , sigma, sigmaA
  , totient, totientA
  , jordan, jordanA
  , ramanujan, ramanujanA
  , moebius, moebiusA, Moebius(..), runMoebius
  , liouville, liouvilleA
    -- * Additive functions
  , additive
  , smallOmega, smallOmegaA
  , bigOmega, bigOmegaA
    -- * Misc
  , carmichael, carmichaelA
  , expMangoldt, expMangoldtA
  , isNFree, isNFreeA, nFrees, nFreesBlock
  ) where

import Data.Coerce
import Data.Euclidean (GcdDomain(divide))
import Data.IntSet (IntSet)
import qualified Data.IntSet as IS
import Data.Maybe
import Data.Set (Set)
import qualified Data.Set as S
import Data.Semigroup

import Math.NumberTheory.ArithmeticFunctions.Class
import Math.NumberTheory.ArithmeticFunctions.Moebius
import Math.NumberTheory.ArithmeticFunctions.NFreedom (nFrees, nFreesBlock)
import Math.NumberTheory.Primes
import Math.NumberTheory.Utils.FromIntegral

import Numeric.Natural

-- | Create a multiplicative function from the function on prime's powers. See examples below.
multiplicative :: Num a => (Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative :: (Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative Prime n -> Word -> a
f = (Prime n -> Word -> Product a)
-> (Product a -> a) -> ArithmeticFunction n a
forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction ((a -> Product a
forall a. a -> Product a
Product (a -> Product a) -> (Word -> a) -> Word -> Product a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((Word -> a) -> Word -> Product a)
-> (Prime n -> Word -> a) -> Prime n -> Word -> Product a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prime n -> Word -> a
f) Product a -> a
forall a. Product a -> a
getProduct

-- | See 'divisorsA'.
divisors :: (UniqueFactorisation n, Ord n) => n -> Set n
divisors :: n -> Set n
divisors = ArithmeticFunction n (Set n) -> n -> Set n
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction n (Set n)
forall n. (Ord n, Num n) => ArithmeticFunction n (Set n)
divisorsA
{-# SPECIALIZE divisors :: Natural -> Set Natural #-}
{-# SPECIALIZE divisors :: Integer -> Set Integer #-}

-- | The set of all (positive) divisors of an argument.
divisorsA :: (Ord n, Num n) => ArithmeticFunction n (Set n)
divisorsA :: ArithmeticFunction n (Set n)
divisorsA = (Prime n -> Word -> SetProduct n)
-> (SetProduct n -> Set n) -> ArithmeticFunction n (Set n)
forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (\Prime n
p -> Set n -> SetProduct n
forall a. Set a -> SetProduct a
SetProduct (Set n -> SetProduct n) -> (Word -> Set n) -> Word -> SetProduct n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. n -> Word -> Set n
forall n. Num n => n -> Word -> Set n
divisorsHelper (Prime n -> n
forall a. Prime a -> a
unPrime Prime n
p)) (n -> Set n -> Set n
forall a. Ord a => a -> Set a -> Set a
S.insert n
1 (Set n -> Set n)
-> (SetProduct n -> Set n) -> SetProduct n -> Set n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SetProduct n -> Set n
forall a. SetProduct a -> Set a
getSetProduct)

divisorsHelper :: Num n => n -> Word -> Set n
divisorsHelper :: n -> Word -> Set n
divisorsHelper n
_ Word
0 = Set n
forall a. Set a
S.empty
divisorsHelper n
p Word
1 = n -> Set n
forall a. a -> Set a
S.singleton n
p
divisorsHelper n
p Word
a = [n] -> Set n
forall a. [a] -> Set a
S.fromDistinctAscList ([n] -> Set n) -> [n] -> Set n
forall a b. (a -> b) -> a -> b
$ n
p n -> [n] -> [n]
forall a. a -> [a] -> [a]
: n
p n -> n -> n
forall a. Num a => a -> a -> a
* n
p n -> [n] -> [n]
forall a. a -> [a] -> [a]
: (Int -> n) -> [Int] -> [n]
forall a b. (a -> b) -> [a] -> [b]
map (n
p n -> Int -> n
forall a b. (Num a, Integral b) => a -> b -> a
^) [Int
3 .. Word -> Int
wordToInt Word
a]
{-# INLINE divisorsHelper #-}

-- | See 'divisorsListA'.
divisorsList :: UniqueFactorisation n => n -> [n]
divisorsList :: n -> [n]
divisorsList = ArithmeticFunction n [n] -> n -> [n]
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction n [n]
forall n. Num n => ArithmeticFunction n [n]
divisorsListA

-- | The unsorted list of all (positive) divisors of an argument, produced in lazy fashion.
divisorsListA :: Num n => ArithmeticFunction n [n]
divisorsListA :: ArithmeticFunction n [n]
divisorsListA = (Prime n -> Word -> ListProduct n)
-> (ListProduct n -> [n]) -> ArithmeticFunction n [n]
forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (\Prime n
p -> [n] -> ListProduct n
forall a. [a] -> ListProduct a
ListProduct ([n] -> ListProduct n) -> (Word -> [n]) -> Word -> ListProduct n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. n -> Word -> [n]
forall n. Num n => n -> Word -> [n]
divisorsListHelper (Prime n -> n
forall a. Prime a -> a
unPrime Prime n
p)) ((n
1 n -> [n] -> [n]
forall a. a -> [a] -> [a]
:) ([n] -> [n]) -> (ListProduct n -> [n]) -> ListProduct n -> [n]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ListProduct n -> [n]
forall a. ListProduct a -> [a]
getListProduct)

divisorsListHelper :: Num n => n -> Word -> [n]
divisorsListHelper :: n -> Word -> [n]
divisorsListHelper n
_ Word
0 = []
divisorsListHelper n
p Word
1 = [n
p]
divisorsListHelper n
p Word
a = n
p n -> [n] -> [n]
forall a. a -> [a] -> [a]
: n
p n -> n -> n
forall a. Num a => a -> a -> a
* n
p n -> [n] -> [n]
forall a. a -> [a] -> [a]
: (Int -> n) -> [Int] -> [n]
forall a b. (a -> b) -> [a] -> [b]
map (n
p n -> Int -> n
forall a b. (Num a, Integral b) => a -> b -> a
^) [Int
3 .. Word -> Int
wordToInt Word
a]
{-# INLINE divisorsListHelper #-}

-- | See 'divisorsSmallA'.
divisorsSmall :: Int -> IntSet
divisorsSmall :: Int -> IntSet
divisorsSmall = ArithmeticFunction Int IntSet -> Int -> IntSet
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction Int IntSet
divisorsSmallA

-- | Same as 'divisors', but with better performance on cost of type restriction.
divisorsSmallA :: ArithmeticFunction Int IntSet
divisorsSmallA :: ArithmeticFunction Int IntSet
divisorsSmallA = (Prime Int -> Word -> IntSetProduct)
-> (IntSetProduct -> IntSet) -> ArithmeticFunction Int IntSet
forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (\Prime Int
p -> IntSet -> IntSetProduct
IntSetProduct (IntSet -> IntSetProduct)
-> (Word -> IntSet) -> Word -> IntSetProduct
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word -> IntSet
divisorsHelperSmall (Prime Int -> Int
forall a. Prime a -> a
unPrime Prime Int
p)) (Int -> IntSet -> IntSet
IS.insert Int
1 (IntSet -> IntSet)
-> (IntSetProduct -> IntSet) -> IntSetProduct -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntSetProduct -> IntSet
getIntSetProduct)

divisorsHelperSmall :: Int -> Word -> IntSet
divisorsHelperSmall :: Int -> Word -> IntSet
divisorsHelperSmall Int
_ Word
0 = IntSet
IS.empty
divisorsHelperSmall Int
p Word
1 = Int -> IntSet
IS.singleton Int
p
divisorsHelperSmall Int
p Word
a = [Int] -> IntSet
IS.fromDistinctAscList ([Int] -> IntSet) -> [Int] -> IntSet
forall a b. (a -> b) -> a -> b
$ Int
p Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
p Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int
p Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^) [Int
3 .. Word -> Int
wordToInt Word
a]
{-# INLINE divisorsHelperSmall #-}

-- | See 'divisorsToA'.
divisorsTo :: (UniqueFactorisation n, Integral n) => n -> n -> Set n
divisorsTo :: n -> n -> Set n
divisorsTo n
to = ArithmeticFunction n (Set n) -> n -> Set n
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction (n -> ArithmeticFunction n (Set n)
forall n.
(UniqueFactorisation n, Integral n) =>
n -> ArithmeticFunction n (Set n)
divisorsToA n
to)

-- | The set of all (positive) divisors up to an inclusive bound.
divisorsToA :: (UniqueFactorisation n, Integral n) => n -> ArithmeticFunction n (Set n)
divisorsToA :: n -> ArithmeticFunction n (Set n)
divisorsToA n
to = (Prime n -> Word -> BoundedSetProduct n)
-> (BoundedSetProduct n -> Set n) -> ArithmeticFunction n (Set n)
forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction Prime n -> Word -> BoundedSetProduct n
forall n. (Ord n, Num n) => Prime n -> Word -> BoundedSetProduct n
f BoundedSetProduct n -> Set n
unwrap
  where f :: Prime n -> Word -> BoundedSetProduct n
f Prime n
p Word
k = (n -> Set n) -> BoundedSetProduct n
forall a. (a -> Set a) -> BoundedSetProduct a
BoundedSetProduct (\n
bound -> n -> n -> Word -> Set n
forall n. (Ord n, Num n) => n -> n -> Word -> Set n
divisorsToHelper n
bound (Prime n -> n
forall a. Prime a -> a
unPrime Prime n
p) Word
k)
        unwrap :: BoundedSetProduct n -> Set n
unwrap (BoundedSetProduct n -> Set n
res) = if n
1 n -> n -> Bool
forall a. Ord a => a -> a -> Bool
<= n
to then n -> Set n -> Set n
forall a. Ord a => a -> Set a -> Set a
S.insert n
1 (n -> Set n
res n
to) else n -> Set n
res n
to

-- | Generate at most @a@ powers of @p@ up to an inclusive bound @b@.
divisorsToHelper :: (Ord n, Num n) => n -> n -> Word -> Set n
divisorsToHelper :: n -> n -> Word -> Set n
divisorsToHelper n
_ n
_ Word
0 = Set n
forall a. Set a
S.empty
divisorsToHelper n
b n
p Word
1 = if n
p n -> n -> Bool
forall a. Ord a => a -> a -> Bool
<= n
b then n -> Set n
forall a. a -> Set a
S.singleton n
p else Set n
forall a. Set a
S.empty
divisorsToHelper n
b n
p Word
a = [n] -> Set n
forall a. [a] -> Set a
S.fromDistinctAscList ([n] -> Set n) -> [n] -> Set n
forall a b. (a -> b) -> a -> b
$ Int -> [n] -> [n]
forall a. Int -> [a] -> [a]
take (Word -> Int
wordToInt Word
a) ([n] -> [n]) -> [n] -> [n]
forall a b. (a -> b) -> a -> b
$ (n -> Bool) -> [n] -> [n]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (n -> n -> Bool
forall a. Ord a => a -> a -> Bool
<=n
b) ([n] -> [n]) -> [n] -> [n]
forall a b. (a -> b) -> a -> b
$ (n -> n) -> n -> [n]
forall a. (a -> a) -> a -> [a]
iterate (n
pn -> n -> n
forall a. Num a => a -> a -> a
*) n
p
{-# INLINE divisorsToHelper #-}

-- | Synonym for 'tau'.
--
-- >>> map divisorCount [1..10]
-- [1,2,2,3,2,4,2,4,3,4]
divisorCount :: (UniqueFactorisation n, Num a) => n -> a
divisorCount :: n -> a
divisorCount = n -> a
forall n a. (UniqueFactorisation n, Num a) => n -> a
tau

-- | See 'tauA'.
tau :: (UniqueFactorisation n, Num a) => n -> a
tau :: n -> a
tau = ArithmeticFunction n a -> n -> a
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction n a
forall a n. Num a => ArithmeticFunction n a
tauA

-- | The number of (positive) divisors of an argument.
--
-- > tauA = multiplicative (\_ k -> k + 1)
tauA :: Num a => ArithmeticFunction n a
tauA :: ArithmeticFunction n a
tauA = (Prime n -> Word -> a) -> ArithmeticFunction n a
forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative ((Prime n -> Word -> a) -> ArithmeticFunction n a)
-> (Prime n -> Word -> a) -> ArithmeticFunction n a
forall a b. (a -> b) -> a -> b
$ (Word -> a) -> Prime n -> Word -> a
forall a b. a -> b -> a
const (Word -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> a) -> (Word -> Word) -> Word -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Word
forall a. Enum a => a -> a
succ)

-- | See 'sigmaA'.
sigma :: (UniqueFactorisation n, Integral n, Num a, GcdDomain a) => Word -> n -> a
sigma :: Word -> n -> a
sigma = ArithmeticFunction n a -> n -> a
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction (ArithmeticFunction n a -> n -> a)
-> (Word -> ArithmeticFunction n a) -> Word -> n -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> ArithmeticFunction n a
forall n a.
(Integral n, Num a, GcdDomain a) =>
Word -> ArithmeticFunction n a
sigmaA
{-# INLINABLE sigma #-}

-- | The sum of the @k@-th powers of (positive) divisors of an argument.
--
-- > sigmaA = multiplicative (\p k -> sum $ map (p ^) [0..k])
-- > sigmaA 0 = tauA
sigmaA :: (Integral n, Num a, GcdDomain a) => Word -> ArithmeticFunction n a
sigmaA :: Word -> ArithmeticFunction n a
sigmaA Word
0 = ArithmeticFunction n a
forall a n. Num a => ArithmeticFunction n a
tauA
sigmaA Word
1 = (Prime n -> Word -> a) -> ArithmeticFunction n a
forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative ((Prime n -> Word -> a) -> ArithmeticFunction n a)
-> (Prime n -> Word -> a) -> ArithmeticFunction n a
forall a b. (a -> b) -> a -> b
$ a -> Word -> a
forall a. (Num a, GcdDomain a) => a -> Word -> a
sigmaHelper (a -> Word -> a) -> (Prime n -> a) -> Prime n -> Word -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. n -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral' (n -> a) -> (Prime n -> n) -> Prime n -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prime n -> n
forall a. Prime a -> a
unPrime
sigmaA Word
a = (Prime n -> Word -> a) -> ArithmeticFunction n a
forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative ((Prime n -> Word -> a) -> ArithmeticFunction n a)
-> (Prime n -> Word -> a) -> ArithmeticFunction n a
forall a b. (a -> b) -> a -> b
$ a -> Word -> a
forall a. (Num a, GcdDomain a) => a -> Word -> a
sigmaHelper (a -> Word -> a) -> (Prime n -> a) -> Prime n -> Word -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ Word -> Int
wordToInt Word
a) (a -> a) -> (Prime n -> a) -> Prime n -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. n -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral' (n -> a) -> (Prime n -> n) -> Prime n -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prime n -> n
forall a. Prime a -> a
unPrime
{-# INLINABLE sigmaA #-}

sigmaHelper :: (Num a, GcdDomain a) => a -> Word -> a
sigmaHelper :: a -> Word -> a
sigmaHelper a
pa Word
1 = a
pa a -> a -> a
forall a. Num a => a -> a -> a
+ a
1
sigmaHelper a
pa Word
2 = a
pa a -> a -> a
forall a. Num a => a -> a -> a
* a
pa a -> a -> a
forall a. Num a => a -> a -> a
+ a
pa a -> a -> a
forall a. Num a => a -> a -> a
+ a
1
sigmaHelper a
pa Word
k = Maybe a -> a
forall a. HasCallStack => Maybe a -> a
fromJust ((a
pa a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ Word -> Int
wordToInt (Word
k Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1) a -> a -> a
forall a. Num a => a -> a -> a
- a
1) a -> a -> Maybe a
forall a. GcdDomain a => a -> a -> Maybe a
`divide` (a
pa a -> a -> a
forall a. Num a => a -> a -> a
- a
1))
{-# INLINE sigmaHelper #-}

-- | See 'totientA'.
totient :: UniqueFactorisation n => n -> n
totient :: n -> n
totient = ArithmeticFunction n n -> n -> n
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction n n
forall n. Num n => ArithmeticFunction n n
totientA
{-# INLINABLE totient #-}

-- | Calculates the totient of a positive number @n@, i.e.
--   the number of @k@ with @1 <= k <= n@ and @'gcd' n k == 1@,
--   in other words, the order of the group of units in @&#8484;/(n)@.
totientA :: Num n => ArithmeticFunction n n
totientA :: ArithmeticFunction n n
totientA = (Prime n -> Word -> n) -> ArithmeticFunction n n
forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative ((Prime n -> Word -> n) -> ArithmeticFunction n n)
-> (Prime n -> Word -> n) -> ArithmeticFunction n n
forall a b. (a -> b) -> a -> b
$ n -> Word -> n
forall n. Num n => n -> Word -> n
jordanHelper (n -> Word -> n) -> (Prime n -> n) -> Prime n -> Word -> n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prime n -> n
forall a. Prime a -> a
unPrime
{-# INLINABLE totientA #-}

-- | See 'jordanA'.
jordan :: UniqueFactorisation n => Word -> n -> n
jordan :: Word -> n -> n
jordan = ArithmeticFunction n n -> n -> n
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction (ArithmeticFunction n n -> n -> n)
-> (Word -> ArithmeticFunction n n) -> Word -> n -> n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> ArithmeticFunction n n
forall n. Num n => Word -> ArithmeticFunction n n
jordanA

-- | Calculates the k-th Jordan function of an argument.
--
-- > jordanA 1 = totientA
jordanA :: Num n => Word -> ArithmeticFunction n n
jordanA :: Word -> ArithmeticFunction n n
jordanA Word
0 = (Prime n -> Word -> n) -> ArithmeticFunction n n
forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative ((Prime n -> Word -> n) -> ArithmeticFunction n n)
-> (Prime n -> Word -> n) -> ArithmeticFunction n n
forall a b. (a -> b) -> a -> b
$ \Prime n
_ Word
_ -> n
0
jordanA Word
1 = ArithmeticFunction n n
forall n. Num n => ArithmeticFunction n n
totientA
jordanA Word
a = (Prime n -> Word -> n) -> ArithmeticFunction n n
forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative ((Prime n -> Word -> n) -> ArithmeticFunction n n)
-> (Prime n -> Word -> n) -> ArithmeticFunction n n
forall a b. (a -> b) -> a -> b
$ n -> Word -> n
forall n. Num n => n -> Word -> n
jordanHelper (n -> Word -> n) -> (Prime n -> n) -> Prime n -> Word -> n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (n -> Int -> n
forall a b. (Num a, Integral b) => a -> b -> a
^ Word -> Int
wordToInt Word
a) (n -> n) -> (Prime n -> n) -> Prime n -> n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prime n -> n
forall a. Prime a -> a
unPrime

jordanHelper :: Num n => n -> Word -> n
jordanHelper :: n -> Word -> n
jordanHelper n
pa Word
1 = n
pa n -> n -> n
forall a. Num a => a -> a -> a
- n
1
jordanHelper n
pa Word
2 = (n
pa n -> n -> n
forall a. Num a => a -> a -> a
- n
1) n -> n -> n
forall a. Num a => a -> a -> a
* n
pa
jordanHelper n
pa Word
k = (n
pa n -> n -> n
forall a. Num a => a -> a -> a
- n
1) n -> n -> n
forall a. Num a => a -> a -> a
* n
pa n -> Int -> n
forall a b. (Num a, Integral b) => a -> b -> a
^ Word -> Int
wordToInt (Word
k Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1)
{-# INLINE jordanHelper #-}

-- | See 'ramanujanA'.
ramanujan :: Integer -> Integer
ramanujan :: Integer -> Integer
ramanujan = ArithmeticFunction Integer Integer -> Integer -> Integer
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction Integer Integer
ramanujanA

-- | Calculates the <https://en.wikipedia.org/wiki/Ramanujan_tau_function Ramanujan tau function>
--   of a positive number @n@, using formulas given <http://www.numbertheory.org/php/tau.html here>
ramanujanA :: ArithmeticFunction Integer Integer
ramanujanA :: ArithmeticFunction Integer Integer
ramanujanA = (Prime Integer -> Word -> Integer)
-> ArithmeticFunction Integer Integer
forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative ((Prime Integer -> Word -> Integer)
 -> ArithmeticFunction Integer Integer)
-> (Prime Integer -> Word -> Integer)
-> ArithmeticFunction Integer Integer
forall a b. (a -> b) -> a -> b
$ Integer -> Word -> Integer
ramanujanHelper (Integer -> Word -> Integer)
-> (Prime Integer -> Integer) -> Prime Integer -> Word -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prime Integer -> Integer
forall a. Prime a -> a
unPrime

ramanujanHelper :: Integer -> Word -> Integer
ramanujanHelper :: Integer -> Word -> Integer
ramanujanHelper Integer
_ Word
0 = Integer
1
ramanujanHelper Integer
2 Word
1 = -Integer
24
ramanujanHelper Integer
p Word
1 = (Integer
65 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer -> Word -> Integer
forall a. (Num a, GcdDomain a) => a -> Word -> a
sigmaHelper (Integer
p Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
11 :: Int)) Word
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
691 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer -> Word -> Integer
forall a. (Num a, GcdDomain a) => a -> Word -> a
sigmaHelper (Integer
p Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
5 :: Int)) Word
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
691 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
252 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* [Integer] -> Integer
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Word -> Integer -> Integer
forall n a.
(UniqueFactorisation n, Integral n, Num a, GcdDomain a) =>
Word -> n -> a
sigma Word
5 Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Word -> Integer -> Integer
forall n a.
(UniqueFactorisation n, Integral n, Num a, GcdDomain a) =>
Word -> n -> a
sigma Word
5 (Integer
pInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
k) | Integer
k <- [Integer
1..(Integer
p Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`quot` Integer
2)]]) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`quot` Integer
756
ramanujanHelper Integer
p Word
k = [Integer] -> Integer
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Integer] -> Integer) -> [Integer] -> Integer
forall a b. (a -> b) -> a -> b
$ (Integer -> Integer -> Integer -> Integer)
-> [Integer] -> [Integer] -> [Integer] -> [Integer]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Integer
a Integer
b Integer
c -> Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
c) [Integer]
paPowers [Integer]
tpPowers [Integer]
binomials
  where pa :: Integer
pa = Integer
p Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
11 :: Int)
        tp :: Integer
tp = Integer -> Word -> Integer
ramanujanHelper Integer
p Word
1
        paPowers :: [Integer]
paPowers = (Integer -> Integer) -> Integer -> [Integer]
forall a. (a -> a) -> a -> [a]
iterate (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (-Integer
pa)) Integer
1
        binomials :: [Integer]
binomials = (Integer -> Integer -> Integer)
-> Integer -> [Integer] -> [Integer]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl (\Integer
acc Integer
j -> Integer
acc Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Integer
k' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
j) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Integer
k' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
j Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`quot` (Integer
k' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
j) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`quot` (Integer
j Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)) Integer
1 [Integer
0 .. Integer
k' Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`quot` Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1]
        k' :: Integer
k' = Word -> Integer
wordToInteger Word
k
        tpPowers :: [Integer]
tpPowers = [Integer] -> [Integer]
forall a. [a] -> [a]
reverse ([Integer] -> [Integer]) -> [Integer] -> [Integer]
forall a b. (a -> b) -> a -> b
$ Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take ([Integer] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
binomials) ([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. Num a => a -> a -> a
* Integer
tpInteger -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^(Int
2::Int)) (if Word -> Bool
forall a. Integral a => a -> Bool
even Word
k then Integer
1 else Integer
tp)
{-# INLINE ramanujanHelper #-}

-- | See 'moebiusA'.
moebius :: UniqueFactorisation n => n -> Moebius
moebius :: n -> Moebius
moebius = ArithmeticFunction n Moebius -> n -> Moebius
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction n Moebius
forall n. ArithmeticFunction n Moebius
moebiusA

-- | Calculates the Möbius function of an argument.
moebiusA :: ArithmeticFunction n Moebius
moebiusA :: ArithmeticFunction n Moebius
moebiusA = (Prime n -> Word -> Moebius)
-> (Moebius -> Moebius) -> ArithmeticFunction n Moebius
forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction ((Word -> Moebius) -> Prime n -> Word -> Moebius
forall a b. a -> b -> a
const Word -> Moebius
forall a. (Eq a, Num a) => a -> Moebius
f) Moebius -> Moebius
forall a. a -> a
id
  where
    f :: a -> Moebius
f a
1 = Moebius
MoebiusN
    f a
0 = Moebius
MoebiusP
    f a
_ = Moebius
MoebiusZ

-- | See 'liouvilleA'.
liouville :: (UniqueFactorisation n, Num a) => n -> a
liouville :: n -> a
liouville = ArithmeticFunction n a -> n -> a
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction n a
forall a n. Num a => ArithmeticFunction n a
liouvilleA

-- | Calculates the Liouville function of an argument.
liouvilleA :: Num a => ArithmeticFunction n a
liouvilleA :: ArithmeticFunction n a
liouvilleA = (Prime n -> Word -> Xor) -> (Xor -> a) -> ArithmeticFunction n a
forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction ((Word -> Xor) -> Prime n -> Word -> Xor
forall a b. a -> b -> a
const ((Word -> Xor) -> Prime n -> Word -> Xor)
-> (Word -> Xor) -> Prime n -> Word -> Xor
forall a b. (a -> b) -> a -> b
$ Bool -> Xor
Xor (Bool -> Xor) -> (Word -> Bool) -> Word -> Xor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Bool
forall a. Integral a => a -> Bool
odd) Xor -> a
forall a. Num a => Xor -> a
runXor

-- | See 'carmichaelA'.
carmichael :: (UniqueFactorisation n, Integral n) => n -> n
carmichael :: n -> n
carmichael = ArithmeticFunction n n -> n -> n
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction n n
forall n. Integral n => ArithmeticFunction n n
carmichaelA
{-# SPECIALIZE carmichael :: Int     -> Int #-}
{-# SPECIALIZE carmichael :: Word    -> Word #-}
{-# SPECIALIZE carmichael :: Integer -> Integer #-}
{-# SPECIALIZE carmichael :: Natural -> Natural #-}

-- | Calculates the Carmichael function for a positive integer, that is,
--   the (smallest) exponent of the group of units in @&#8484;/(n)@.
carmichaelA :: Integral n => ArithmeticFunction n n
carmichaelA :: ArithmeticFunction n n
carmichaelA = (Prime n -> Word -> LCM n)
-> (LCM n -> n) -> ArithmeticFunction n n
forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (\Prime n
p -> n -> LCM n
forall a. a -> LCM a
LCM (n -> LCM n) -> (Word -> n) -> Word -> LCM n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. n -> Word -> n
forall a. (Eq a, Num a) => a -> Word -> a
f (Prime n -> n
forall a. Prime a -> a
unPrime Prime n
p)) LCM n -> n
forall a. LCM a -> a
getLCM
  where
    f :: a -> Word -> a
f a
2 Word
1 = a
1
    f a
2 Word
2 = a
2
    f a
2 Word
k = a
2 a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ Word -> Int
wordToInt (Word
k Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
2)
    f a
p Word
1 = a
p a -> a -> a
forall a. Num a => a -> a -> a
- a
1
    f a
p Word
2 = (a
p a -> a -> a
forall a. Num a => a -> a -> a
- a
1) a -> a -> a
forall a. Num a => a -> a -> a
* a
p
    f a
p Word
k = (a
p a -> a -> a
forall a. Num a => a -> a -> a
- a
1) a -> a -> a
forall a. Num a => a -> a -> a
* a
p a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ Word -> Int
wordToInt (Word
k Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1)

-- | Create an additive function from the function on prime's powers. See examples below.
additive :: Num a => (Prime n -> Word -> a) -> ArithmeticFunction n a
additive :: (Prime n -> Word -> a) -> ArithmeticFunction n a
additive Prime n -> Word -> a
f = (Prime n -> Word -> Sum a)
-> (Sum a -> a) -> ArithmeticFunction n a
forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction ((a -> Sum a
forall a. a -> Sum a
Sum (a -> Sum a) -> (Word -> a) -> Word -> Sum a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((Word -> a) -> Word -> Sum a)
-> (Prime n -> Word -> a) -> Prime n -> Word -> Sum a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prime n -> Word -> a
f) Sum a -> a
forall a. Sum a -> a
getSum

-- | See 'smallOmegaA'.
smallOmega :: (UniqueFactorisation n, Num a) => n -> a
smallOmega :: n -> a
smallOmega = ArithmeticFunction n a -> n -> a
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction n a
forall a n. Num a => ArithmeticFunction n a
smallOmegaA

-- | Number of distinct prime factors.
--
-- > smallOmegaA = additive (\_ _ -> 1)
smallOmegaA :: Num a => ArithmeticFunction n a
smallOmegaA :: ArithmeticFunction n a
smallOmegaA = (Prime n -> Word -> a) -> ArithmeticFunction n a
forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
additive ((Prime n -> Word -> a) -> ArithmeticFunction n a)
-> (Prime n -> Word -> a) -> ArithmeticFunction n a
forall a b. (a -> b) -> a -> b
$ (Word -> a) -> Prime n -> Word -> a
forall a b. a -> b -> a
const ((Word -> a) -> Prime n -> Word -> a)
-> (Word -> a) -> Prime n -> Word -> a
forall a b. (a -> b) -> a -> b
$ a -> Word -> a
forall a b. a -> b -> a
const a
1

-- | See 'bigOmegaA'.
bigOmega :: UniqueFactorisation n => n -> Word
bigOmega :: n -> Word
bigOmega = ArithmeticFunction n Word -> n -> Word
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction n Word
forall n. ArithmeticFunction n Word
bigOmegaA

-- | Number of prime factors, counted with multiplicity.
--
-- > bigOmegaA = additive (\_ k -> k)
bigOmegaA :: ArithmeticFunction n Word
bigOmegaA :: ArithmeticFunction n Word
bigOmegaA = (Prime n -> Word -> Word) -> ArithmeticFunction n Word
forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
additive ((Prime n -> Word -> Word) -> ArithmeticFunction n Word)
-> (Prime n -> Word -> Word) -> ArithmeticFunction n Word
forall a b. (a -> b) -> a -> b
$ (Word -> Word) -> Prime n -> Word -> Word
forall a b. a -> b -> a
const Word -> Word
forall a. a -> a
id

-- | See 'expMangoldtA'.
expMangoldt :: UniqueFactorisation n => n -> n
expMangoldt :: n -> n
expMangoldt = ArithmeticFunction n n -> n -> n
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction n n
forall n. Num n => ArithmeticFunction n n
expMangoldtA

-- | The exponent of von Mangoldt function. Use @log expMangoldtA@ to recover von Mangoldt function itself.
expMangoldtA :: Num n => ArithmeticFunction n n
expMangoldtA :: ArithmeticFunction n n
expMangoldtA = (Prime n -> Word -> Mangoldt n)
-> (Mangoldt n -> n) -> ArithmeticFunction n n
forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (Mangoldt n -> Word -> Mangoldt n
forall a b. a -> b -> a
const (Mangoldt n -> Word -> Mangoldt n)
-> (Prime n -> Mangoldt n) -> Prime n -> Word -> Mangoldt n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. n -> Mangoldt n
forall a. a -> Mangoldt a
MangoldtOne (n -> Mangoldt n) -> (Prime n -> n) -> Prime n -> Mangoldt n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prime n -> n
forall a. Prime a -> a
unPrime) Mangoldt n -> n
forall a. Num a => Mangoldt a -> a
runMangoldt

data Mangoldt a
  = MangoldtZero
  | MangoldtOne a
  | MangoldtMany

runMangoldt :: Num a => Mangoldt a -> a
runMangoldt :: Mangoldt a -> a
runMangoldt Mangoldt a
m = case Mangoldt a
m of
  Mangoldt a
MangoldtZero  -> a
1
  MangoldtOne a
a -> a
a
  Mangoldt a
MangoldtMany  -> a
1

instance Semigroup (Mangoldt a) where
  Mangoldt a
MangoldtZero <> :: Mangoldt a -> Mangoldt a -> Mangoldt a
<> Mangoldt a
a = Mangoldt a
a
  Mangoldt a
a <> Mangoldt a
MangoldtZero = Mangoldt a
a
  Mangoldt a
_ <> Mangoldt a
_ = Mangoldt a
forall a. Mangoldt a
MangoldtMany

instance Monoid (Mangoldt a) where
  mempty :: Mangoldt a
mempty  = Mangoldt a
forall a. Mangoldt a
MangoldtZero
  mappend :: Mangoldt a -> Mangoldt a -> Mangoldt a
mappend = Mangoldt a -> Mangoldt a -> Mangoldt a
forall a. Semigroup a => a -> a -> a
(<>)

-- | See 'isNFreeA'.
isNFree :: UniqueFactorisation n => Word -> n -> Bool
isNFree :: Word -> n -> Bool
isNFree Word
n = ArithmeticFunction n Bool -> n -> Bool
forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction (Word -> ArithmeticFunction n Bool
forall n. Word -> ArithmeticFunction n Bool
isNFreeA Word
n)

-- | Check if an integer is @n@-free. An integer @x@ is @n@-free if in its
-- factorisation into prime factors, no factor has an exponent larger than or
-- equal to @n@.
isNFreeA :: Word -> ArithmeticFunction n Bool
isNFreeA :: Word -> ArithmeticFunction n Bool
isNFreeA Word
n = (Prime n -> Word -> All)
-> (All -> Bool) -> ArithmeticFunction n Bool
forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (\Prime n
_ Word
pow -> Bool -> All
All (Bool -> All) -> Bool -> All
forall a b. (a -> b) -> a -> b
$ Word
pow Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Word
n) All -> Bool
getAll

newtype LCM a = LCM { LCM a -> a
getLCM :: a }

instance Integral a => Semigroup (LCM a) where
  <> :: LCM a -> LCM a -> LCM a
(<>) = (a -> a -> a) -> LCM a -> LCM a -> LCM a
coerce (a -> a -> a
forall a. Integral a => a -> a -> a
lcm :: a -> a -> a)

instance Integral a => Monoid (LCM a) where
  mempty :: LCM a
mempty  = a -> LCM a
forall a. a -> LCM a
LCM a
1
  mappend :: LCM a -> LCM a -> LCM a
mappend = LCM a -> LCM a -> LCM a
forall a. Semigroup a => a -> a -> a
(<>)

newtype Xor = Xor { Xor -> Bool
_getXor :: Bool }

runXor :: Num a => Xor -> a
runXor :: Xor -> a
runXor Xor
m = case Xor
m of
  Xor Bool
False ->  a
1
  Xor Bool
True  -> -a
1

instance Semigroup Xor where
   <> :: Xor -> Xor -> Xor
(<>) = (Bool -> Bool -> Bool) -> Xor -> Xor -> Xor
coerce (Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
(/=) :: Bool -> Bool -> Bool)

instance Monoid Xor where
  mempty :: Xor
mempty  = Bool -> Xor
Xor Bool
False
  mappend :: Xor -> Xor -> Xor
mappend = Xor -> Xor -> Xor
forall a. Semigroup a => a -> a -> a
(<>)

newtype SetProduct a = SetProduct { SetProduct a -> Set a
getSetProduct :: Set a }

instance (Num a, Ord a) => Semigroup (SetProduct a) where
  SetProduct Set a
s1 <> :: SetProduct a -> SetProduct a -> SetProduct a
<> SetProduct Set a
s2 = Set a -> SetProduct a
forall a. Set a -> SetProduct a
SetProduct (Set a -> SetProduct a) -> Set a -> SetProduct a
forall a b. (a -> b) -> a -> b
$ Set a
s1 Set a -> Set a -> Set a
forall a. Semigroup a => a -> a -> a
<> Set a
s2 Set a -> Set a -> Set a
forall a. Semigroup a => a -> a -> a
<> (a -> Set a) -> Set a -> Set a
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\a
n -> (a -> a) -> Set a -> Set a
forall a b. (a -> b) -> Set a -> Set b
S.mapMonotonic (a -> a -> a
forall a. Num a => a -> a -> a
* a
n) Set a
s2) Set a
s1

instance (Num a, Ord a) => Monoid (SetProduct a) where
  mempty :: SetProduct a
mempty  = Set a -> SetProduct a
forall a. Set a -> SetProduct a
SetProduct Set a
forall a. Monoid a => a
mempty
  mappend :: SetProduct a -> SetProduct a -> SetProduct a
mappend = SetProduct a -> SetProduct a -> SetProduct a
forall a. Semigroup a => a -> a -> a
(<>)

newtype ListProduct a = ListProduct { ListProduct a -> [a]
getListProduct :: [a] }

instance Num a => Semigroup (ListProduct a) where
  ListProduct [a]
s1 <> :: ListProduct a -> ListProduct a -> ListProduct a
<> ListProduct [a]
s2 = [a] -> ListProduct a
forall a. [a] -> ListProduct a
ListProduct ([a] -> ListProduct a) -> [a] -> ListProduct a
forall a b. (a -> b) -> a -> b
$ [a]
s1 [a] -> [a] -> [a]
forall a. Semigroup a => a -> a -> a
<> [a]
s2 [a] -> [a] -> [a]
forall a. Semigroup a => a -> a -> a
<> (a -> [a]) -> [a] -> [a]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\a
n -> (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a -> a -> a
forall a. Num a => a -> a -> a
* a
n) [a]
s2) [a]
s1

instance Num a => Monoid (ListProduct a) where
  mempty :: ListProduct a
mempty  = [a] -> ListProduct a
forall a. [a] -> ListProduct a
ListProduct [a]
forall a. Monoid a => a
mempty
  mappend :: ListProduct a -> ListProduct a -> ListProduct a
mappend = ListProduct a -> ListProduct a -> ListProduct a
forall a. Semigroup a => a -> a -> a
(<>)

-- Represent as a Reader monad
newtype BoundedSetProduct a = BoundedSetProduct { BoundedSetProduct a -> a -> Set a
_getBoundedSetProduct :: a -> Set a }

takeWhileLE :: Ord a => a -> Set a -> Set a
takeWhileLE :: a -> Set a -> Set a
takeWhileLE a
b Set a
xs = if Bool
m then a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
S.insert a
b Set a
ls else Set a
ls
  where (Set a
ls, Bool
m, Set a
_) = a -> Set a -> (Set a, Bool, Set a)
forall a. Ord a => a -> Set a -> (Set a, Bool, Set a)
S.splitMember a
b Set a
xs

instance (Ord a, Num a) => Semigroup (BoundedSetProduct a) where
  BoundedSetProduct a -> Set a
f1 <> :: BoundedSetProduct a -> BoundedSetProduct a -> BoundedSetProduct a
<> BoundedSetProduct a -> Set a
f2 = (a -> Set a) -> BoundedSetProduct a
forall a. (a -> Set a) -> BoundedSetProduct a
BoundedSetProduct a -> Set a
f
    where f :: a -> Set a
f a
b = Set a
s1 Set a -> Set a -> Set a
forall a. Semigroup a => a -> a -> a
<> Set a
s2 Set a -> Set a -> Set a
forall a. Semigroup a => a -> a -> a
<> (a -> Set a) -> Set a -> Set a
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\a
n -> a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
takeWhileLE a
b (Set a -> Set a) -> Set a -> Set a
forall a b. (a -> b) -> a -> b
$ (a -> a) -> Set a -> Set a
forall a b. (a -> b) -> Set a -> Set b
S.mapMonotonic (a -> a -> a
forall a. Num a => a -> a -> a
* a
n) Set a
s2) Set a
s1
            where s1 :: Set a
s1 = a -> Set a
f1 a
b
                  s2 :: Set a
s2 = a -> Set a
f2 a
b

instance (Ord a, Num a) => Monoid (BoundedSetProduct a) where
  mempty :: BoundedSetProduct a
mempty = (a -> Set a) -> BoundedSetProduct a
forall a. (a -> Set a) -> BoundedSetProduct a
BoundedSetProduct a -> Set a
forall a. Monoid a => a
mempty
  mappend :: BoundedSetProduct a -> BoundedSetProduct a -> BoundedSetProduct a
mappend = BoundedSetProduct a -> BoundedSetProduct a -> BoundedSetProduct a
forall a. Semigroup a => a -> a -> a
(<>)

newtype IntSetProduct = IntSetProduct { IntSetProduct -> IntSet
getIntSetProduct :: IntSet }

instance Semigroup IntSetProduct where
  IntSetProduct IntSet
s1 <> :: IntSetProduct -> IntSetProduct -> IntSetProduct
<> IntSetProduct IntSet
s2 = IntSet -> IntSetProduct
IntSetProduct (IntSet -> IntSetProduct) -> IntSet -> IntSetProduct
forall a b. (a -> b) -> a -> b
$ [IntSet] -> IntSet
forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IS.unions ([IntSet] -> IntSet) -> [IntSet] -> IntSet
forall a b. (a -> b) -> a -> b
$ IntSet
s1 IntSet -> [IntSet] -> [IntSet]
forall a. a -> [a] -> [a]
: IntSet
s2 IntSet -> [IntSet] -> [IntSet]
forall a. a -> [a] -> [a]
: (Int -> IntSet) -> [Int] -> [IntSet]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
n -> (Int -> Int) -> IntSet -> IntSet
IS.map (Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) IntSet
s2) (IntSet -> [Int]
IS.toAscList IntSet
s1)

instance Monoid IntSetProduct where
  mempty :: IntSetProduct
mempty  = IntSet -> IntSetProduct
IntSetProduct IntSet
forall a. Monoid a => a
mempty
  mappend :: IntSetProduct -> IntSetProduct -> IntSetProduct
mappend = IntSetProduct -> IntSetProduct -> IntSetProduct
forall a. Semigroup a => a -> a -> a
(<>)