-- |
-- Module:      Math.NumberTheory.Roots.General
-- Copyright:   (c) 2011 Daniel Fischer, 2016-2020 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Daniel Fischer <daniel.is.fischer@googlemail.com>
--
-- Calculating integer roots and determining perfect powers.
-- The algorithms are moderately efficient.
--

{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE CPP           #-}
{-# LANGUAGE MagicHash     #-}
{-# LANGUAGE ViewPatterns  #-}

module Math.NumberTheory.Roots.General
    ( integerRoot
    , exactRoot
    , isKthPower
    , isPerfectPower
    , highestPower
    ) where

#include "MachDeps.h"

import Data.Bits (countTrailingZeros, shiftL, shiftR)
import Data.List (foldl', sortBy)
import Data.Maybe (isJust)
import GHC.Exts (Int(..), Word(..), word2Int#, int2Double#, double2Int#, isTrue#, Ptr(..), indexWord16OffAddr#, (/##), (**##))
#if MIN_VERSION_base(4,16,0)
import GHC.Exts (word16ToWord#)
#endif
#ifdef WORDS_BIGENDIAN
import GHC.Exts (byteSwap16#)
#endif
import Numeric.Natural (Natural)

#ifdef MIN_VERSION_integer_gmp
import GHC.Exts (int2Word#, quotInt#, (<#), (*#), (-#), (+#))
import GHC.Integer.GMP.Internals (Integer(..), shiftLInteger, shiftRInteger)
import GHC.Integer.Logarithms (integerLog2#)
#define IS S#
#else
import GHC.Exts (plusWord#, minusWord#, timesWord#, quotWord#, ltWord#)
import GHC.Num.Integer (Integer(..), integerLog2#, integerShiftR#, integerShiftL#)
#endif

import qualified Math.NumberTheory.Roots.Squares as P2
import qualified Math.NumberTheory.Roots.Cubes as P3
import qualified Math.NumberTheory.Roots.Fourth as P4
import Math.NumberTheory.Primes.Small
import Math.NumberTheory.Utils.FromIntegral (wordToInt)

-- | For a positive power \( k \) and
-- a given \( n \)
-- return the integer \( k \)-th root \( \lfloor \sqrt[k]{n} \rfloor \).
-- Throw an error if \( k \le 0 \) or if \( n \le 0 \) and \( k \) is even.
--
-- >>> integerRoot 6 65
-- 2
-- >>> integerRoot 5 243
-- 3
-- >>> integerRoot 4 624
-- 4
-- >>> integerRoot 3 (-124)
-- -5
-- >>> integerRoot 1 5
-- 5
{-# SPECIALISE integerRoot :: Int -> Int -> Int,
                              Int -> Word -> Word,
                              Int -> Integer -> Integer,
                              Int -> Natural -> Natural,
                              Word -> Int -> Int,
                              Word -> Word -> Word,
                              Word -> Integer -> Integer,
                              Word -> Natural -> Natural,
                              Integer -> Integer -> Integer,
                              Natural -> Natural -> Natural
  #-}
integerRoot :: (Integral a, Integral b) => b -> a -> a
integerRoot :: b -> a -> a
integerRoot b
1 a
n         = a
n
integerRoot b
2 a
n         = a -> a
forall a. Integral a => a -> a
P2.integerSquareRoot a
n
integerRoot b
3 a
n         = a -> a
forall a. Integral a => a -> a
P3.integerCubeRoot a
n
integerRoot b
4 a
n         = a -> a
forall a. Integral a => a -> a
P4.integerFourthRoot a
n
integerRoot b
k a
n
  | b
k b -> b -> Bool
forall a. Ord a => a -> a -> Bool
< b
1             = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"integerRoot: negative exponent or exponent 0"
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0 Bool -> Bool -> Bool
&& b -> Bool
forall a. Integral a => a -> Bool
even b
k   = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"integerRoot: negative radicand for even exponent"
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0             =
    let r :: a
r = a -> a
forall a. Num a => a -> a
negate (a -> a) -> (Integer -> a) -> Integer -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> a
forall a. Num a => Integer -> a
fromInteger (Integer -> a) -> (Integer -> Integer) -> Integer -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Integer -> Integer
forall a b. (Integral a, Integral b) => b -> a -> a
integerRoot b
k (Integer -> Integer) -> (Integer -> Integer) -> Integer -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Integer
forall a. Num a => a -> a
negate (Integer -> a) -> Integer -> a
forall a b. (a -> b) -> a -> b
$ a -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n
    in if a
ra -> b -> a
forall a b. (Num a, Integral b) => a -> b -> a
^b
k a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
n then a
r else (a
ra -> a -> a
forall a. Num a => a -> a -> a
-a
1)
  | a
n a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0            = a
0
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
31            = a
1
  | Bool
kTooLarge         = a
1
  | Bool
otherwise         = Integer -> a
forall a. Num a => Integer -> a
fromInteger (Integer -> a) -> Integer -> a
forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> Integer -> Integer
newtonK (b -> Integer
forall a. Integral a => a -> Integer
toInteger b
k) (a -> Integer
forall a. Integral a => a -> Integer
toInteger a
n) Integer
a
    where
      a :: Integer
a  = Int -> Integer -> Integer
appKthRoot (b -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral b
k) (a -> Integer
forall a. Integral a => a -> Integer
toInteger a
n)
      kTooLarge :: Bool
kTooLarge = (b -> Integer
forall a. Integral a => a -> Integer
toInteger b
k Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= a -> Integer
forall a. Integral a => a -> Integer
toInteger (b -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral b
k a -> a -> a
forall a. a -> a -> a
`asTypeOf` a
n))    -- k doesn't fit in n's type
                  Bool -> Bool -> Bool
|| (b -> Integer
forall a. Integral a => a -> Integer
toInteger b
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int))  -- 2^k doesn't fit in Integer
#ifdef MIN_VERSION_integer_gmp
                  Bool -> Bool -> Bool
|| (Int# -> Int
I# (Integer -> Int#
integerLog2# (a -> Integer
forall a. Integral a => a -> Integer
toInteger a
n)) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< b -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral b
k) -- n < 2^k
#else
                  || (W# (integerLog2# (toInteger n)) < fromIntegral k) -- n < 2^k
#endif

-- | For a positive exponent \( k \)
-- calculate the exact integer \( k \)-th root of the second argument if it exists,
-- otherwise return 'Nothing'.
--
-- >>> map (uncurry exactRoot) [(6, 65), (5, 243), (4, 624), (3, -124), (1, 5)]
-- [Nothing,Just 3,Nothing,Nothing,Just 5]
exactRoot :: (Integral a, Integral b) => b -> a -> Maybe a
exactRoot :: b -> a -> Maybe a
exactRoot b
1 a
n = a -> Maybe a
forall a. a -> Maybe a
Just a
n
exactRoot b
2 a
n = a -> Maybe a
forall a. Integral a => a -> Maybe a
P2.exactSquareRoot a
n
exactRoot b
3 a
n = a -> Maybe a
forall a. Integral a => a -> Maybe a
P3.exactCubeRoot a
n
exactRoot b
4 a
n = a -> Maybe a
forall a. Integral a => a -> Maybe a
P4.exactFourthRoot a
n
exactRoot b
k a
n
  | a
n a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1          = a -> Maybe a
forall a. a -> Maybe a
Just a
1
  | b
k b -> b -> Bool
forall a. Ord a => a -> a -> Bool
< b
1           = Maybe a
forall a. Maybe a
Nothing
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0 Bool -> Bool -> Bool
&& b -> Bool
forall a. Integral a => a -> Bool
even b
k = Maybe a
forall a. Maybe a
Nothing
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0           = let m :: a
m = a -> a
forall a. Num a => a -> a
negate a
n in
                      if a
m a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0
                        then (Integer -> a) -> Maybe Integer -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Integer -> a
forall a. Num a => Integer -> a
fromInteger (Integer -> a) -> (Integer -> Integer) -> Integer -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Integer
forall a. Num a => a -> a
negate) (b -> Integer -> Maybe Integer
forall a b. (Integral a, Integral b) => b -> a -> Maybe a
exactRoot b
k (Integer -> Integer
forall a. Num a => a -> a
negate (a -> Integer
forall a. Integral a => a -> Integer
toInteger a
n)))
                        else (a -> a) -> Maybe a -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. Num a => a -> a
negate (b -> a -> Maybe a
forall a b. (Integral a, Integral b) => b -> a -> Maybe a
exactRoot b
k a
m)
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
2           = a -> Maybe a
forall a. a -> Maybe a
Just a
n
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
31          = Maybe a
forall a. Maybe a
Nothing
  | Bool
kTooLarge       = Maybe a
forall a. Maybe a
Nothing
  | Bool
otherwise       = case b
k b -> b -> b
forall a. Integral a => a -> a -> a
`rem` b
12 of
                        b
0 | Bool
c4 Bool -> Bool -> Bool
&& Bool
c3 Bool -> Bool -> Bool
&& Bool
ok -> a -> Maybe a
forall a. a -> Maybe a
Just a
r
                          | Bool
otherwise -> Maybe a
forall a. Maybe a
Nothing
                        b
2 | Bool
c2 Bool -> Bool -> Bool
&& Bool
ok -> a -> Maybe a
forall a. a -> Maybe a
Just a
r
                          | Bool
otherwise -> Maybe a
forall a. Maybe a
Nothing
                        b
3 | Bool
c3 Bool -> Bool -> Bool
&& Bool
ok -> a -> Maybe a
forall a. a -> Maybe a
Just a
r
                          | Bool
otherwise -> Maybe a
forall a. Maybe a
Nothing
                        b
4 | Bool
c4 Bool -> Bool -> Bool
&& Bool
ok -> a -> Maybe a
forall a. a -> Maybe a
Just a
r
                          | Bool
otherwise -> Maybe a
forall a. Maybe a
Nothing
                        b
6 | Bool
c3 Bool -> Bool -> Bool
&& Bool
c2 Bool -> Bool -> Bool
&& Bool
ok -> a -> Maybe a
forall a. a -> Maybe a
Just a
r
                          | Bool
otherwise -> Maybe a
forall a. Maybe a
Nothing
                        b
8 | Bool
c4 Bool -> Bool -> Bool
&& Bool
ok -> a -> Maybe a
forall a. a -> Maybe a
Just a
r
                          | Bool
otherwise -> Maybe a
forall a. Maybe a
Nothing
                        b
9 | Bool
c3 Bool -> Bool -> Bool
&& Bool
ok -> a -> Maybe a
forall a. a -> Maybe a
Just a
r
                          | Bool
otherwise -> Maybe a
forall a. Maybe a
Nothing
                        b
10 | Bool
c2 Bool -> Bool -> Bool
&& Bool
ok -> a -> Maybe a
forall a. a -> Maybe a
Just a
r
                           | Bool
otherwise -> Maybe a
forall a. Maybe a
Nothing
                        b
_ | Bool
ok -> a -> Maybe a
forall a. a -> Maybe a
Just a
r
                          | Bool
otherwise -> Maybe a
forall a. Maybe a
Nothing

    where
      k' :: Int
      k' :: Int
k' = b -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral b
k
      r :: a
r  = Int -> a -> a
forall a b. (Integral a, Integral b) => b -> a -> a
integerRoot Int
k' a
n
      c2 :: Bool
c2 = a -> Bool
forall a. Integral a => a -> Bool
P2.isPossibleSquare a
n
      c3 :: Bool
c3 = a -> Bool
forall a. Integral a => a -> Bool
P3.isPossibleCube a
n
      c4 :: Bool
c4 = a -> Bool
forall a. Integral a => a -> Bool
P4.isPossibleFourthPower a
n
      ok :: Bool
ok = a
ra -> b -> a
forall a b. (Num a, Integral b) => a -> b -> a
^b
k a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
n
      kTooLarge :: Bool
kTooLarge = (b -> Integer
forall a. Integral a => a -> Integer
toInteger b
k Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= a -> Integer
forall a. Integral a => a -> Integer
toInteger (b -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral b
k a -> a -> a
forall a. a -> a -> a
`asTypeOf` a
n))    -- k doesn't fit in n's type
                  Bool -> Bool -> Bool
|| (b -> Integer
forall a. Integral a => a -> Integer
toInteger b
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int))  -- 2^k doesn't fit in Integer
#ifdef MIN_VERSION_integer_gmp
                  Bool -> Bool -> Bool
|| (Int# -> Int
I# (Integer -> Int#
integerLog2# (a -> Integer
forall a. Integral a => a -> Integer
toInteger a
n)) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< b -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral b
k) -- n < 2^k
#else
                  || (W# (integerLog2# (toInteger n)) < fromIntegral k) -- n < 2^k
#endif

-- | For a positive exponent \( k \) test whether the second argument
-- is a perfect \( k \)-th power.
--
-- >>> map (uncurry isKthPower) [(6, 65), (5, 243), (4, 624), (3, -124), (1, 5)]
-- [False,True,False,False,True]
isKthPower :: (Integral a, Integral b) => b -> a -> Bool
isKthPower :: b -> a -> Bool
isKthPower b
k a
n = Maybe a -> Bool
forall a. Maybe a -> Bool
isJust (b -> a -> Maybe a
forall a b. (Integral a, Integral b) => b -> a -> Maybe a
exactRoot b
k a
n)

-- | Test whether the argument is a non-trivial perfect power
-- (e. g., square, cube, etc.).
--
-- >>> map isPerfectPower [0..10]
-- [True,True,False,False,True,False,False,False,True,True,False]
-- >>> map isPerfectPower [-10..0]
-- [False,False,True,False,False,False,False,False,False,True,True]
isPerfectPower :: Integral a => a -> Bool
isPerfectPower :: a -> Bool
isPerfectPower a
n
  | a
n a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 Bool -> Bool -> Bool
|| a
n a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1    = Bool
True
  | Bool
otherwise           = Word
k Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
> Word
1
    where
      (a
_,Word
k) = a -> (a, Word)
forall a. Integral a => a -> (a, Word)
highestPower a
n

-- | For \( n \not\in \{ -1, 0, 1 \} \)
-- find the largest exponent \( k \) for which
-- an exact integer \( k \)-th root \( r \) exists.
-- Return \( (r, k) \).
--
-- For \( n \in \{ -1, 0, 1 \} \) arbitrarily large exponents exist;
-- by arbitrary convention 'highestPower' returns \( (n, 3) \).
--
-- >>> map highestPower [0..10]
-- [(0,3),(1,3),(2,1),(3,1),(2,2),(5,1),(6,1),(7,1),(2,3),(3,2),(10,1)]
-- >>> map highestPower [-10..0]
-- [(-10,1),(-9,1),(-2,3),(-7,1),(-6,1),(-5,1),(-4,1),(-3,1),(-2,1),(-1,3),(0,3)]
highestPower :: Integral a => a -> (a, Word)
highestPower :: a -> (a, Word)
highestPower a
n'
  | Integer -> Integer
forall a. Num a => a -> a
abs Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
1  = (a
n', Word
3)
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0       = case Integer -> (Integer, Word)
integerHighPower (Integer -> Integer
forall a. Num a => a -> a
negate Integer
n) of
                    (Integer
r,Word
e) -> case Word -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros Word
e of
                               Int
k -> (a -> a
forall a. Num a => a -> a
negate (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$ Integer -> a
forall a. Num a => Integer -> a
fromInteger (Int -> Integer -> Integer
sqr Int
k Integer
r), Word
e Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`shiftR` Int
k)
  | Bool
otherwise   = case Integer -> (Integer, Word)
integerHighPower Integer
n of
                    (Integer
r,Word
e) -> (Integer -> a
forall a. Num a => Integer -> a
fromInteger Integer
r, Word
e)
    where
      n :: Integer
      n :: Integer
n = a -> Integer
forall a. Integral a => a -> Integer
toInteger a
n'

      sqr :: Int -> Integer -> Integer
      sqr :: Int -> Integer -> Integer
sqr Int
0 Integer
m = Integer
m
      sqr Int
k Integer
m = Int -> Integer -> Integer
sqr (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (Integer
mInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
m)

------------------------------------------------------------------------------------------
--                                  Auxiliary functions                                 --
------------------------------------------------------------------------------------------

newtonK :: Integer -> Integer -> Integer -> Integer
newtonK :: Integer -> Integer -> Integer -> Integer
newtonK Integer
k Integer
n Integer
a = Integer -> Integer
go (Integer -> Integer
step Integer
a)
  where
    step :: Integer -> Integer
step Integer
m = ((Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
m Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`quot` Integer
m Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ (Integer
k 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
    go :: Integer -> Integer
go Integer
m
      | Integer
l Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
m     = Integer -> Integer
go Integer
l
      | Bool
otherwise = Integer
m
        where
          l :: Integer
l = Integer -> Integer
step Integer
m

-- find an approximation to the k-th root
-- here, k > 4 and n > 31
appKthRoot :: Int -> Integer -> Integer
appKthRoot :: Int -> Integer -> Integer
appKthRoot (I# Int#
k#) (IS Int#
n#) = Int# -> Integer
IS (Double# -> Int#
double2Int# (Int# -> Double#
int2Double# Int#
n# Double# -> Double# -> Double#
**## (Double#
1.0## Double# -> Double# -> Double#
/## Int# -> Double#
int2Double# Int#
k#)))
#ifdef MIN_VERSION_integer_gmp
appKthRoot k :: Int
k@(I# Int#
k#) Integer
n
  | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
256 = Integer
1 Integer -> Int# -> Integer
`shiftLInteger` (Integer -> Int#
integerLog2# Integer
n Int# -> Int# -> Int#
`quotInt#` Int#
k# Int# -> Int# -> Int#
+# Int#
1#)
  | Bool
otherwise =
    case Integer -> Int#
integerLog2# Integer
n of
      Int#
l# -> case Int#
l# Int# -> Int# -> Int#
`quotInt#` Int#
k# of
              Int#
0# -> Integer
1
              Int#
1# -> Integer
3
              Int#
2# -> Integer
5
              Int#
3# -> Integer
11
              Int#
h# | Int# -> Bool
isTrue# (Int#
h# Int# -> Int# -> Int#
<# Int#
500#) ->
                   Double -> Integer
forall a b. (RealFrac a, Integral b) => a -> b
floor (Int -> Double -> Double
forall a. RealFloat a => Int -> a -> a
scaleFloat (Int# -> Int
I# Int#
h#)
                          (Integer -> Double
forall a. Num a => Integer -> a
fromInteger (Integer
n Integer -> Int# -> Integer
`shiftRInteger` (Int#
h# Int# -> Int# -> Int#
*# Int#
k#)) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (Double
1Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) :: Double))
                 | Bool
otherwise ->
                   Double -> Integer
forall a b. (RealFrac a, Integral b) => a -> b
floor (Int -> Double -> Double
forall a. RealFloat a => Int -> a -> a
scaleFloat Int
400 (Integer -> Double
forall a. Num a => Integer -> a
fromInteger (Integer
n Integer -> Int# -> Integer
`shiftRInteger` (Int#
h# Int# -> Int# -> Int#
*# Int#
k#)) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (Double
1Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) :: Double))
                          Integer -> Int# -> Integer
`shiftLInteger` (Int#
h# Int# -> Int# -> Int#
-# Int#
400#)
#else
appKthRoot k@(fromIntegral -> W# k#) n
  | k >= 256 = 1 `integerShiftL#` (integerLog2# n `quotWord#` k# `plusWord#` 1##)
  | otherwise =
    case integerLog2# n of
      l# -> case l# `quotWord#` k# of
              0## -> 1
              1## -> 3
              2## -> 5
              3## -> 11
              h# | isTrue# (h# `ltWord#` 500##) ->
                   floor (scaleFloat (I# (word2Int# h#))
                          (fromInteger (n `integerShiftR#` (h# `timesWord#` k#)) ** (1/fromIntegral k) :: Double))
                 | otherwise ->
                   floor (scaleFloat 400 (fromInteger (n `integerShiftR#` (h# `timesWord#` k#)) ** (1/fromIntegral k) :: Double))
                          `integerShiftL#` (h# `minusWord#` 400##)
#endif

-- assumption: argument is > 1
integerHighPower :: Integer -> (Integer, Word)
integerHighPower :: Integer -> (Integer, Word)
integerHighPower Integer
n
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
4       = (Integer
n,Word
1)
  | Bool
otherwise   = case Integer -> Integer -> (Word, Integer)
splitOff Integer
2 Integer
n of
                    (Word
e2,Integer
m) | Integer
m Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1     -> (Integer
2,Word
e2)
                           | Bool
otherwise  -> Word
-> [(Integer, Word)]
-> Integer
-> Integer
-> [Integer]
-> (Integer, Word)
findHighPower Word
e2 (if Word
e2 Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
0 then [] else [(Integer
2,Word
e2)]) Integer
m Integer
r [Integer]
smallOddPrimes
                             where
                               r :: Integer
r = Integer -> Integer
forall a. Integral a => a -> a
P2.integerSquareRoot Integer
m

findHighPower :: Word -> [(Integer, Word)] -> Integer -> Integer -> [Integer] -> (Integer, Word)
findHighPower :: Word
-> [(Integer, Word)]
-> Integer
-> Integer
-> [Integer]
-> (Integer, Word)
findHighPower Word
1 [(Integer, Word)]
pws Integer
m Integer
_ [Integer]
_ = ((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
m [Integer
pInteger -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^Word
e | (Integer
p,Word
e) <- [(Integer, Word)]
pws], Word
1)
findHighPower Word
e [(Integer, Word)]
pws Integer
1 Integer
_ [Integer]
_ = ((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 -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^(Word
ex Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
e) | (Integer
p,Word
ex) <- [(Integer, Word)]
pws], Word
e)
findHighPower Word
e [(Integer, Word)]
pws Integer
m Integer
s (Integer
p:[Integer]
ps)
  | Integer
s Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
p       = Word
-> [(Integer, Word)]
-> Integer
-> Integer
-> [Integer]
-> (Integer, Word)
findHighPower Word
1 [(Integer, Word)]
pws Integer
m Integer
s []
  | Bool
otherwise   =
    case Integer -> Integer -> (Word, Integer)
splitOff Integer
p Integer
m of
      (Word
0,Integer
_) -> Word
-> [(Integer, Word)]
-> Integer
-> Integer
-> [Integer]
-> (Integer, Word)
findHighPower Word
e [(Integer, Word)]
pws Integer
m Integer
s [Integer]
ps
      (Word
k,Integer
r) -> Word
-> [(Integer, Word)]
-> Integer
-> Integer
-> [Integer]
-> (Integer, Word)
findHighPower (Word -> Word -> Word
forall a. Integral a => a -> a -> a
gcd Word
k Word
e) ((Integer
p,Word
k)(Integer, Word) -> [(Integer, Word)] -> [(Integer, Word)]
forall a. a -> [a] -> [a]
:[(Integer, Word)]
pws) Integer
r (Integer -> Integer
forall a. Integral a => a -> a
P2.integerSquareRoot Integer
r) [Integer]
ps
findHighPower Word
e [(Integer, Word)]
pws Integer
m Integer
_ [] = Word -> [(Integer, Word)] -> Integer -> (Integer, Word)
finishPower Word
e [(Integer, Word)]
pws Integer
m

splitOff :: Integer -> Integer -> (Word, Integer)
splitOff :: Integer -> Integer -> (Word, Integer)
splitOff !Integer
_ Integer
0 = (Word
0, Integer
0) -- prevent infinite loop
splitOff Integer
p Integer
n = Word -> Integer -> (Word, Integer)
forall a. Num a => a -> Integer -> (a, Integer)
go Word
0 Integer
n
  where
    go :: a -> Integer -> (a, Integer)
go !a
k Integer
m = case Integer
m Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Integer
p of
      (Integer
q, Integer
0) -> a -> Integer -> (a, Integer)
go (a
k a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) Integer
q
      (Integer, Integer)
_      -> (a
k, Integer
m)
{-# INLINABLE splitOff #-}

smallOddPrimes :: [Integer]
smallOddPrimes :: [Integer]
smallOddPrimes
  = (Integer -> Bool) -> [Integer] -> [Integer]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
spBound)
  ([Integer] -> [Integer]) -> [Integer] -> [Integer]
forall a b. (a -> b) -> a -> b
$ (Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (\(I# Int#
k#) -> Int# -> Integer
IS (Word# -> Int#
word2Int# (
#if MIN_VERSION_base(4,16,0)
#ifdef WORDS_BIGENDIAN
  byteSwap16# (word16ToWord# (indexWord16OffAddr# smallPrimesAddr# k#))
#else
  word16ToWord# (indexWord16OffAddr# smallPrimesAddr# k#)
#endif
#else
#ifdef WORDS_BIGENDIAN
  byteSwap16# (indexWord16OffAddr# smallPrimesAddr# k#)
#else
  Addr# -> Int# -> Word#
indexWord16OffAddr# Addr#
smallPrimesAddr# Int#
k#
#endif
#endif
  )))
    [Int
1 .. Int
smallPrimesLength Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
  where
    !(Ptr Addr#
smallPrimesAddr#) = Ptr Word16
smallPrimesPtr

spBEx :: Word
spBEx :: Word
spBEx = Word
14

spBound :: Integer
spBound :: Integer
spBound = Integer
2Integer -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^Word
spBEx

-- n large, has no prime divisors < spBound
finishPower :: Word -> [(Integer, Word)] -> Integer -> (Integer, Word)
finishPower :: Word -> [(Integer, Word)] -> Integer -> (Integer, Word)
finishPower Word
e [(Integer, Word)]
pws Integer
n
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< (Integer
1 Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftL` Word -> Int
wordToInt (Word
2Word -> Word -> Word
forall a. Num a => a -> a -> a
*Word
spBEx))  = ((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
n [Integer
pInteger -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^Word
ex | (Integer
p,Word
ex) <- [(Integer, Word)]
pws], Word
1)    -- n is prime
  | Word
e Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
0  = Word -> Integer -> (Integer, Word)
rawPower Word
maxExp Integer
n
  | Bool
otherwise = [Word] -> (Integer, Word)
go [Word]
divs
    where
#ifdef MIN_VERSION_integer_gmp
      maxExp :: Word
maxExp = (Word# -> Word
W# (Int# -> Word#
int2Word# (Integer -> Int#
integerLog2# Integer
n))) Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
spBEx
#else
      maxExp = (W# (integerLog2# n)) `quot` spBEx
#endif
      divs :: [Word]
divs = Word -> Word -> [Word]
divisorsTo Word
maxExp Word
e
      go :: [Word] -> (Integer, Word)
go [] = ((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
n [Integer
pInteger -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^Word
ex | (Integer
p,Word
ex) <- [(Integer, Word)]
pws], Word
1)
      go (Word
d:[Word]
ds) = case Word -> Integer -> Maybe Integer
forall a b. (Integral a, Integral b) => b -> a -> Maybe a
exactRoot Word
d Integer
n of
                    Just Integer
r -> ((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
r [Integer
pInteger -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^(Word
ex Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
d) | (Integer
p,Word
ex) <- [(Integer, Word)]
pws], Word
d)
                    Maybe Integer
Nothing -> [Word] -> (Integer, Word)
go [Word]
ds

rawPower :: Word -> Integer -> (Integer, Word)
rawPower :: Word -> Integer -> (Integer, Word)
rawPower Word
mx Integer
n
  | Word
mx Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Word
2      = (Integer
n,Word
1)
  | Word
mx Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
2     = case Integer -> Maybe Integer
forall a. Integral a => a -> Maybe a
P2.exactSquareRoot Integer
n of
                    Just Integer
r  -> (Integer
r,Word
2)
                    Maybe Integer
Nothing -> (Integer
n,Word
1)
rawPower Word
mx Integer
n = case Integer -> Maybe Integer
forall a. Integral a => a -> Maybe a
P4.exactFourthRoot Integer
n of
                  Just Integer
r -> case Word -> Integer -> (Integer, Word)
rawPower (Word
mx Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
4) Integer
r of
                              (Integer
m,Word
e) -> (Integer
m, Word
4Word -> Word -> Word
forall a. Num a => a -> a -> a
*Word
e)
                  Maybe Integer
Nothing -> case Integer -> Maybe Integer
forall a. Integral a => a -> Maybe a
P2.exactSquareRoot Integer
n of
                               Just Integer
r -> case Word -> Integer -> (Integer, Word)
rawOddPower (Word
mx Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
2) Integer
r of
                                           (Integer
m,Word
e) -> (Integer
m, Word
2Word -> Word -> Word
forall a. Num a => a -> a -> a
*Word
e)
                               Maybe Integer
Nothing -> Word -> Integer -> (Integer, Word)
rawOddPower Word
mx Integer
n

rawOddPower :: Word -> Integer -> (Integer, Word)
rawOddPower :: Word -> Integer -> (Integer, Word)
rawOddPower Word
mx Integer
n
  | Word
mx Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Word
3       = (Integer
n,Word
1)
rawOddPower Word
mx Integer
n = case Integer -> Maybe Integer
forall a. Integral a => a -> Maybe a
P3.exactCubeRoot Integer
n of
                     Just Integer
r -> case Word -> Integer -> (Integer, Word)
rawOddPower (Word
mx Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
3) Integer
r of
                                 (Integer
m,Word
e) -> (Integer
m, Word
3Word -> Word -> Word
forall a. Num a => a -> a -> a
*Word
e)
                     Maybe Integer
Nothing -> Word -> Integer -> (Integer, Word)
badPower Word
mx Integer
n

badPower :: Word -> Integer -> (Integer, Word)
badPower :: Word -> Integer -> (Integer, Word)
badPower Word
mx Integer
n
  | Word
mx Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Word
5      = (Integer
n,Word
1)
  | Bool
otherwise   = Word -> Word -> Integer -> [Word] -> (Integer, Word)
forall t a.
(Integral t, Integral a) =>
a -> a -> t -> [a] -> (t, a)
go Word
1 Word
mx Integer
n ((Word -> Bool) -> [Word] -> [Word]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Word
mx) ([Word] -> [Word]) -> [Word] -> [Word]
forall a b. (a -> b) -> a -> b
$ (Word -> Word -> Word) -> Word -> [Word] -> [Word]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Word -> Word -> Word
forall a. Num a => a -> a -> a
(+) Word
5 ([Word] -> [Word]) -> [Word] -> [Word]
forall a b. (a -> b) -> a -> b
$ [Word] -> [Word]
forall a. [a] -> [a]
cycle [Word
2,Word
4])
    where
      go :: a -> a -> t -> [a] -> (t, a)
go !a
e a
b t
m (a
k:[a]
ks)
        | a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
k     = (t
m,a
e)
        | Bool
otherwise = case a -> t -> Maybe t
forall a b. (Integral a, Integral b) => b -> a -> Maybe a
exactRoot a
k t
m of
                        Just t
r -> a -> a -> t -> [a] -> (t, a)
go (a
ea -> a -> a
forall a. Num a => a -> a -> a
*a
k) (a
b a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
k) t
r (a
ka -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ks)
                        Maybe t
Nothing -> a -> a -> t -> [a] -> (t, a)
go a
e a
b t
m [a]
ks
      go a
e a
_ t
m []   = (t
m,a
e)

-- | List divisors of n, which are <= mx.
divisorsTo :: Word -> Word -> [Word]
divisorsTo :: Word -> Word -> [Word]
divisorsTo Word
mx Word
n = (Word -> Word -> Ordering) -> [Word] -> [Word]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy ((Word -> Word -> Ordering) -> Word -> Word -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip Word -> Word -> Ordering
forall a. Ord a => a -> a -> Ordering
compare) ([Word] -> [Word]) -> [Word] -> [Word]
forall a b. (a -> b) -> a -> b
$ [Word] -> Word -> [Word] -> [Word]
go [Word
1] Word
n (Word
2 Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: Word
3 Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: Word
5 Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: [Word]
prs)
  where
    -- unP p m = (k, m / p ^ k), where k is as large as possible such that p ^ k still divides m
    unP :: Word -> Word -> (Word, Word)
    unP :: Word -> Word -> (Word, Word)
unP Word
p Word
m = Word -> Word -> (Word, Word)
goP Word
0 Word
m
      where
        goP :: Word -> Word -> (Word, Word)
        goP :: Word -> Word -> (Word, Word)
goP !Word
i Word
j = case Word
j Word -> Word -> (Word, Word)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Word
p of
                     (Word
q,Word
r) | Word
r Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
0 -> Word -> Word -> (Word, Word)
goP (Word
iWord -> Word -> Word
forall a. Num a => a -> a -> a
+Word
1) Word
q
                           | Bool
otherwise -> (Word
i,Word
j)

    mset :: Word -> [Word] -> [Word]
mset Word
k = (Word -> Bool) -> [Word] -> [Word]
forall a. (a -> Bool) -> [a] -> [a]
filter (Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Word
mx) ([Word] -> [Word]) -> ([Word] -> [Word]) -> [Word] -> [Word]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word -> Word) -> [Word] -> [Word]
forall a b. (a -> b) -> [a] -> [b]
map (Word -> Word -> Word
forall a. Num a => a -> a -> a
* Word
k)

    prs :: [Word]
    prs :: [Word]
prs = Word
7 Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: (Word -> Bool) -> [Word] -> [Word]
forall a. (a -> Bool) -> [a] -> [a]
filter Word -> Bool
prm ((Word -> Word -> Word) -> Word -> [Word] -> [Word]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Word -> Word -> Word
forall a. Num a => a -> a -> a
(+) Word
11 ([Word] -> [Word]) -> [Word] -> [Word]
forall a b. (a -> b) -> a -> b
$ [Word] -> [Word]
forall a. [a] -> [a]
cycle [Word
2, Word
4, Word
2, Word
4, Word
6, Word
2, Word
6, Word
4])
    prm :: Word -> Bool
    prm :: Word -> Bool
prm Word
k = [Word] -> Bool
td [Word]
prs
      where
        td :: [Word] -> Bool
td (Word
p:[Word]
ps) = (Word
pWord -> Word -> Word
forall a. Num a => a -> a -> a
*Word
p Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
> Word
k) Bool -> Bool -> Bool
|| (Word
k Word -> Word -> Word
forall a. Integral a => a -> a -> a
`rem` Word
p Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
/= Word
0 Bool -> Bool -> Bool
&& [Word] -> Bool
td [Word]
ps)
        td []     = Bool
True

    go :: [Word] -> Word -> [Word] -> [Word]
go ![Word]
st Word
m (Word
p:[Word]
ps)
      | Word
m Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
1  = [Word]
st
      | Word
m Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Word
pWord -> Word -> Word
forall a. Num a => a -> a -> a
*Word
p = [Word]
st [Word] -> [Word] -> [Word]
forall a. [a] -> [a] -> [a]
++ Word -> [Word] -> [Word]
mset Word
m [Word]
st
      | Bool
otherwise =
        case Word -> Word -> (Word, Word)
unP Word
p Word
m of
          (Word
0,Word
_) -> [Word] -> Word -> [Word] -> [Word]
go [Word]
st Word
m [Word]
ps
          -- iterate f x = [x, f x, f (f x)...]
          (Word
k,Word
r) -> [Word] -> Word -> [Word] -> [Word]
go ([[Word]] -> [Word]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Int -> [[Word]] -> [[Word]]
forall a. Int -> [a] -> [a]
take (Word -> Int
wordToInt Word
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (([Word] -> [Word]) -> [Word] -> [[Word]]
forall a. (a -> a) -> a -> [a]
iterate (Word -> [Word] -> [Word]
mset Word
p) [Word]
st))) Word
r [Word]
ps
    go [Word]
st Word
m [] = [Word] -> Word -> [Word] -> [Word]
go [Word]
st Word
m [Word
mWord -> Word -> Word
forall a. Num a => a -> a -> a
+Word
1]