module Math.Utils where

import Data.Bits
import GHC.Num.Integer (integerLog2)

{- |
>>> floorSqrt 0
0
>>> floorSqrt 1
1
>>> floorSqrt 2
1
>>> floorSqrt 4
2
>>> floorSqrt (12345 * 12345)
12345
>>> floorSqrt (2^52 + 2^27)
67108864
>>> floorSqrt maxBound
3037000499
>>> floorSqrt (-1)
0
-}
floorSqrt :: Int -> Int
floorSqrt :: Int -> Int
floorSqrt Int
n
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Int
0
  | Bool
otherwise =
      let !k :: Int
k = Int
32 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR (Int -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Int
1
          !x0 :: Int
x0 = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftL Int
1 Int
k
          !x1 :: Int
x1 = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR (Int
x0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR Int
n Int
k) Int
1
       in Int -> Int -> Int
go Int
x0 Int
x1
  where
    go :: Int -> Int -> Int
go !Int
x !Int
x'
      | Int
x' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x = Int -> Int -> Int
go Int
x' (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR (Int
x' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Int
forall a. Integral a => a -> a -> a
quot Int
n Int
x') Int
1)
      | Bool
otherwise = Int
x

{- |
>>> integerFloorSqrt 1
1
>>> integerFloorSqrt 2
1
>>> integerFloorSqrt 4
2
>>> integerFloorSqrt (12345 * 12345)
12345
>>> integerFloorSqrt (-1)
0
-}
integerFloorSqrt :: Integer -> Integer
integerFloorSqrt :: Integer -> Integer
integerFloorSqrt Integer
n | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0 = Integer
0
integerFloorSqrt Integer
n = Integer -> Integer -> Integer
go Integer
x0 Integer
x1
  where
    !k :: Int
k = Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Int) -> Word -> Int
forall a b. (a -> b) -> a -> b
$ Integer -> Word
integerLog2 Integer
n Word -> Int -> Word
forall a. Bits a => a -> Int -> a
!>>. Int
1 Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1
    !x0 :: Integer
x0 = Integer
1 Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
!<<. Int
k
    !x1 :: Integer
x1 = (Integer
x0 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
n Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
!>>. Int
k) Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
!>>. Int
1
    go :: Integer -> Integer -> Integer
go !Integer
x !Integer
x'
      | Integer
x' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
x = Integer -> Integer -> Integer
go Integer
x' ((Integer
x' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
quot Integer
n Integer
x') Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
!>>. Int
1)
      | Bool
otherwise = Integer
x