{-# LANGUAGE MagicHash #-}
module Math.NumberTheory.Roots.Fourth
( integerFourthRoot
, integerFourthRoot'
, exactFourthRoot
, isFourthPower
, isFourthPower'
, isPossibleFourthPower
) where
import Data.Bits (finiteBitSize, (.&.))
import GHC.Exts (Int#, Ptr(..), uncheckedIShiftRA#, int2Double#, double2Int#, isTrue#, sqrtDouble#, (<#), (*#), (-#))
import GHC.Integer.GMP.Internals (Integer(..), shiftLInteger, shiftRInteger, sizeofBigNat#)
import GHC.Integer.Logarithms (integerLog2#)
import Numeric.Natural (Natural)
import Math.NumberTheory.Utils.BitMask (indexBitSet)
{-# SPECIALISE integerFourthRoot :: Int -> Int,
Word -> Word,
Integer -> Integer,
Natural -> Natural
#-}
integerFourthRoot :: Integral a => a -> a
integerFourthRoot n
| n < 0 = error "integerFourthRoot: negative argument"
| otherwise = integerFourthRoot' n
{-# RULES
"integerFourthRoot'/Int" integerFourthRoot' = biSqrtInt
"integerFourthRoot'/Word" integerFourthRoot' = biSqrtWord
"integerFourthRoot'/Integer" integerFourthRoot' = biSqrtIgr
#-}
{-# INLINE [1] integerFourthRoot' #-}
integerFourthRoot' :: Integral a => a -> a
integerFourthRoot' 0 = 0
integerFourthRoot' n = newton4 n (approxBiSqrt n)
{-# SPECIALISE exactFourthRoot :: Int -> Maybe Int,
Word -> Maybe Word,
Integer -> Maybe Integer,
Natural -> Maybe Natural
#-}
exactFourthRoot :: Integral a => a -> Maybe a
exactFourthRoot 0 = Just 0
exactFourthRoot n
| n < 0 = Nothing
| isPossibleFourthPower n && r2*r2 == n = Just r
| otherwise = Nothing
where
r = integerFourthRoot' n
r2 = r*r
{-# SPECIALISE isFourthPower :: Int -> Bool,
Word -> Bool,
Integer -> Bool,
Natural -> Bool
#-}
isFourthPower :: Integral a => a -> Bool
isFourthPower 0 = True
isFourthPower n = n > 0 && isFourthPower' n
{-# SPECIALISE isFourthPower' :: Int -> Bool,
Word -> Bool,
Integer -> Bool,
Natural -> Bool
#-}
isFourthPower' :: Integral a => a -> Bool
isFourthPower' n = isPossibleFourthPower n && r2*r2 == n
where
r = integerFourthRoot' n
r2 = r*r
{-# SPECIALISE isPossibleFourthPower :: Int -> Bool,
Word -> Bool,
Integer -> Bool,
Natural -> Bool
#-}
isPossibleFourthPower :: Integral a => a -> Bool
isPossibleFourthPower n'
= indexBitSet mask256 (fromInteger (n .&. 255))
&& indexBitSet mask425 (fromInteger (n `rem` 425))
&& indexBitSet mask377 (fromInteger (n `rem` 377))
where
n = toInteger n'
{-# SPECIALISE newton4 :: Integer -> Integer -> Integer #-}
newton4 :: Integral a => a -> a -> a
newton4 n a = go (step a)
where
step k = (3*k + n `quot` (k*k*k)) `quot` 4
go k
| m < k = go m
| otherwise = k
where
m = step k
{-# SPECIALISE approxBiSqrt :: Integer -> Integer #-}
approxBiSqrt :: Integral a => a -> a
approxBiSqrt = fromInteger . appBiSqrt . fromIntegral
appBiSqrt :: Integer -> Integer
appBiSqrt (S# i#) = S# (double2Int# (sqrtDouble# (sqrtDouble# (int2Double# i#))))
appBiSqrt n@(Jp# bn#)
| isTrue# ((sizeofBigNat# bn#) <# thresh#) =
floor (sqrt . sqrt $ fromInteger n :: Double)
| otherwise = case integerLog2# n of
l# -> case uncheckedIShiftRA# l# 2# -# 47# of
h# -> case shiftRInteger n (4# *# h#) of
m -> case floor (sqrt $ sqrt $ fromInteger m :: Double) of
r -> shiftLInteger r h#
where
thresh# :: Int#
thresh# = if finiteBitSize (0 :: Word) == 64 then 5# else 9#
appBiSqrt _ = error "integerFourthRoot': negative argument"
mask256 :: Ptr Word
mask256 = Ptr "\ETX\NUL\ETX\NUL\STX\NUL\STX\NUL\STX\NUL\STX\NUL\STX\NUL\STX\NUL\STX\NUL\STX\NUL\STX\NUL\STX\NUL\STX\NUL\STX\NUL\STX\NUL\STX\NUL"#
mask425 :: Ptr Word
mask425 = Ptr "\ETX\NUL!\NUL\NUL\NUL\f\NUL\NUL\NULB\NUL \EOT\NUL\NUL\NUL\SOH\NUL\NUL@\b\NUL\132\NUL\SOH\NUL \STX\NUL\NUL\b\SOH\128\DLE\NUL\NUL\NUL\EOT\NUL\NUL\NUL!\NUL\DLE\STX\128\NUL\128\NUL\NUL\NUL \NUL"#
mask377 :: Ptr Word
mask377 = Ptr "\ETX\NUL\SOH \NUL\NUL0\NUL\STXD\130@\NUL\b \NUL\NUL\b\EOT\SOH \ACK\NUL\NUL@\DLE\NUL\NUL\NUL\NUL\NUL\SOH!\NUL\NUL@\NUL\NUL\NUL\n@\NUL\b\NUL\NUL\DLE \NUL"#
biSqRootIntLimit :: Int
biSqRootIntLimit = if finiteBitSize (0 :: Word) == 64 then 55108 else 215
biSqrtInt :: Int -> Int
biSqrtInt 0 = 0
biSqrtInt n
| r > biSqRootIntLimit = biSqRootIntLimit
| n < r4 = r-1
| otherwise = r
where
x :: Double
x = fromIntegral n
r = truncate (sqrt (sqrt x))
r2 = r*r
r4 = r2*r2
biSqRootWordLimit :: Word
biSqRootWordLimit = if finiteBitSize (0 :: Word) == 64 then 65535 else 255
biSqrtWord :: Word -> Word
biSqrtWord 0 = 0
biSqrtWord n
| r > biSqRootWordLimit = biSqRootWordLimit
| n < r4 = r-1
| otherwise = r
where
x :: Double
x = fromIntegral n
r = truncate (sqrt (sqrt x))
r2 = r*r
r4 = r2*r2
biSqrtIgr :: Integer -> Integer
biSqrtIgr 0 = 0
biSqrtIgr n = newton4 n (approxBiSqrt n)