{-|
Module      : Data.FastDigits
Description : Integer-to-digits conversion.
Copyright   : (c) Andrew Lelechenko, 2015-2020
License     : GPL-3
Maintainer  : andrew.lelechenko@gmail.com

Convert an integer to digits and back.
This library is both asymptotically (O(n^1.4) vs. O(n^2))
and practically (2x-40x for typical inputs)
faster than "Data.Digits".
-}

{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE MagicHash     #-}
{-# LANGUAGE UnboxedTuples #-}

module Data.FastDigits
  ( digits
  , undigits
  , digitsUnsigned
  ) where

import Data.Bits (finiteBitSize)
import GHC.Exts (Word#, Word(..), uncheckedShiftRL#, and#, timesWord2#, minusWord#, quotRemWord#, timesWord#, Int(..), iShiftRL#, isTrue#, word2Int#, (>#), (*#))
import Data.FastDigits.Internal (selectPower)
import GHC.Natural (Natural(..))
import GHC.Num.BigNat (BigNat(..), bigNatIsZero, bigNatQuotRemWord#, bigNatSize#, BigNat#)

digitsNatural :: Word# -> BigNat# -> [Word]
digitsNatural :: Word# -> BigNat# -> [Word]
digitsNatural Word#
base = BigNat# -> [Word]
f
  where
    f :: BigNat# -> [Word]
f BigNat#
n
      | BigNat# -> Bool
bigNatIsZero BigNat#
n = []
      | Bool
otherwise      = let !(# BigNat#
q, Word#
r #) = BigNat#
n BigNat# -> Word# -> (# BigNat#, Word# #)
`bigNatQuotRemWord#` Word#
base in
                         Word# -> Word
W# Word#
r forall a. a -> [a] -> [a]
: BigNat# -> [Word]
f BigNat#
q

digitsWord :: Word# -> Word# -> [Word]
digitsWord :: Word# -> Word# -> [Word]
digitsWord Word#
2## = Word# -> [Word]
g
  where
    g :: Word# -> [Word]
    g :: Word# -> [Word]
g Word#
0## = []
    g Word#
n   = Word# -> Word
W# (Word#
n Word# -> Word# -> Word#
`and#` Word#
1##) forall a. a -> [a] -> [a]
: Word# -> [Word]
g (Word#
n Word# -> Int# -> Word#
`uncheckedShiftRL#` Int#
1#)
digitsWord Word#
10##
  | forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) forall a. Eq a => a -> a -> Bool
== Int
64
  = Word# -> [Word]
f
  where
    f :: Word# -> [Word]
    f :: Word# -> [Word]
f Word#
0## = []
    f Word#
n   = let !(# Word#
hi, Word#
_ #) = Word#
n Word# -> Word# -> (# Word#, Word# #)
`timesWord2#` Word#
14757395258967641293## in
            let q :: Word#
q = Word#
hi Word# -> Int# -> Word#
`uncheckedShiftRL#` Int#
3# in
            let r :: Word#
r = Word#
n Word# -> Word# -> Word#
`minusWord#` (Word#
q Word# -> Word# -> Word#
`timesWord#` Word#
10##) in
            Word# -> Word
W# Word#
r forall a. a -> [a] -> [a]
: Word# -> [Word]
f Word#
q
digitsWord Word#
base = Word# -> [Word]
f
  where
    f :: Word# -> [Word]
    f :: Word# -> [Word]
f Word#
0## = []
    f Word#
n   = let !(# Word#
q, Word#
r #) = Word#
n Word# -> Word# -> (# Word#, Word# #)
`quotRemWord#` Word#
base in
            Word# -> Word
W# Word#
r forall a. a -> [a] -> [a]
: Word# -> [Word]
f Word#
q

-- | For a given base and expected length of list of digits
--   return the list of digits and padding till expected length.
digitsWordL :: Word# -> Word# -> Word# -> (# [Word], Word# #)
digitsWordL :: Word# -> Word# -> Word# -> (# [Word], Word# #)
digitsWordL Word#
2## Word#
power = Word# -> (# [Word], Word# #)
g
  where
    g :: Word# -> (# [Word], Word# #)
    g :: Word# -> (# [Word], Word# #)
g Word#
0## = (# [], Word#
power #)
    g Word#
n   = (# Word# -> Word
W# (Word#
n Word# -> Word# -> Word#
`and#` Word#
1##) forall a. a -> [a] -> [a]
: [Word]
fq, Word#
lq Word# -> Word# -> Word#
`minusWord#` Word#
1## #)
      where
        !(# [Word]
fq, Word#
lq #) = Word# -> (# [Word], Word# #)
g (Word#
n Word# -> Int# -> Word#
`uncheckedShiftRL#` Int#
1#)
digitsWordL Word#
10## Word#
power
  | forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) forall a. Eq a => a -> a -> Bool
== Int
64
  = Word# -> (# [Word], Word# #)
f
  where
    f :: Word# -> (# [Word], Word# #)
    f :: Word# -> (# [Word], Word# #)
f Word#
0## = (# [], Word#
power #)
    f Word#
n   = (# Word# -> Word
W# Word#
r forall a. a -> [a] -> [a]
: [Word]
fq, Word#
lq Word# -> Word# -> Word#
`minusWord#` Word#
1## #)
      where
        !(# Word#
hi, Word#
_ #) = Word#
n Word# -> Word# -> (# Word#, Word# #)
`timesWord2#` Word#
14757395258967641293##
        q :: Word#
q = Word#
hi Word# -> Int# -> Word#
`uncheckedShiftRL#` Int#
3#
        r :: Word#
r = Word#
n Word# -> Word# -> Word#
`minusWord#` (Word#
q Word# -> Word# -> Word#
`timesWord#` Word#
10##)
        !(# [Word]
fq, Word#
lq #) = Word# -> (# [Word], Word# #)
f Word#
q
digitsWordL Word#
base Word#
power = Word# -> (# [Word], Word# #)
f
  where
    f :: Word# -> (# [Word], Word# #)
    f :: Word# -> (# [Word], Word# #)
f Word#
0## = (# [], Word#
power #)
    f Word#
n   = (# Word# -> Word
W# Word#
r forall a. a -> [a] -> [a]
: [Word]
fq, Word#
lq Word# -> Word# -> Word#
`minusWord#` Word#
1## #)
      where
        !(#  Word#
q,  Word#
r #) = Word#
n Word# -> Word# -> (# Word#, Word# #)
`quotRemWord#` Word#
base
        !(# [Word]
fq, Word#
lq #) = Word# -> (# [Word], Word# #)
f Word#
q

-- | For a given base, power and precalculated base^power
--   take an integer and return the list of its digits.
digitsNatural' :: Word# -> Word# -> Word# -> BigNat# -> [Word]
digitsNatural' :: Word# -> Word# -> Word# -> BigNat# -> [Word]
digitsNatural' Word#
base Word#
power Word#
poweredBase = BigNat# -> [Word]
f
  where
    f :: BigNat# -> [Word]
    f :: BigNat# -> [Word]
f BigNat#
n = let !(# BigNat#
q, Word#
r #) = BigNat#
n BigNat# -> Word# -> (# BigNat#, Word# #)
`bigNatQuotRemWord#` Word#
poweredBase in
      if BigNat# -> Bool
bigNatIsZero BigNat#
q
        then Word# -> Word# -> [Word]
digitsWord Word#
base Word#
r
        else let !(# [Word]
fr, Word#
lr #) = Word# -> Word# -> Word# -> (# [Word], Word# #)
digitsWordL Word#
base Word#
power Word#
r in
          [Word]
fr forall a. [a] -> [a] -> [a]
++ forall a. Int -> a -> [a]
replicate (Int# -> Int
I# (Word# -> Int#
word2Int# Word#
lr)) Word
0 forall a. [a] -> [a] -> [a]
++ BigNat# -> [Word]
f BigNat#
q

padUpTo :: Int -> [Word] -> [Word]
padUpTo :: Int -> [Word] -> [Word]
padUpTo !Int
n [] = forall a. Int -> a -> [a]
replicate Int
n Word
0
padUpTo !Int
n (Word
x : [Word]
xs) = Word
x forall a. a -> [a] -> [a]
: Int -> [Word] -> [Word]
padUpTo (Int
n forall a. Num a => a -> a -> a
- Int
1) [Word]
xs

-- | Return digits of a non-negative number in reverse order.
digitsUnsigned
  :: Word    -- ^ Precondition that base is ≥2 is not checked
  -> Natural
  -> [Word]
digitsUnsigned :: Word -> Natural -> [Word]
digitsUnsigned (W# Word#
base) (NatS# Word#
n) = Word# -> Word# -> [Word]
digitsWord Word#
base Word#
n
digitsUnsigned (W# Word#
base) (NatJ# n :: BigNat
n@(BN# BigNat#
n#))
  | Int#
halfSize <- BigNat# -> Int#
bigNatSize# BigNat#
n# Int# -> Int# -> Int#
`iShiftRL#` Int#
1#
  , Int# -> Bool
isTrue# (Int#
halfSize Int# -> Int# -> Int#
># Int#
128#)
  = let pow :: Int
pow = Int# -> Int
I# (Word# -> Int#
word2Int# Word#
power Int# -> Int# -> Int#
*# Int#
halfSize) in
    let (Natural
nHi, Natural
nLo) = BigNat -> Natural
NatJ# BigNat
n forall a. Integral a => a -> a -> (a, a)
`quotRem` (Word# -> Natural
NatS# Word#
poweredBase forall a b. (Num a, Integral b) => a -> b -> a
^ (Int# -> Int
I# Int#
halfSize)) in
    Int -> [Word] -> [Word]
padUpTo Int
pow (Word -> Natural -> [Word]
digitsUnsigned (Word# -> Word
W# Word#
base) Natural
nLo) forall a. [a] -> [a] -> [a]
++ Word -> Natural -> [Word]
digitsUnsigned (Word# -> Word
W# Word#
base) Natural
nHi
  | Bool
otherwise
  = case Word#
power of
    Word#
1## -> Word# -> BigNat# -> [Word]
digitsNatural Word#
base BigNat#
n#
    Word#
_   -> Word# -> Word# -> Word# -> BigNat# -> [Word]
digitsNatural' Word#
base Word#
power Word#
poweredBase BigNat#
n#
  where
    !(# Word#
power, Word#
poweredBase #) = Word# -> (# Word#, Word# #)
selectPower Word#
base

-- | Return digits of a non-negative number in reverse order.
--   Throw an error if number is negative or base is below 2.
--
--   > digits 10 123 = [3, 2, 1]
--   > digits 10 0   = []
digits
  :: Int     -- ^ The base to use
  -> Integer -- ^ The number to convert
  -> [Int]   -- ^ Digits in reverse order
digits :: Int -> Integer -> [Int]
digits Int
base Integer
n
  | Int
base forall a. Ord a => a -> a -> Bool
< Int
2  = forall a. HasCallStack => [Char] -> a
error [Char]
"Base must be > 1"
  | Integer
n forall a. Ord a => a -> a -> Bool
< Integer
0     = forall a. HasCallStack => [Char] -> a
error [Char]
"Number must be non-negative"
  | Bool
otherwise = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (Integral a, Num b) => a -> b
fromIntegral
              forall a b. (a -> b) -> a -> b
$ Word -> Natural -> [Word]
digitsUnsigned (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
base) (forall a. Num a => Integer -> a
fromInteger Integer
n)

-- | Return an integer, built from given digits in reverse order.
--   Condition 0 ≤ digit < base is not checked.
undigits :: (Integral a, Integral b)
         => a       -- ^ The base to use
         -> [b]     -- ^ The list of digits to convert
         -> Integer
undigits :: forall a b. (Integral a, Integral b) => a -> [b] -> Integer
undigits a
base' = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\b
d Integer
acc -> Integer
acc forall a. Num a => a -> a -> a
* Integer
base forall a. Num a => a -> a -> a
+ forall a. Integral a => a -> Integer
toInteger b
d) Integer
0
  where
    base :: Integer
base = forall a. Integral a => a -> Integer
toInteger a
base'
{-# SPECIALIZE undigits :: Word    -> [Word]    -> Integer #-}
{-# SPECIALIZE undigits :: Int     -> [Int]     -> Integer #-}
{-# SPECIALIZE undigits :: Integer -> [Integer] -> Integer #-}