{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
module Math.NumberTheory.Roots.Squares
(
integerSquareRoot
, integerSquareRoot'
, integerSquareRootRem
, integerSquareRootRem'
, exactSquareRoot
, isSquare
, isSquare'
, isPossibleSquare
) where
import Data.Bits (finiteBitSize, (.&.))
import GHC.Exts (Ptr(..))
import Numeric.Natural (Natural)
import Math.NumberTheory.Roots.Squares.Internal
import Math.NumberTheory.Utils.BitMask (indexBitSet)
{-# SPECIALISE integerSquareRoot :: Int -> Int,
Word -> Word,
Integer -> Integer,
Natural -> Natural
#-}
integerSquareRoot :: Integral a => a -> a
integerSquareRoot :: a -> a
integerSquareRoot a
n
| a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0 = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"integerSquareRoot: negative argument"
| Bool
otherwise = a -> a
forall a. Integral a => a -> a
integerSquareRoot' a
n
{-# RULES
"integerSquareRoot'/Int" integerSquareRoot' = isqrtInt'
"integerSquareRoot'/Word" integerSquareRoot' = isqrtWord
"integerSquareRoot'/Integer" integerSquareRoot' = isqrtInteger
#-}
{-# INLINE [1] integerSquareRoot' #-}
integerSquareRoot' :: Integral a => a -> a
integerSquareRoot' :: a -> a
integerSquareRoot' = a -> a
forall a. Integral a => a -> a
isqrtA
{-# SPECIALISE integerSquareRootRem ::
Int -> (Int, Int),
Word -> (Word, Word),
Integer -> (Integer, Integer),
Natural -> (Natural, Natural)
#-}
integerSquareRootRem :: Integral a => a -> (a, a)
integerSquareRootRem :: a -> (a, a)
integerSquareRootRem a
n
| a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0 = [Char] -> (a, a)
forall a. HasCallStack => [Char] -> a
error [Char]
"integerSquareRootRem: negative argument"
| Bool
otherwise = a -> (a, a)
forall a. Integral a => a -> (a, a)
integerSquareRootRem' a
n
{-# RULES
"integerSquareRootRem'/Integer" integerSquareRootRem' = karatsubaSqrt
#-}
{-# INLINE [1] integerSquareRootRem' #-}
integerSquareRootRem' :: Integral a => a -> (a, a)
integerSquareRootRem' :: a -> (a, a)
integerSquareRootRem' a
n = (a
s, a
n a -> a -> a
forall a. Num a => a -> a -> a
- a
s a -> a -> a
forall a. Num a => a -> a -> a
* a
s)
where
s :: a
s = a -> a
forall a. Integral a => a -> a
integerSquareRoot' a
n
{-# SPECIALISE exactSquareRoot :: Int -> Maybe Int,
Word -> Maybe Word,
Integer -> Maybe Integer,
Natural -> Maybe Natural
#-}
exactSquareRoot :: Integral a => a -> Maybe a
exactSquareRoot :: a -> Maybe a
exactSquareRoot a
n
| a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0
, a -> Bool
forall a. Integral a => a -> Bool
isPossibleSquare a
n
, (a
r, a
0) <- a -> (a, a)
forall a. Integral a => a -> (a, a)
integerSquareRootRem' a
n = a -> Maybe a
forall a. a -> Maybe a
Just a
r
| Bool
otherwise = Maybe a
forall a. Maybe a
Nothing
{-# SPECIALISE isSquare :: Int -> Bool,
Word -> Bool,
Integer -> Bool,
Natural -> Bool
#-}
isSquare :: Integral a => a -> Bool
isSquare :: a -> Bool
isSquare a
n = a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0 Bool -> Bool -> Bool
&& a -> Bool
forall a. Integral a => a -> Bool
isSquare' a
n
{-# SPECIALISE isSquare' :: Int -> Bool,
Word -> Bool,
Integer -> Bool,
Natural -> Bool
#-}
isSquare' :: Integral a => a -> Bool
isSquare' :: a -> Bool
isSquare' a
n
| a -> Bool
forall a. Integral a => a -> Bool
isPossibleSquare a
n
, (a
_, a
0) <- a -> (a, a)
forall a. Integral a => a -> (a, a)
integerSquareRootRem' a
n = Bool
True
| Bool
otherwise = Bool
False
{-# SPECIALISE isPossibleSquare :: Int -> Bool,
Word -> Bool,
Integer -> Bool,
Natural -> Bool
#-}
isPossibleSquare :: Integral a => a -> Bool
isPossibleSquare :: a -> Bool
isPossibleSquare a
n'
= Ptr Word -> Int -> Bool
indexBitSet Ptr Word
mask256 (Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer
n Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. Integer
255))
Bool -> Bool -> Bool
&& Ptr Word -> Int -> Bool
indexBitSet Ptr Word
mask693 (Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` Integer
693))
Bool -> Bool -> Bool
&& Ptr Word -> Int -> Bool
indexBitSet Ptr Word
mask325 (Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` Integer
325))
where
n :: Integer
n = a -> Integer
forall a. Integral a => a -> Integer
toInteger a
n'
mask256 :: Ptr Word
mask256 :: Ptr Word
mask256 = Addr# -> Ptr Word
forall a. Addr# -> Ptr a
Ptr Addr#
"\DC3\STX\ETX\STX\DC2\STX\STX\STX\DC3\STX\STX\STX\DC2\STX\STX\STX\DC2\STX\ETX\STX\DC2\STX\STX\STX\DC2\STX\STX\STX\DC2\STX\STX\STX"#
mask693 :: Ptr Word
mask693 :: Ptr Word
mask693 = Addr# -> Ptr Word
forall a. Addr# -> Ptr a
Ptr Addr#
"\DC3\STXA\STX0\NUL\STX\EOTI\NUL\STX\t\CAN\NUL\NULB\164\NUL\DC1\EOT\b\STX\NUL@P\128@\NUL\STX\t\128 \SOH\DLE\NUL\SOH\130$\NUL\128\DC4(\NUL\NUL\SOH\DC2\NUL\f\STX\DC4\SOH\NUL \b\NUL\"\NUL\128\EOT`\144\NUL\b\129\NULE\DC2\DLE@\STX\EOT\NUL\129\NUL\t\b\EOT\SOH\194\128\NUL\DLE\EOT\NUL\DLE\NUL\NUL"#
mask325 :: Ptr Word
mask325 :: Ptr Word
mask325 = Addr# -> Ptr Word
forall a. Addr# -> Ptr a
Ptr Addr#
"\DC3B\SOH&\144\NUL\n!%\140\STXH0\SOH\DC4BJ\b\ENQ\144@\STX(\132\148\DLE\n \131\EOTP\f)!\DC4@\STX\EM\160\DLE\DC2"#
isqrtInt' :: Int -> Int
isqrtInt' :: Int -> Int
isqrtInt' Int
n
| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
r = Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1
| Bool
otherwise = Int
r
where
!r :: Int
r = (Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
truncate :: Double -> Int) (Double -> Int) -> (Double -> Double) -> Double -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
isqrtWord :: Word -> Word
isqrtWord :: Word -> Word
isqrtWord Word
n
| Word
n Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< (Word
rWord -> Word -> Word
forall a. Num a => a -> a -> a
*Word
r)
Bool -> Bool -> Bool
|| Word -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
64 Bool -> Bool -> Bool
&& Word
r Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
4294967296
= Word
rWord -> Word -> Word
forall a. Num a => a -> a -> a
-Word
1
| Bool
otherwise = Word
r
where
!r :: Word
r = (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Int -> Word) (Int -> Word) -> (Double -> Int) -> Double -> Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
truncate :: Double -> Int) (Double -> Int) -> (Double -> Double) -> Double -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Word) -> Double -> Word
forall a b. (a -> b) -> a -> b
$ Word -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n
{-# INLINE isqrtInteger #-}
isqrtInteger :: Integer -> Integer
isqrtInteger :: Integer -> Integer
isqrtInteger = (Integer, Integer) -> Integer
forall a b. (a, b) -> a
fst ((Integer, Integer) -> Integer)
-> (Integer -> (Integer, Integer)) -> Integer -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> (Integer, Integer)
karatsubaSqrt