{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}
{-# OPTIONS_GHC -O2 #-}
{-# OPTIONS_GHC -optc-O3 #-}
module Data.FastDigits
( digits
, undigits
, digitsUnsigned
) where
import GHC.Exts
import GHC.Integer.GMP.Internals
import GHC.Natural
import Unsafe.Coerce
import Data.FastDigits.Internal
digitsNatural :: GmpLimb# -> BigNat -> [Word]
digitsNatural base = f
where
f n
| isZeroBigNat n = []
| otherwise = let (# q, r #) = n `quotRemBigNatWord` base in
W# r : f q
digitsWord :: Word# -> Word# -> [Word]
digitsWord 2## = g
where
g :: Word# -> [Word]
g 0## = []
g n = W# (n `and#` 1##) : g (n `uncheckedShiftRL#` 1#)
digitsWord base = f
where
f :: Word# -> [Word]
f 0## = []
f n = let (# q, r #) = n `quotRemWord#` base in
W# r : f q
digitsWordL :: Word# -> Word# -> Word# -> (# [Word], Word# #)
digitsWordL 2## power = g
where
g :: Word# -> (# [Word], Word# #)
g 0## = (# [], power #)
g n = (# W# (n `and#` 1##) : fq, lq `minusWord#` 1## #)
where
(# fq, lq #) = g (n `uncheckedShiftRL#` 1#)
digitsWordL base power = f
where
f :: Word# -> (# [Word], Word# #)
f 0## = (# [], power #)
f n = (# W# r : fq, lq `minusWord#` 1## #)
where
(# q, r #) = n `quotRemWord#` base
(# fq, lq #) = f q
digitsNatural' :: Word# -> Word# -> Word# -> BigNat -> [Word]
digitsNatural' base power poweredBase = f
where
f :: BigNat -> [Word]
f n = let (# q, r #) = n `quotRemBigNatWord` poweredBase in
if isZeroBigNat q
then digitsWord base r
else let (# fr, lr #) = digitsWordL base power r in
fr ++ replicate (I# (unsafeCoerce# lr)) 0 ++ f q
padUpTo :: Int -> [Word] -> [Word]
padUpTo !n [] = replicate n 0
padUpTo !n (x : xs) = x : padUpTo (n - 1) xs
digitsUnsigned
:: Word
-> Natural
-> [Word]
digitsUnsigned (W# base) (NatS# n) = digitsWord base n
digitsUnsigned (W# base) (NatJ# n)
| halfSize <- sizeofBigNat# n `iShiftRL#` 1#
, isTrue# (halfSize ># 128#)
= let pow = I# (word2Int# power *# halfSize) in
let (nHi, nLo) = NatJ# n `quotRem` (NatS# poweredBase ^ (I# halfSize)) in
padUpTo pow (digitsUnsigned (W# base) nLo) ++ digitsUnsigned (W# base) nHi
| otherwise
= case power of
1## -> digitsNatural base n
_ -> digitsNatural' base power poweredBase n
where
(# power, poweredBase #) = selectPower base
digits
:: Int
-> Integer
-> [Int]
digits base n
| base < 2 = error "Base must be > 1"
| n < 0 = error "Number must be non-negative"
| otherwise = unsafeCoerce
$ digitsUnsigned (unsafeCoerce base) (unsafeCoerce n)
undigits :: (Integral a, Integral b)
=> a
-> [b]
-> Integer
undigits base' = foldr (\d acc -> acc * base + toInteger d) 0
where
base = toInteger base'
{-# SPECIALIZE undigits :: Word -> [Word] -> Integer #-}
{-# SPECIALIZE undigits :: Int -> [Int] -> Integer #-}
{-# SPECIALIZE undigits :: Integer -> [Integer] -> Integer #-}