{-# LANGUAGE MagicHash, BangPatterns, PatternGuards, CPP, FlexibleContexts #-}
module Math.NumberTheory.Powers.Squares
(
integerSquareRoot
, integerSquareRoot'
, integerSquareRootRem
, integerSquareRootRem'
, exactSquareRoot
, isSquare
, isSquare'
, isPossibleSquare
, isPossibleSquare2
) where
#include "MachDeps.h"
import Control.Monad.ST
import Data.Bits
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Unboxed.Mutable as MV
import Numeric.Natural
import Math.NumberTheory.Powers.Squares.Internal
{-# SPECIALISE integerSquareRoot :: Int -> Int,
Word -> Word,
Integer -> Integer,
Natural -> Natural
#-}
integerSquareRoot :: Integral a => a -> a
integerSquareRoot n
| n < 0 = error "integerSquareRoot: negative argument"
| otherwise = integerSquareRoot' n
{-# RULES
"integerSquareRoot'/Int" integerSquareRoot' = isqrtInt'
"integerSquareRoot'/Word" integerSquareRoot' = isqrtWord
"integerSquareRoot'/Integer" integerSquareRoot' = isqrtInteger
#-}
{-# INLINE [1] integerSquareRoot' #-}
integerSquareRoot' :: Integral a => a -> a
integerSquareRoot' = isqrtA
{-# SPECIALISE integerSquareRootRem ::
Int -> (Int, Int),
Word -> (Word, Word),
Integer -> (Integer, Integer),
Natural -> (Natural, Natural)
#-}
integerSquareRootRem :: Integral a => a -> (a, a)
integerSquareRootRem n
| n < 0 = error "integerSquareRootRem: negative argument"
| otherwise = integerSquareRootRem' n
{-# RULES
"integerSquareRootRem'/Integer" integerSquareRootRem' = karatsubaSqrt
#-}
{-# INLINE [1] integerSquareRootRem' #-}
integerSquareRootRem' :: Integral a => a -> (a, a)
integerSquareRootRem' n = (s, n - s * s)
where
s = integerSquareRoot' n
{-# SPECIALISE exactSquareRoot :: Int -> Maybe Int,
Word -> Maybe Word,
Integer -> Maybe Integer,
Natural -> Maybe Natural
#-}
exactSquareRoot :: Integral a => a -> Maybe a
exactSquareRoot n
| n >= 0
, isPossibleSquare n
, (r, 0) <- integerSquareRootRem' n = Just r
| otherwise = Nothing
{-# SPECIALISE isSquare :: Int -> Bool,
Word -> Bool,
Integer -> Bool,
Natural -> Bool
#-}
isSquare :: Integral a => a -> Bool
isSquare n = n >= 0 && isSquare' n
{-# SPECIALISE isSquare' :: Int -> Bool,
Word -> Bool,
Integer -> Bool,
Natural -> Bool
#-}
isSquare' :: Integral a => a -> Bool
isSquare' n
| isPossibleSquare n
, (_, 0) <- integerSquareRootRem' n = True
| otherwise = False
{-# SPECIALISE isPossibleSquare :: Int -> Bool,
Word -> Bool,
Integer -> Bool,
Natural -> Bool
#-}
isPossibleSquare :: Integral a => a -> Bool
isPossibleSquare n
= V.unsafeIndex sr256 ((fromIntegral n) .&. 255)
&& V.unsafeIndex sr693 (fromIntegral (n `rem` 693))
&& V.unsafeIndex sr325 (fromIntegral (n `rem` 325))
{-# SPECIALISE isPossibleSquare2 :: Int -> Bool,
Word -> Bool,
Integer -> Bool,
Natural -> Bool
#-}
isPossibleSquare2 :: Integral a => a -> Bool
isPossibleSquare2 n
= V.unsafeIndex sr256 ((fromIntegral n) .&. 255)
&& V.unsafeIndex sr819 (fromIntegral (n `rem` 819))
&& V.unsafeIndex sr1025 (fromIntegral (n `rem` 1025))
&& V.unsafeIndex sr2047 (fromIntegral (n `rem` 2047))
&& V.unsafeIndex sr4097 (fromIntegral (n `rem` 4097))
&& V.unsafeIndex sr341 (fromIntegral (n `rem` 341))
sqRemArray :: Int -> V.Vector Bool
sqRemArray md = runST $ do
ar <- MV.replicate md False
let !stop = (md `quot` 2) + 1
fill k
| k < stop = MV.unsafeWrite ar ((k*k) `rem` md) True >> fill (k+1)
| otherwise = return ()
MV.unsafeWrite ar 0 True
MV.unsafeWrite ar 1 True
fill 2
V.unsafeFreeze ar
sr256 :: V.Vector Bool
sr256 = sqRemArray 256
sr819 :: V.Vector Bool
sr819 = sqRemArray 819
sr4097 :: V.Vector Bool
sr4097 = sqRemArray 4097
sr341 :: V.Vector Bool
sr341 = sqRemArray 341
sr1025 :: V.Vector Bool
sr1025 = sqRemArray 1025
sr2047 :: V.Vector Bool
sr2047 = sqRemArray 2047
sr693 :: V.Vector Bool
sr693 = sqRemArray 693
sr325 :: V.Vector Bool
sr325 = sqRemArray 325
isqrtInt' :: Int -> Int
isqrtInt' n
| n < r*r = r-1
| otherwise = r
where
!r = (truncate :: Double -> Int) . sqrt $ fromIntegral n
isqrtWord :: Word -> Word
isqrtWord n
| n < (r*r)
#if WORD_SIZE_IN_BITS == 64
|| r == 4294967296
#endif
= r-1
| otherwise = r
where
!r = (fromIntegral :: Int -> Word) . (truncate :: Double -> Int) . sqrt $ fromIntegral n
{-# INLINE isqrtInteger #-}
isqrtInteger :: Integer -> Integer
isqrtInteger = fst . karatsubaSqrt