integer-roots-1.0.2.0: Math/NumberTheory/Roots/General.hs
-- |
-- 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 1 n = n
integerRoot 2 n = P2.integerSquareRoot n
integerRoot 3 n = P3.integerCubeRoot n
integerRoot 4 n = P4.integerFourthRoot n
integerRoot k n
| k < 1 = error "integerRoot: negative exponent or exponent 0"
| n < 0 && even k = error "integerRoot: negative radicand for even exponent"
| n < 0 =
let r = negate . fromInteger . integerRoot k . negate $ fromIntegral n
in if r^k == n then r else (r-1)
| n == 0 = 0
| n < 31 = 1
| kTooLarge = 1
| otherwise = fromInteger $ newtonK (toInteger k) (toInteger n) a
where
a = appKthRoot (fromIntegral k) (toInteger n)
kTooLarge = (toInteger k /= toInteger (fromIntegral k `asTypeOf` n)) -- k doesn't fit in n's type
|| (toInteger k > toInteger (maxBound :: Int)) -- 2^k doesn't fit in Integer
#ifdef MIN_VERSION_integer_gmp
|| (I# (integerLog2# (toInteger n)) < fromIntegral 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 1 n = Just n
exactRoot 2 n = P2.exactSquareRoot n
exactRoot 3 n = P3.exactCubeRoot n
exactRoot 4 n = P4.exactFourthRoot n
exactRoot k n
| n == 1 = Just 1
| k < 1 = Nothing
| n < 0 && even k = Nothing
| n < 0 = let m = negate n in
if m < 0
then fmap (fromInteger . negate) (exactRoot k (negate (toInteger n)))
else fmap negate (exactRoot k m)
| n < 2 = Just n
| n < 31 = Nothing
| kTooLarge = Nothing
| otherwise = case k `rem` 12 of
0 | c4 && c3 && ok -> Just r
| otherwise -> Nothing
2 | c2 && ok -> Just r
| otherwise -> Nothing
3 | c3 && ok -> Just r
| otherwise -> Nothing
4 | c4 && ok -> Just r
| otherwise -> Nothing
6 | c3 && c2 && ok -> Just r
| otherwise -> Nothing
8 | c4 && ok -> Just r
| otherwise -> Nothing
9 | c3 && ok -> Just r
| otherwise -> Nothing
10 | c2 && ok -> Just r
| otherwise -> Nothing
_ | ok -> Just r
| otherwise -> Nothing
where
k' :: Int
k' = fromIntegral k
r = integerRoot k' n
c2 = P2.isPossibleSquare n
c3 = P3.isPossibleCube n
c4 = P4.isPossibleFourthPower n
ok = r^k == n
kTooLarge = (toInteger k /= toInteger (fromIntegral k `asTypeOf` n)) -- k doesn't fit in n's type
|| (toInteger k > toInteger (maxBound :: Int)) -- 2^k doesn't fit in Integer
#ifdef MIN_VERSION_integer_gmp
|| (I# (integerLog2# (toInteger n)) < fromIntegral 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 k n = isJust (exactRoot k 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 n
| n == 0 || n == 1 = True
| otherwise = k > 1
where
(_,k) = highestPower 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 n'
| abs n <= 1 = (n', 3)
| n < 0 = case integerHighPower (negate n) of
(r,e) -> case countTrailingZeros e of
k -> (negate $ fromInteger (sqr k r), e `shiftR` k)
| otherwise = case integerHighPower n of
(r,e) -> (fromInteger r, e)
where
n :: Integer
n = toInteger n'
sqr :: Int -> Integer -> Integer
sqr 0 m = m
sqr k m = sqr (k-1) (m*m)
------------------------------------------------------------------------------------------
-- Auxiliary functions --
------------------------------------------------------------------------------------------
newtonK :: Integer -> Integer -> Integer -> Integer
newtonK k n a = go (step a)
where
step m = ((k - 1) * m + n `quot` m ^ (k - 1)) `quot` k
go m
| l < m = go l
| otherwise = m
where
l = step m
-- find an approximation to the k-th root
-- here, k > 4 and n > 31
appKthRoot :: Int -> Integer -> Integer
appKthRoot (I# k#) (IS n#) = IS (double2Int# (int2Double# n# **## (1.0## /## int2Double# k#)))
#ifdef MIN_VERSION_integer_gmp
appKthRoot k@(I# k#) n
| k >= 256 = 1 `shiftLInteger` (integerLog2# n `quotInt#` k# +# 1#)
| otherwise =
case integerLog2# n of
l# -> case l# `quotInt#` k# of
0# -> 1
1# -> 3
2# -> 5
3# -> 11
h# | isTrue# (h# <# 500#) ->
floor (scaleFloat (I# h#)
(fromInteger (n `shiftRInteger` (h# *# k#)) ** (1/fromIntegral k) :: Double))
| otherwise ->
floor (scaleFloat 400 (fromInteger (n `shiftRInteger` (h# *# k#)) ** (1/fromIntegral k) :: Double))
`shiftLInteger` (h# -# 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 n
| n < 4 = (n,1)
| otherwise = case splitOff 2 n of
(e2,m) | m == 1 -> (2,e2)
| otherwise -> findHighPower e2 (if e2 == 0 then [] else [(2,e2)]) m r smallOddPrimes
where
r = P2.integerSquareRoot m
findHighPower :: Word -> [(Integer, Word)] -> Integer -> Integer -> [Integer] -> (Integer, Word)
findHighPower 1 pws m _ _ = (foldl' (*) m [p^e | (p,e) <- pws], 1)
findHighPower e pws 1 _ _ = (foldl' (*) 1 [p^(ex `quot` e) | (p,ex) <- pws], e)
findHighPower e pws m s (p:ps)
| s < p = findHighPower 1 pws m s []
| otherwise =
case splitOff p m of
(0,_) -> findHighPower e pws m s ps
(k,r) -> findHighPower (gcd k e) ((p,k):pws) r (P2.integerSquareRoot r) ps
findHighPower e pws m _ [] = finishPower e pws m
splitOff :: Integer -> Integer -> (Word, Integer)
splitOff !_ 0 = (0, 0) -- prevent infinite loop
splitOff p n = go 0 n
where
go !k m = case m `quotRem` p of
(q, 0) -> go (k + 1) q
_ -> (k, m)
{-# INLINABLE splitOff #-}
smallOddPrimes :: [Integer]
smallOddPrimes
= takeWhile (< spBound)
$ map (\(I# k#) -> IS (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
indexWord16OffAddr# smallPrimesAddr# k#
#endif
#endif
)))
[1 .. smallPrimesLength - 1]
where
!(Ptr smallPrimesAddr#) = smallPrimesPtr
spBEx :: Word
spBEx = 14
spBound :: Integer
spBound = 2^spBEx
-- n large, has no prime divisors < spBound
finishPower :: Word -> [(Integer, Word)] -> Integer -> (Integer, Word)
finishPower e pws n
| n < (1 `shiftL` wordToInt (2*spBEx)) = (foldl' (*) n [p^ex | (p,ex) <- pws], 1) -- n is prime
| e == 0 = rawPower maxExp n
| otherwise = go divs
where
#ifdef MIN_VERSION_integer_gmp
maxExp = (W# (int2Word# (integerLog2# n))) `quot` spBEx
#else
maxExp = (W# (integerLog2# n)) `quot` spBEx
#endif
divs = divisorsTo maxExp e
go [] = (foldl' (*) n [p^ex | (p,ex) <- pws], 1)
go (d:ds) = case exactRoot d n of
Just r -> (foldl' (*) r [p^(ex `quot` d) | (p,ex) <- pws], d)
Nothing -> go ds
rawPower :: Word -> Integer -> (Integer, Word)
rawPower mx n
| mx < 2 = (n,1)
| mx == 2 = case P2.exactSquareRoot n of
Just r -> (r,2)
Nothing -> (n,1)
rawPower mx n = case P4.exactFourthRoot n of
Just r -> case rawPower (mx `quot` 4) r of
(m,e) -> (m, 4*e)
Nothing -> case P2.exactSquareRoot n of
Just r -> case rawOddPower (mx `quot` 2) r of
(m,e) -> (m, 2*e)
Nothing -> rawOddPower mx n
rawOddPower :: Word -> Integer -> (Integer, Word)
rawOddPower mx n
| mx < 3 = (n,1)
rawOddPower mx n = case P3.exactCubeRoot n of
Just r -> case rawOddPower (mx `quot` 3) r of
(m,e) -> (m, 3*e)
Nothing -> badPower mx n
badPower :: Word -> Integer -> (Integer, Word)
badPower mx n
| mx < 5 = (n,1)
| otherwise = go 1 mx n (takeWhile (<= mx) $ scanl (+) 5 $ cycle [2,4])
where
go !e b m (k:ks)
| b < k = (m,e)
| otherwise = case exactRoot k m of
Just r -> go (e*k) (b `quot` k) r (k:ks)
Nothing -> go e b m ks
go e _ m [] = (m,e)
-- | List divisors of n, which are <= mx.
divisorsTo :: Word -> Word -> [Word]
divisorsTo mx n = sortBy (flip compare) $ go [1] n (2 : 3 : 5 : 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 p m = goP 0 m
where
goP :: Word -> Word -> (Word, Word)
goP !i j = case j `quotRem` p of
(q,r) | r == 0 -> goP (i+1) q
| otherwise -> (i,j)
mset k = filter (<= mx) . map (* k)
prs :: [Word]
prs = 7 : filter prm (scanl (+) 11 $ cycle [2, 4, 2, 4, 6, 2, 6, 4])
prm :: Word -> Bool
prm k = td prs
where
td (p:ps) = (p*p > k) || (k `rem` p /= 0 && td ps)
td [] = True
go !st m (p:ps)
| m == 1 = st
| m < p*p = st ++ mset m st
| otherwise =
case unP p m of
(0,_) -> go st m ps
-- iterate f x = [x, f x, f (f x)...]
(k,r) -> go (concat (take (wordToInt k + 1) (iterate (mset p) st))) r ps
go st m [] = go st m [m+1]