{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE UnboxedSums #-}
{-# LANGUAGE UnboxedTuples #-}
{-# OPTIONS_GHC -Wno-unused-imports -fobject-code #-}

#include "MachDeps.h"

module Numeric.Floating.IEEE.Internal.IntegerInternals
  ( integerToIntMaybe
  , naturalToWordMaybe
  , unsafeShiftLInteger
  , unsafeShiftRInteger
  , roundingMode
  , countTrailingZerosInteger
  , integerIsPowerOf2
  , integerLog2IsPowerOf2
  ) where
import           Data.Bits
import           GHC.Exts (Int#, Word#, ctz#, int2Word#, plusWord#, quotRemInt#,
                           uncheckedShiftL#, word2Int#, (+#), (-#))
import           GHC.Int (Int (I#))
import           GHC.Word (Word (W#))
import           MyPrelude
import           Numeric.Floating.IEEE.Internal.Base
import           Numeric.Natural
#if defined(MIN_VERSION_ghc_bignum)
import qualified GHC.Num.BigNat
import           GHC.Num.Integer (Integer (IN, IP, IS))
import qualified GHC.Num.Integer
import           GHC.Num.Natural (Natural (NS))
#elif defined(MIN_VERSION_integer_gmp)
import qualified GHC.Integer
import           GHC.Integer.GMP.Internals (Integer (Jn#, Jp#, S#),
                                            indexBigNat#)
import qualified GHC.Integer.Logarithms.Internals
import           GHC.Natural (Natural (NatS#))
#define IN Jn#
#define IP Jp#
#define IS S#
#define NS NatS#
#else
import           Math.NumberTheory.Logarithms (integerLog2')
#endif

-- $setup
-- >>> :m + Data.Int Data.Bits Test.QuickCheck
-- >>> :{
--   -- Workaround for https://github.com/sol/doctest/issues/160:
--   import Numeric.Floating.IEEE.Internal.IntegerInternals
-- :}

integerToIntMaybe :: Integer -> Maybe Int
naturalToWordMaybe :: Natural -> Maybe Word

-- The instance 'Bits Integer' is not very optimized...
unsafeShiftLInteger :: Integer -> Int -> Integer
unsafeShiftRInteger :: Integer -> Int -> Integer

-- |
-- Assumption: @n > 0@, @e >= 0@, and @integerLog2 n >= e@
--
-- Returns @compare (n \`'rem'\` 2^(e+1)) (2^e)@.
roundingMode :: Integer -- ^ @n@
             -> Int -- ^ @e@
             -> Ordering

-- |
-- 'Integer' version of 'countTrailingZeros'.
-- The argument must not be zero.
--
-- prop> \(NonZero x) -> countTrailingZerosInteger (toInteger x) === countTrailingZeros (x :: Int64)
-- >>> countTrailingZerosInteger 7
-- 0
-- >>> countTrailingZerosInteger 8
-- 3
countTrailingZerosInteger :: Integer -> Int

-- |
-- Returns @Just (integerLog2 x)@ if the argument @x@ is a power of 2, and @Nothing@ otherwise.
-- The argument @x@ must be strictly positive.
integerIsPowerOf2 :: Integer -> Maybe Int

-- |
-- Returns @(integerLog2 x, isJust (integerIsPowerOf2 x))@.
-- The argument @x@ must be strictly positive.
integerLog2IsPowerOf2 :: Integer -> (Int, Bool)

#if defined(MIN_VERSION_ghc_bignum) || defined(MIN_VERSION_integer_gmp)

integerToIntMaybe :: Integer -> Maybe Int
integerToIntMaybe (IS Int#
x) = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int# -> Int
I# Int#
x)
integerToIntMaybe Integer
_      = Maybe Int
forall a. Maybe a
Nothing -- relies on Integer's invariant
{-# INLINE [0] integerToIntMaybe #-}

naturalToWordMaybe :: Natural -> Maybe Word
naturalToWordMaybe (NS Word#
x) = Word -> Maybe Word
forall a. a -> Maybe a
Just (Word# -> Word
W# Word#
x)
naturalToWordMaybe Natural
_      = Maybe Word
forall a. Maybe a
Nothing -- relies on Natural's invariant
{-# INLINE [0] naturalToWordMaybe #-}

integerToIntMaybe2 :: Bool -> Integer -> Maybe Int
integerToIntMaybe2 :: Bool -> Integer -> Maybe Int
integerToIntMaybe2 Bool
_ (IS Int#
x) = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int# -> Int
I# Int#
x)
integerToIntMaybe2 Bool
_ Integer
_      = Maybe Int
forall a. Maybe a
Nothing
{-# INLINE [0] integerToIntMaybe2 #-}

naturalToWordMaybe2 :: Bool -> Natural -> Maybe Word
naturalToWordMaybe2 :: Bool -> Natural -> Maybe Word
naturalToWordMaybe2 Bool
_ (NS Word#
x) = Word -> Maybe Word
forall a. a -> Maybe a
Just (Word# -> Word
W# Word#
x)
naturalToWordMaybe2 Bool
_ Natural
_      = Maybe Word
forall a. Maybe a
Nothing
{-# INLINE [0] naturalToWordMaybe2 #-}

minBoundIntAsInteger :: Integer
minBoundIntAsInteger :: Integer
minBoundIntAsInteger = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
forall a. Bounded a => a
minBound :: Int)
{-# INLINE minBoundIntAsInteger #-}

maxBoundIntAsInteger :: Integer
maxBoundIntAsInteger :: Integer
maxBoundIntAsInteger = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
forall a. Bounded a => a
maxBound :: Int)
{-# INLINE maxBoundIntAsInteger #-}

maxBoundWordAsNatural :: Natural
maxBoundWordAsNatural :: Natural
maxBoundWordAsNatural = Word -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word
forall a. Bounded a => a
maxBound :: Word)
{-# INLINE maxBoundWordAsNatural #-}

{-# RULES
"integerToIntMaybe" [~0] forall x.
  integerToIntMaybe x = integerToIntMaybe2 (minBoundIntAsInteger <= x && x <= maxBoundIntAsInteger) x
"integerToIntMaybe2/small" forall x.
  integerToIntMaybe2 True x = Just (fromIntegral x)
"integerToIntMaybe2/large" forall x.
  integerToIntMaybe2 False x = Nothing
"naturalToWordMaybe" [~0] forall x.
  naturalToWordMaybe x = naturalToWordMaybe2 (x <= maxBoundWordAsNatural) x
"naturalToWordIntMaybe2/small" forall x.
  naturalToWordMaybe2 True x = Just (fromIntegral x)
"naturalToWordIntMaybe2/large" forall x.
  naturalToWordMaybe2 False x = Nothing
  #-}

#else

integerToIntMaybe = toIntegralSized
naturalToWordMaybe = toIntegralSized
{-# INLINE integerToIntMaybe #-}
{-# INLINE naturalToWordMaybe #-}

#endif

#if defined(MIN_VERSION_ghc_bignum)

unsafeShiftLInteger :: Integer -> Int -> Integer
unsafeShiftLInteger Integer
x (I# Int#
i) = Integer -> Word# -> Integer
GHC.Num.Integer.integerShiftL# Integer
x (Int# -> Word#
int2Word# Int#
i)
unsafeShiftRInteger :: Integer -> Int -> Integer
unsafeShiftRInteger Integer
x (I# Int#
i) = Integer -> Word# -> Integer
GHC.Num.Integer.integerShiftR# Integer
x (Int# -> Word#
int2Word# Int#
i)

#elif defined(MIN_VERSION_integer_gmp)

unsafeShiftLInteger x (I# i) = GHC.Integer.shiftLInteger x i
unsafeShiftRInteger x (I# i) = GHC.Integer.shiftRInteger x i

#else

unsafeShiftLInteger = unsafeShiftL
unsafeShiftRInteger = unsafeShiftR

#endif

{-# INLINE unsafeShiftLInteger #-}
{-# INLINE unsafeShiftRInteger #-}

#if defined(MIN_VERSION_ghc_bignum) || defined(MIN_VERSION_integer_gmp)

countTrailingZerosInteger# :: Integer -> Word#
countTrailingZerosInteger# :: Integer -> Word#
countTrailingZerosInteger# (IS Int#
x) = Word# -> Word#
ctz# (Int# -> Word#
int2Word# Int#
x)
countTrailingZerosInteger# (IN ByteArray#
bn) = Integer -> Word#
countTrailingZerosInteger# (ByteArray# -> Integer
IP ByteArray#
bn)
countTrailingZerosInteger# (IP ByteArray#
bn) = Int# -> Word# -> Word#
loop Int#
0# Word#
0##
  where
    loop :: Int# -> Word# -> Word#
loop Int#
i Word#
acc =
      let
#if defined(MIN_VERSION_ghc_bignum)
        !bn_i :: Word#
bn_i = ByteArray# -> Int# -> Word#
GHC.Num.BigNat.bigNatIndex# ByteArray#
bn Int#
i -- `i < bigNatSize# bn` must hold
#else
        !bn_i = indexBigNat# bn i -- `i < sizeOfBigNat# bn` must hold
#endif
      in case Word#
bn_i of
           Word#
0## -> Int# -> Word# -> Word#
loop (Int#
i Int# -> Int# -> Int#
+# Int#
1#) (Word#
acc Word# -> Word# -> Word#
`plusWord#` WORD_SIZE_IN_BITS##)
           Word#
w   -> Word#
acc Word# -> Word# -> Word#
`plusWord#` Word# -> Word#
ctz# Word#
w

countTrailingZerosInteger :: Integer -> Int
countTrailingZerosInteger Integer
0 = [Char] -> Int
forall a. HasCallStack => [Char] -> a
error [Char]
"countTrailingZerosInteger: zero"
countTrailingZerosInteger Integer
x = Int# -> Int
I# (Word# -> Int#
word2Int# (Integer -> Word#
countTrailingZerosInteger# Integer
x))
{-# INLINE countTrailingZerosInteger #-}

#else

countTrailingZerosInteger 0 = error "countTrailingZerosInteger: zero"
countTrailingZerosInteger x = integerLog2' (x `xor` (x - 1))
{-# INLINE countTrailingZerosInteger #-}

#endif

#if defined(MIN_VERSION_ghc_bignum)

roundingMode# :: Integer -> Int# -> Ordering
roundingMode# :: Integer -> Int# -> Ordering
roundingMode# (IS Int#
x) Int#
t = let !w :: Word#
w = Int# -> Word#
int2Word# Int#
x
                         in Word -> Word -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Word# -> Word
W# (Word#
w Word# -> Int# -> Word#
`uncheckedShiftL#` (WORD_SIZE_IN_BITS# Word# -> Word
-# 1# -# t))) (W# (1## `uncheckedShiftL#` (WORD_SIZE_IN_BITS# -# 1#)))
roundingMode# (IN ByteArray#
bn) Int#
t = Integer -> Int# -> Ordering
roundingMode# (ByteArray# -> Integer
IP ByteArray#
bn) Int#
t -- unexpected
roundingMode# (IP ByteArray#
bn) Int#
t = case Int#
t Int# -> Int# -> (# Int#, Int# #)
`quotRemInt#` WORD_SIZE_IN_BITS# of
                            -- 0 <= r < WORD_SIZE_IN_BITS
                            (# Int#
s, Int#
r #) -> let !w :: Word#
w = ByteArray# -> Int# -> Word#
GHC.Num.BigNat.bigNatIndex# ByteArray#
bn Int#
s
                                              -- w `shiftL` (WORD_SIZE_IN_BITS - r - 1) vs. 1 `shiftL` (WORD_SIZE_IN_BITS - 1)
                                          in Word -> Word -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Word# -> Word
W# (Word#
w Word# -> Int# -> Word#
`uncheckedShiftL#` (WORD_SIZE_IN_BITS# Word# -> Word
-# 1# -# r))) (W# (1## `uncheckedShiftL#` (WORD_SIZE_IN_BITS# -# 1#)))
                                             Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Int# -> Ordering
loop Int#
s
  where
    loop :: Int# -> Ordering
loop Int#
0# = Ordering
EQ
    loop Int#
i = case ByteArray# -> Int# -> Word#
GHC.Num.BigNat.bigNatIndex# ByteArray#
bn (Int#
i Int# -> Int# -> Int#
-# Int#
1#) of
               Word#
0## -> Int# -> Ordering
loop (Int#
i Int# -> Int# -> Int#
-# Int#
1#)
               Word#
_   -> Ordering
GT

roundingMode :: Integer -> Int -> Ordering
roundingMode Integer
x (I# Int#
t) = Integer -> Int# -> Ordering
roundingMode# Integer
x Int#
t
{-# INLINE roundingMode #-}

integerIsPowerOf2 :: Integer -> Maybe Int
integerIsPowerOf2 Integer
x = case Integer -> (# (# #) | Word# #)
GHC.Num.Integer.integerIsPowerOf2# Integer
x of
                        (# (# #)
_ | #) -> Maybe Int
forall a. Maybe a
Nothing
                        (# | Word#
w #) -> Int -> Maybe Int
forall a. a -> Maybe a
Just (Int# -> Int
I# (Word# -> Int#
word2Int# Word#
w))
{-# INLINE integerIsPowerOf2 #-}

integerLog2IsPowerOf2 :: Integer -> (Int, Bool)
integerLog2IsPowerOf2 Integer
x = case Integer -> (# (# #) | Word# #)
GHC.Num.Integer.integerIsPowerOf2# Integer
x of
                            (# (# #)
_ | #) -> (Int# -> Int
I# (Word# -> Int#
word2Int# (Integer -> Word#
GHC.Num.Integer.integerLog2# Integer
x)), Bool
False)
                            (# | Word#
w #) -> (Int# -> Int
I# (Word# -> Int#
word2Int# Word#
w), Bool
True)
{-# INLINE integerLog2IsPowerOf2 #-}

#elif defined(MIN_VERSION_integer_gmp)

roundingMode x (I# t#) = case GHC.Integer.Logarithms.Internals.roundingMode# x t# of
                           0# -> LT -- round toward zero
                           1# -> EQ -- half
                           _  -> GT -- 2#: round away from zero
{-# INLINE roundingMode #-}

integerIsPowerOf2 x = case GHC.Integer.Logarithms.Internals.integerLog2IsPowerOf2# x of
                        (# l, 0# #) -> Just (I# l)
                        (# _, _ #)  -> Nothing
{-# INLINE integerIsPowerOf2 #-}

integerLog2IsPowerOf2 x = case GHC.Integer.Logarithms.Internals.integerLog2IsPowerOf2# x of
                            (# l, 0# #) -> (I# l, True)
                            (# l, _ #)  -> (I# l, False)
{-# INLINE integerLog2IsPowerOf2 #-}

#else

roundingMode x t = compare (x .&. (bit (t + 1) - 1)) (bit t)
{-# INLINE roundingMode #-}

integerIsPowerOf2 x = if x .&. (x - 1) == 0 then
                        Just (integerLog2' x)
                      else
                        Nothing

integerLog2IsPowerOf2 x = (integerLog2' x, x .&. (x - 1) == 0)
{-# INLINE integerLog2IsPowerOf2 #-}

#endif