{-# LANGUAGE CPP, MagicHash, UnboxedTuples, BangPatterns #-}
{-# OPTIONS_HADDOCK hide #-}
module Math.NumberTheory.Utils
( shiftToOddCount
, shiftToOdd
, shiftToOdd#
, shiftToOddCount#
, bitCountWord
, bitCountInt
, bitCountWord#
, uncheckedShiftR
, splitOff
, splitOff#
, mergeBy
, recipMod
) where
#include "MachDeps.h"
import GHC.Base
import GHC.Integer
import GHC.Integer.GMP.Internals
import GHC.Natural
import Data.Bits
uncheckedShiftR :: Word -> Int -> Word
uncheckedShiftR (W# w#) (I# i#) = W# (uncheckedShiftRL# w# i#)
{-# RULES
"shiftToOddCount/Int" shiftToOddCount = shiftOCInt
"shiftToOddCount/Word" shiftToOddCount = shiftOCWord
"shiftToOddCount/Integer" shiftToOddCount = shiftOCInteger
"shiftToOddCount/Natural" shiftToOddCount = shiftOCNatural
#-}
{-# INLINE [1] shiftToOddCount #-}
shiftToOddCount :: Integral a => a -> (Int, a)
shiftToOddCount n = case shiftOCInteger (fromIntegral n) of
(z, o) -> (z, fromInteger o)
shiftOCWord :: Word -> (Int, Word)
shiftOCWord (W# w#) = case shiftToOddCount# w# of
(# z# , u# #) -> (I# z#, W# u#)
shiftOCInt :: Int -> (Int, Int)
shiftOCInt (I# i#) = case shiftToOddCount# (int2Word# i#) of
(# z#, u# #) -> (I# z#, I# (word2Int# u#))
shiftOCInteger :: Integer -> (Int, Integer)
shiftOCInteger n@(S# i#) =
case shiftToOddCount# (int2Word# i#) of
(# z#, w# #)
| isTrue# (z# ==# 0#) -> (0, n)
| otherwise -> (I# z#, S# (word2Int# w#))
shiftOCInteger n@(Jp# bn#) = case bigNatZeroCount bn# of
0# -> (0, n)
z# -> (I# z#, n `shiftRInteger` z#)
shiftOCInteger n@(Jn# bn#) = case bigNatZeroCount bn# of
0# -> (0, n)
z# -> (I# z#, n `shiftRInteger` z#)
shiftOCNatural :: Natural -> (Int, Natural)
shiftOCNatural n@(NatS# i#) =
case shiftToOddCount# i# of
(# z#, w# #)
| isTrue# (z# ==# 0#) -> (0, n)
| otherwise -> (I# z#, NatS# w#)
shiftOCNatural n@(NatJ# bn#) = case bigNatZeroCount bn# of
0# -> (0, n)
z# -> (I# z#, NatJ# (bn# `shiftRBigNat` z#))
bigNatZeroCount :: BigNat -> Int#
bigNatZeroCount bn# = count 0# 0#
where
count a# i# =
case indexBigNat# bn# i# of
0## -> count (a# +# WORD_SIZE_IN_BITS#) (i# +# 1#)
w# -> a# +# word2Int# (ctz# w#)
{-# RULES
"shiftToOdd/Int" shiftToOdd = shiftOInt
"shiftToOdd/Word" shiftToOdd = shiftOWord
"shiftToOdd/Integer" shiftToOdd = shiftOInteger
#-}
{-# INLINE [1] shiftToOdd #-}
shiftToOdd :: Integral a => a -> a
shiftToOdd n = fromInteger (shiftOInteger (fromIntegral n))
shiftOInt :: Int -> Int
shiftOInt (I# i#) = I# (word2Int# (shiftToOdd# (int2Word# i#)))
shiftOWord :: Word -> Word
shiftOWord (W# w#) = W# (shiftToOdd# w#)
shiftOInteger :: Integer -> Integer
shiftOInteger (S# i#) = S# (word2Int# (shiftToOdd# (int2Word# i#)))
shiftOInteger n@(Jn# bn#) = case bigNatZeroCount bn# of
0# -> n
z# -> n `shiftRInteger` z#
shiftOInteger n@(Jp# bn#) = case bigNatZeroCount bn# of
0# -> n
z# -> n `shiftRInteger` z#
shiftToOdd# :: Word# -> Word#
shiftToOdd# w# = uncheckedShiftRL# w# (word2Int# (ctz# w#))
shiftToOddCount# :: Word# -> (# Int#, Word# #)
shiftToOddCount# w# = case word2Int# (ctz# w#) of
k# -> (# k#, uncheckedShiftRL# w# k# #)
bitCountWord# :: Word# -> Int#
bitCountWord# w# = case bitCountWord (W# w#) of
I# i# -> i#
bitCountWord :: Word -> Int
bitCountWord = popCount
bitCountInt :: Int -> Int
bitCountInt = popCount
splitOff :: Integral a => a -> a -> (Int, a)
splitOff _ 0 = (0, 0)
splitOff p n = go 0 n
where
go !k m = case m `quotRem` p of
(q, 0) -> go (k + 1) q
_ -> (k, m)
{-# INLINABLE splitOff #-}
splitOff# :: Word# -> Word# -> (# Int#, Word# #)
splitOff# _ 0## = (# 0#, 0## #)
splitOff# p n = go 0# n
where
go k m = case m `quotRemWord#` p of
(# q, 0## #) -> go (k +# 1#) q
_ -> (# k, m #)
{-# INLINABLE splitOff# #-}
mergeBy :: (a -> a -> Ordering) -> [a] -> [a] -> [a]
mergeBy cmp = loop
where
loop [] ys = ys
loop xs [] = xs
loop (x:xs) (y:ys)
= case cmp x y of
GT -> y : loop (x:xs) ys
_ -> x : loop xs (y:ys)
recipMod :: Integer -> Integer -> Maybe Integer
recipMod x m = case recipModInteger (x `mod` m) m of
0 -> Nothing
y -> Just y