{-|
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 GHC.Integer.GMP.Internals (GmpLimb#, BigNat, quotRemBigNatWord, isZeroBigNat, sizeofBigNat#)
import Data.FastDigits.Internal (selectPower)
import GHC.Natural (Natural(..))

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

digitsWord :: Word# -> Word# -> [Word]
digitsWord :: GmpLimb# -> GmpLimb# -> [Word]
digitsWord GmpLimb#
2## = GmpLimb# -> [Word]
g
  where
    g :: Word# -> [Word]
    g :: GmpLimb# -> [Word]
g GmpLimb#
0## = []
    g GmpLimb#
n   = GmpLimb# -> Word
W# (GmpLimb#
n GmpLimb# -> GmpLimb# -> GmpLimb#
`and#` GmpLimb#
1##) Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: GmpLimb# -> [Word]
g (GmpLimb#
n GmpLimb# -> Int# -> GmpLimb#
`uncheckedShiftRL#` Int#
1#)
digitsWord GmpLimb#
10##
  | Word -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
64
  = GmpLimb# -> [Word]
f
  where
    f :: Word# -> [Word]
    f :: GmpLimb# -> [Word]
f GmpLimb#
0## = []
    f GmpLimb#
n   = let !(# GmpLimb#
hi, GmpLimb#
_ #) = GmpLimb#
n GmpLimb# -> GmpLimb# -> (# GmpLimb#, GmpLimb# #)
`timesWord2#` GmpLimb#
14757395258967641293## in
            let q :: GmpLimb#
q = GmpLimb#
hi GmpLimb# -> Int# -> GmpLimb#
`uncheckedShiftRL#` Int#
3# in
            let r :: GmpLimb#
r = GmpLimb#
n GmpLimb# -> GmpLimb# -> GmpLimb#
`minusWord#` (GmpLimb#
q GmpLimb# -> GmpLimb# -> GmpLimb#
`timesWord#` GmpLimb#
10##) in
            GmpLimb# -> Word
W# GmpLimb#
r Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: GmpLimb# -> [Word]
f GmpLimb#
q
digitsWord GmpLimb#
base = GmpLimb# -> [Word]
f
  where
    f :: Word# -> [Word]
    f :: GmpLimb# -> [Word]
f GmpLimb#
0## = []
    f GmpLimb#
n   = let !(# GmpLimb#
q, GmpLimb#
r #) = GmpLimb#
n GmpLimb# -> GmpLimb# -> (# GmpLimb#, GmpLimb# #)
`quotRemWord#` GmpLimb#
base in
            GmpLimb# -> Word
W# GmpLimb#
r Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: GmpLimb# -> [Word]
f GmpLimb#
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 :: GmpLimb# -> GmpLimb# -> GmpLimb# -> (# [Word], GmpLimb# #)
digitsWordL GmpLimb#
2## GmpLimb#
power = GmpLimb# -> (# [Word], GmpLimb# #)
g
  where
    g :: Word# -> (# [Word], Word# #)
    g :: GmpLimb# -> (# [Word], GmpLimb# #)
g GmpLimb#
0## = (# [], GmpLimb#
power #)
    g GmpLimb#
n   = (# GmpLimb# -> Word
W# (GmpLimb#
n GmpLimb# -> GmpLimb# -> GmpLimb#
`and#` GmpLimb#
1##) Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: [Word]
fq, GmpLimb#
lq GmpLimb# -> GmpLimb# -> GmpLimb#
`minusWord#` GmpLimb#
1## #)
      where
        !(# [Word]
fq, GmpLimb#
lq #) = GmpLimb# -> (# [Word], GmpLimb# #)
g (GmpLimb#
n GmpLimb# -> Int# -> GmpLimb#
`uncheckedShiftRL#` Int#
1#)
digitsWordL GmpLimb#
10## GmpLimb#
power
  | Word -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
64
  = GmpLimb# -> (# [Word], GmpLimb# #)
f
  where
    f :: Word# -> (# [Word], Word# #)
    f :: GmpLimb# -> (# [Word], GmpLimb# #)
f GmpLimb#
0## = (# [], GmpLimb#
power #)
    f GmpLimb#
n   = (# GmpLimb# -> Word
W# GmpLimb#
r Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: [Word]
fq, GmpLimb#
lq GmpLimb# -> GmpLimb# -> GmpLimb#
`minusWord#` GmpLimb#
1## #)
      where
        !(# GmpLimb#
hi, GmpLimb#
_ #) = GmpLimb#
n GmpLimb# -> GmpLimb# -> (# GmpLimb#, GmpLimb# #)
`timesWord2#` GmpLimb#
14757395258967641293##
        q :: GmpLimb#
q = GmpLimb#
hi GmpLimb# -> Int# -> GmpLimb#
`uncheckedShiftRL#` Int#
3#
        r :: GmpLimb#
r = GmpLimb#
n GmpLimb# -> GmpLimb# -> GmpLimb#
`minusWord#` (GmpLimb#
q GmpLimb# -> GmpLimb# -> GmpLimb#
`timesWord#` GmpLimb#
10##)
        !(# [Word]
fq, GmpLimb#
lq #) = GmpLimb# -> (# [Word], GmpLimb# #)
f GmpLimb#
q
digitsWordL GmpLimb#
base GmpLimb#
power = GmpLimb# -> (# [Word], GmpLimb# #)
f
  where
    f :: Word# -> (# [Word], Word# #)
    f :: GmpLimb# -> (# [Word], GmpLimb# #)
f GmpLimb#
0## = (# [], GmpLimb#
power #)
    f GmpLimb#
n   = (# GmpLimb# -> Word
W# GmpLimb#
r Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: [Word]
fq, GmpLimb#
lq GmpLimb# -> GmpLimb# -> GmpLimb#
`minusWord#` GmpLimb#
1## #)
      where
        !(#  GmpLimb#
q,  GmpLimb#
r #) = GmpLimb#
n GmpLimb# -> GmpLimb# -> (# GmpLimb#, GmpLimb# #)
`quotRemWord#` GmpLimb#
base
        !(# [Word]
fq, GmpLimb#
lq #) = GmpLimb# -> (# [Word], GmpLimb# #)
f GmpLimb#
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' :: GmpLimb# -> GmpLimb# -> GmpLimb# -> BigNat -> [Word]
digitsNatural' GmpLimb#
base GmpLimb#
power GmpLimb#
poweredBase = BigNat -> [Word]
f
  where
    f :: BigNat -> [Word]
    f :: BigNat -> [Word]
f BigNat
n = let !(# BigNat
q, GmpLimb#
r #) = BigNat
n BigNat -> GmpLimb# -> (# BigNat, GmpLimb# #)
`quotRemBigNatWord` GmpLimb#
poweredBase in
      if BigNat -> Bool
isZeroBigNat BigNat
q
        then GmpLimb# -> GmpLimb# -> [Word]
digitsWord GmpLimb#
base GmpLimb#
r
        else let !(# [Word]
fr, GmpLimb#
lr #) = GmpLimb# -> GmpLimb# -> GmpLimb# -> (# [Word], GmpLimb# #)
digitsWordL GmpLimb#
base GmpLimb#
power GmpLimb#
r in
          [Word]
fr [Word] -> [Word] -> [Word]
forall a. [a] -> [a] -> [a]
++ Int -> Word -> [Word]
forall a. Int -> a -> [a]
replicate (Int# -> Int
I# (GmpLimb# -> Int#
word2Int# GmpLimb#
lr)) Word
0 [Word] -> [Word] -> [Word]
forall a. [a] -> [a] -> [a]
++ BigNat -> [Word]
f BigNat
q

padUpTo :: Int -> [Word] -> [Word]
padUpTo :: Int -> [Word] -> [Word]
padUpTo !Int
n [] = Int -> Word -> [Word]
forall a. Int -> a -> [a]
replicate Int
n Word
0
padUpTo !Int
n (Word
x : [Word]
xs) = Word
x Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: Int -> [Word] -> [Word]
padUpTo (Int
n Int -> Int -> Int
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# GmpLimb#
base) (NatS# GmpLimb#
n) = GmpLimb# -> GmpLimb# -> [Word]
digitsWord GmpLimb#
base GmpLimb#
n
digitsUnsigned (W# GmpLimb#
base) (NatJ# BigNat
n)
  | Int#
halfSize <- BigNat -> Int#
sizeofBigNat# 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# (GmpLimb# -> Int#
word2Int# GmpLimb#
power Int# -> Int# -> Int#
*# Int#
halfSize) in
    let (Natural
nHi, Natural
nLo) = BigNat -> Natural
NatJ# BigNat
n Natural -> Natural -> (Natural, Natural)
forall a. Integral a => a -> a -> (a, a)
`quotRem` (GmpLimb# -> Natural
NatS# GmpLimb#
poweredBase Natural -> Int -> Natural
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 (GmpLimb# -> Word
W# GmpLimb#
base) Natural
nLo) [Word] -> [Word] -> [Word]
forall a. [a] -> [a] -> [a]
++ Word -> Natural -> [Word]
digitsUnsigned (GmpLimb# -> Word
W# GmpLimb#
base) Natural
nHi
  | Bool
otherwise
  = case GmpLimb#
power of
    GmpLimb#
1## -> GmpLimb# -> BigNat -> [Word]
digitsNatural GmpLimb#
base BigNat
n
    GmpLimb#
_   -> GmpLimb# -> GmpLimb# -> GmpLimb# -> BigNat -> [Word]
digitsNatural' GmpLimb#
base GmpLimb#
power GmpLimb#
poweredBase BigNat
n
  where
    !(# GmpLimb#
power, GmpLimb#
poweredBase #) = GmpLimb# -> (# GmpLimb#, GmpLimb# #)
selectPower GmpLimb#
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2  = [Char] -> [Int]
forall a. HasCallStack => [Char] -> a
error [Char]
"Base must be > 1"
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0     = [Char] -> [Int]
forall a. HasCallStack => [Char] -> a
error [Char]
"Number must be non-negative"
  | Bool
otherwise = (Word -> Int) -> [Word] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral
              ([Word] -> [Int]) -> [Word] -> [Int]
forall a b. (a -> b) -> a -> b
$ Word -> Natural -> [Word]
digitsUnsigned (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
base) (Integer -> Natural
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 :: a -> [b] -> Integer
undigits a
base' = (b -> Integer -> Integer) -> Integer -> [b] -> Integer
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\b
d Integer
acc -> Integer
acc Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
base Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ b -> Integer
forall a. Integral a => a -> Integer
toInteger b
d) Integer
0
  where
    base :: Integer
base = a -> Integer
forall a. Integral a => a -> Integer
toInteger a
base'
{-# SPECIALIZE undigits :: Word    -> [Word]    -> Integer #-}
{-# SPECIALIZE undigits :: Int     -> [Int]     -> Integer #-}
{-# SPECIALIZE undigits :: Integer -> [Integer] -> Integer #-}