-- |
-- Module:      Math.NumberTheory.Roots.Squares
-- Copyright:   (c) 2011 Daniel Fischer, 2016-2020 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Daniel Fischer <daniel.is.fischer@googlemail.com>
--
-- Functions dealing with fourth powers. Efficient calculation of integer fourth
-- roots and efficient testing for being a square's square.

{-# LANGUAGE CPP       #-}
{-# LANGUAGE MagicHash #-}

module Math.NumberTheory.Roots.Fourth
    ( integerFourthRoot
    , integerFourthRoot'
    , exactFourthRoot
    , isFourthPower
    , isFourthPower'
    , isPossibleFourthPower
    ) where

import Data.Bits (finiteBitSize, (.&.))
import GHC.Exts (Int#, Ptr(..), int2Double#, double2Int#, isTrue#, sqrtDouble#, (<#))
import Numeric.Natural (Natural)

#ifdef MIN_VERSION_integer_gmp
import GHC.Exts (uncheckedIShiftRA#, (*#), (-#))
import GHC.Integer.GMP.Internals (Integer(..), shiftLInteger, shiftRInteger, sizeofBigNat#)
import GHC.Integer.Logarithms (integerLog2#)
#define IS S#
#define IP Jp#
#define bigNatSize sizeofBigNat
#else
import GHC.Exts (uncheckedShiftRL#, minusWord#, timesWord#)
import GHC.Num.BigNat (bigNatSize#)
import GHC.Num.Integer (Integer(..), integerLog2#, integerShiftR#, integerShiftL#)
#endif

import Math.NumberTheory.Utils.BitMask (indexBitSet)

-- | Calculate the integer fourth root of a nonnegative number,
--   that is, the largest integer @r@ with @r^4 <= n@.
--   Throws an error on negaitve input.
{-# SPECIALISE integerFourthRoot :: Int -> Int,
                                    Word -> Word,
                                    Integer -> Integer,
                                    Natural -> Natural
  #-}
integerFourthRoot :: Integral a => a -> a
integerFourthRoot :: a -> a
integerFourthRoot 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]
"integerFourthRoot: negative argument"
    | Bool
otherwise = a -> a
forall a. Integral a => a -> a
integerFourthRoot' a
n

-- | Calculate the integer fourth root of a nonnegative number,
--   that is, the largest integer @r@ with @r^4 <= n@.
--   The condition is /not/ checked.
{-# RULES
"integerFourthRoot'/Int"     integerFourthRoot' = biSqrtInt
"integerFourthRoot'/Word"    integerFourthRoot' = biSqrtWord
"integerFourthRoot'/Integer" integerFourthRoot' = biSqrtIgr
  #-}
{-# INLINE [1] integerFourthRoot' #-}
integerFourthRoot' :: Integral a => a -> a
integerFourthRoot' :: a -> a
integerFourthRoot' a
0 = a
0
integerFourthRoot' a
n = a -> a -> a
forall a. Integral a => a -> a -> a
newton4 a
n (a -> a
forall a. Integral a => a -> a
approxBiSqrt a
n)

-- | Returns @Nothing@ if @n@ is not a fourth power,
--   @Just r@ if @n == r^4@ and @r >= 0@.
{-# SPECIALISE exactFourthRoot :: Int -> Maybe Int,
                                  Word -> Maybe Word,
                                  Integer -> Maybe Integer,
                                  Natural -> Maybe Natural
  #-}
exactFourthRoot :: Integral a => a -> Maybe a
exactFourthRoot :: a -> Maybe a
exactFourthRoot a
0 = a -> Maybe a
forall a. a -> Maybe a
Just a
0
exactFourthRoot a
n
    | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0     = Maybe a
forall a. Maybe a
Nothing
    | a -> Bool
forall a. Integral a => a -> Bool
isPossibleFourthPower a
n Bool -> Bool -> Bool
&& a
r2a -> a -> a
forall a. Num a => a -> a -> a
*a
r2 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
n = a -> Maybe a
forall a. a -> Maybe a
Just a
r
    | Bool
otherwise = Maybe a
forall a. Maybe a
Nothing
      where
        r :: a
r = a -> a
forall a. Integral a => a -> a
integerFourthRoot' a
n
        r2 :: a
r2 = a
ra -> a -> a
forall a. Num a => a -> a -> a
*a
r

-- | Test whether an integer is a fourth power.
--   First nonnegativity is checked, then the unchecked
--   test is called.
{-# SPECIALISE isFourthPower :: Int -> Bool,
                                Word -> Bool,
                                Integer -> Bool,
                                Natural -> Bool
  #-}
isFourthPower :: Integral a => a -> Bool
isFourthPower :: a -> Bool
isFourthPower a
0 = Bool
True
isFourthPower 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
isFourthPower' a
n

-- | Test whether a nonnegative number is a fourth power.
--   The condition is /not/ checked. If a number passes the
--   'isPossibleFourthPower' test, its integer fourth root
--   is calculated.
{-# SPECIALISE isFourthPower' :: Int -> Bool,
                                 Word -> Bool,
                                 Integer -> Bool,
                                 Natural -> Bool
  #-}
isFourthPower' :: Integral a => a -> Bool
isFourthPower' :: a -> Bool
isFourthPower' a
n = a -> Bool
forall a. Integral a => a -> Bool
isPossibleFourthPower a
n Bool -> Bool -> Bool
&& a
r2a -> a -> a
forall a. Num a => a -> a -> a
*a
r2 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
n
  where
    r :: a
r = a -> a
forall a. Integral a => a -> a
integerFourthRoot' a
n
    r2 :: a
r2 = a
ra -> a -> a
forall a. Num a => a -> a -> a
*a
r

-- | Test whether a nonnegative number is a possible fourth power.
--   The condition is /not/ checked.
--   This eliminates about 99.958% of numbers.
{-# SPECIALISE isPossibleFourthPower :: Int -> Bool,
                                        Word -> Bool,
                                        Integer -> Bool,
                                        Natural -> Bool
  #-}
isPossibleFourthPower :: Integral a => a -> Bool
isPossibleFourthPower :: a -> Bool
isPossibleFourthPower 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
mask425 (Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` Integer
425))
  Bool -> Bool -> Bool
&& Ptr Word -> Int -> Bool
indexBitSet Ptr Word
mask377 (Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` Integer
377))
  where
    n :: Integer
n = a -> Integer
forall a. Integral a => a -> Integer
toInteger a
n'

{-# SPECIALISE newton4 :: Integer -> Integer -> Integer #-}
newton4 :: Integral a => a -> a -> a
newton4 :: a -> a -> a
newton4 a
n a
a = a -> a
go (a -> a
step a
a)
      where
        step :: a -> a
step a
k = (a
3a -> a -> a
forall a. Num a => a -> a -> a
*a
k a -> a -> a
forall a. Num a => a -> a -> a
+ a
n a -> a -> a
forall a. Integral a => a -> a -> a
`quot` (a
ka -> a -> a
forall a. Num a => a -> a -> a
*a
ka -> a -> a
forall a. Num a => a -> a -> a
*a
k)) a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
4
        go :: a -> a
go a
k
            | a
m a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
k     = a -> a
go a
m
            | Bool
otherwise = a
k
              where
                m :: a
m = a -> a
step a
k

{-# SPECIALISE approxBiSqrt :: Integer -> Integer #-}
approxBiSqrt :: Integral a => a -> a
approxBiSqrt :: a -> a
approxBiSqrt = Integer -> a
forall a. Num a => Integer -> a
fromInteger (Integer -> a) -> (a -> Integer) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Integer
appBiSqrt (Integer -> Integer) -> (a -> Integer) -> a -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- Find a fairly good approximation to the fourth root.
-- About 48 bits should be correct for large Integers.
appBiSqrt :: Integer -> Integer
appBiSqrt :: Integer -> Integer
appBiSqrt (IS Int#
i#) = Int# -> Integer
IS (Double# -> Int#
double2Int# (Double# -> Double#
sqrtDouble# (Double# -> Double#
sqrtDouble# (Int# -> Double#
int2Double# Int#
i#))))
appBiSqrt n :: Integer
n@(IP bn#)
    | Int# -> Bool
isTrue# ((bigNatSize# bn#) <# thresh#) =
          Double -> Integer
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> (Double -> Double) -> Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Integer -> Double
forall a. Num a => Integer -> a
fromInteger Integer
n :: Double)
    | Bool
otherwise = case Integer -> Int#
integerLog2# Integer
n of
#ifdef MIN_VERSION_integer_gmp
                    Int#
l# -> case Int# -> Int# -> Int#
uncheckedIShiftRA# Int#
l# Int#
2# Int# -> Int# -> Int#
-# Int#
47# of
                            Int#
h# -> case Integer -> Int# -> Integer
shiftRInteger Integer
n (Int#
4# Int# -> Int# -> Int#
*# Int#
h#) of
                                    Integer
m -> case Double -> Integer
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Integer -> Double
forall a. Num a => Integer -> a
fromInteger Integer
m :: Double) of
                                            Integer
r -> Integer -> Int# -> Integer
shiftLInteger Integer
r Int#
h#
#else
                    l# -> case uncheckedShiftRL# l# 2# `minusWord#` 47## of
                            h# -> case integerShiftR# n (4## `timesWord#` h#) of
                                    m -> case floor (sqrt $ sqrt $ fromInteger m :: Double) of
                                            r -> integerShiftL# r h#
#endif
    where
        -- threshold for shifting vs. direct fromInteger
        -- we shift when we expect more than 256 bits
        thresh# :: Int#
        thresh# :: Int#
thresh# = if Word -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
64 then Int#
5# else Int#
9#

-- There's already a check for negative in integerFourthRoot,
-- but integerFourthRoot' is exported directly too.
appBiSqrt Integer
_ = [Char] -> Integer
forall a. HasCallStack => [Char] -> a
error [Char]
"integerFourthRoot': negative argument"

-----------------------------------------------------------------------------
-- Generated by 'Math.NumberTheory.Utils.BitMask.vectorToAddrLiteral'

mask256 :: Ptr Word
mask256 :: Ptr Word
mask256 = Addr# -> Ptr Word
forall a. Addr# -> Ptr a
Ptr Addr#
"\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 Word
mask425 = Addr# -> Ptr Word
forall a. Addr# -> Ptr a
Ptr Addr#
"\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 Word
mask377 = Addr# -> Ptr Word
forall a. Addr# -> Ptr a
Ptr Addr#
"\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"#

-- biSqRes256 :: V.Vector Bool
-- biSqRes256 = runST $ do
--     ar <- MV.replicate 256 False
--     let note 257 = return ()
--         note i = MV.unsafeWrite ar i True >> note (i+16)
--     MV.unsafeWrite ar 0 True
--     MV.unsafeWrite ar 16 True
--     note 1
--     V.unsafeFreeze ar

-- biSqRes425 :: V.Vector Bool
-- biSqRes425 = runST $ do
--     ar <- MV.replicate 425 False
--     let note 154 = return ()
--         note i = MV.unsafeWrite ar ((i*i*i*i) `rem` 425) True >> note (i+1)
--     note 0
--     V.unsafeFreeze ar

-- biSqRes377 :: V.Vector Bool
-- biSqRes377 = runST $ do
--     ar <- MV.replicate 377 False
--     let note 144 = return ()
--         note i = MV.unsafeWrite ar ((i*i*i*i) `rem` 377) True >> note (i+1)
--     note 0
--     V.unsafeFreeze ar

-----------------------------------------------------------------------------
-- Specialisations for Int, Word, and Integer

biSqRootIntLimit :: Int
biSqRootIntLimit :: Int
biSqRootIntLimit = if Word -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
64 then Int
55108 else Int
215

biSqrtInt :: Int -> Int
biSqrtInt :: Int -> Int
biSqrtInt Int
0 = Int
0
biSqrtInt Int
n
    | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
biSqRootIntLimit = Int
biSqRootIntLimit
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r4    = Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1
    | Bool
otherwise = Int
r
      where
        x :: Double
        x :: Double
x = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
        -- timed faster than x**0.25, never too small
        r :: Int
r = Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
truncate (Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double
forall a. Floating a => a -> a
sqrt Double
x))
        r2 :: Int
r2 = Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
r
        r4 :: Int
r4 = Int
r2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
r2

biSqRootWordLimit :: Word
biSqRootWordLimit :: Word
biSqRootWordLimit = if Word -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
64 then Word
65535 else Word
255

biSqrtWord :: Word -> Word
biSqrtWord :: Word -> Word
biSqrtWord Word
0 = Word
0
biSqrtWord Word
n
    | Word
r Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
> Word
biSqRootWordLimit = Word
biSqRootWordLimit
    | Word
n Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Word
r4    = Word
rWord -> Word -> Word
forall a. Num a => a -> a -> a
-Word
1
    | Bool
otherwise = Word
r
      where
        x :: Double
        x :: Double
x = Word -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n
        r :: Word
r = Double -> Word
forall a b. (RealFrac a, Integral b) => a -> b
truncate (Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double
forall a. Floating a => a -> a
sqrt Double
x))
        r2 :: Word
r2 = Word
rWord -> Word -> Word
forall a. Num a => a -> a -> a
*Word
r
        r4 :: Word
r4 = Word
r2Word -> Word -> Word
forall a. Num a => a -> a -> a
*Word
r2

biSqrtIgr :: Integer -> Integer
biSqrtIgr :: Integer -> Integer
biSqrtIgr Integer
0 = Integer
0
biSqrtIgr Integer
n = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
newton4 Integer
n (Integer -> Integer
forall a. Integral a => a -> a
approxBiSqrt Integer
n)