{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -ddump-simpl -dsuppress-all -ddump-to-file #-}
-- | The naive left fold to convert digits to integer is quadratic
-- as multiplying (big) 'Integer's is not a constant time operation.
--
-- This module provides sub-quadratic algorithm for conversion of 'Text'
-- or 'ByteString' into 'Integer'.
--
-- For example for a text of 262144 9 digits, fold implementation
-- takes 1.5 seconds, and 'textToInteger' just 26 milliseconds on my machine.
-- Difference is already noticeable around 100-200 digits.
--
-- In particular 'read' is correct (i.e. faster) than @List.foldl'@ (better complexity),
-- 'stringToInteger' is a bit faster than 'read' (same complexity, lower coeffcient).
--
module Data.Integer.Conversion (
    textToInteger,
    byteStringToInteger,
    stringToInteger,
    stringToIntegerWithLen,
) where

import Control.Monad.ST     (ST, runST)
import Data.ByteString      (ByteString)
import Data.Char            (ord)
import Data.Primitive.Array (MutableArray, newArray, readArray, writeArray)
import Data.Text.Internal   (Text (..))
import Data.Word            (Word8)

import qualified Data.ByteString as BS
import qualified Data.List       as L
import qualified Data.Text       as T

-- $setup
-- >>> :set -XOverloadedStrings

-------------------------------------------------------------------------------
-- Text
-------------------------------------------------------------------------------

-- | Convert 'Text' to 'Integer'.
--
-- Semantically same as @T.foldl' (\acc c -> acc * 10 + toInteger (ord c - 48)) 0@,
-- but this is more efficient.
--
-- >>> textToInteger "123456789"
-- 123456789
--
-- For non-decimal inputs some nonsense is calculated
--
-- >>> textToInteger "foobar"
-- 6098556
--
textToInteger :: Text -> Integer
textToInteger :: Text -> Integer
textToInteger t :: Text
t@(Text Array
_arr Int
_off Int
len)
    -- len >= 20000 = algorithmL 10 (T.length t) [ toInteger (ord c - 48) | c <- T.unpack t ]
    | Int
len forall a. Ord a => a -> a -> Bool
>= Int
40    = Text -> Integer
complexTextToInteger Text
t
    | Bool
otherwise    = Text -> Integer
simpleTextToInteger Text
t

simpleTextToInteger :: Text -> Integer
simpleTextToInteger :: Text -> Integer
simpleTextToInteger = forall a. (a -> Char -> a) -> a -> Text -> a
T.foldl' (\Integer
acc Char
c -> Integer
acc forall a. Num a => a -> a -> a
* Integer
10 forall a. Num a => a -> a -> a
+ Char -> Integer
fromChar Char
c) Integer
0

-- Text doesn't have cheap length:
--
-- * We can (over)estimate the size of the needed buffer by the length of text's underlying bytearray.
-- * As we don't know whether the length is even or odd, we cannot do the first pass,
--   so we just copy the contents of given Text as is first.
--
complexTextToInteger :: Text -> Integer
complexTextToInteger :: Text -> Integer
complexTextToInteger t0 :: Text
t0@(Text Array
_ Int
_ Int
len) = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
    MutableArray s Integer
arr <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
len Integer
integer0 -- we overestimate the size here
    forall s. MutableArray s Integer -> Text -> Int -> ST s Integer
loop MutableArray s Integer
arr Text
t0 Int
0
  where
    loop :: MutableArray s Integer -> Text -> Int -> ST s Integer
    loop :: forall s. MutableArray s Integer -> Text -> Int -> ST s Integer
loop !MutableArray s Integer
arr !Text
t !Int
o = case Text -> Maybe (Char, Text)
T.uncons Text
t of
        Just (Char
c, Text
t') -> do
            forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s Integer
arr Int
o forall a b. (a -> b) -> a -> b
$! Char -> Integer
fromChar Char
c
            forall s. MutableArray s Integer -> Text -> Int -> ST s Integer
loop MutableArray s Integer
arr Text
t' (Int
o forall a. Num a => a -> a -> a
+ Int
1)
        Maybe (Char, Text)
Nothing -> forall s. MutableArray s Integer -> Int -> Integer -> ST s Integer
algorithm MutableArray s Integer
arr Int
o Integer
10

fromChar :: Char -> Integer
fromChar :: Char -> Integer
fromChar Char
c = forall a. Integral a => a -> Integer
toInteger (Char -> Int
ord Char
c forall a. Num a => a -> a -> a
- Int
48 :: Int)
{-# INLINE fromChar #-}

-------------------------------------------------------------------------------
-- ByteString
-------------------------------------------------------------------------------

-- | Convert 'ByteString' to 'Integer'.
--
-- Semantically same as @BS.foldl' (\acc c -> acc * 10 + toInteger c - 48) 0@,
-- but this is more efficient.
--
-- >>> byteStringToInteger "123456789"
-- 123456789
--
-- For non-decimal inputs some nonsense is calculated
--
-- >>> byteStringToInteger "foobar"
-- 6098556
--
byteStringToInteger :: ByteString -> Integer
byteStringToInteger :: ByteString -> Integer
byteStringToInteger ByteString
bs
    -- len >= 20000 = algorithmL 10 len [ toInteger w - 48 | w <- BS.unpack bs ]
    | Int
len forall a. Ord a => a -> a -> Bool
>= Int
40    = Int -> ByteString -> Integer
complexByteStringToInteger Int
len ByteString
bs
    | Bool
otherwise    = ByteString -> Integer
simpleByteStringToInteger ByteString
bs
  where
    !len :: Int
len = ByteString -> Int
BS.length ByteString
bs

simpleByteStringToInteger :: BS.ByteString -> Integer
simpleByteStringToInteger :: ByteString -> Integer
simpleByteStringToInteger = forall a. (a -> Word8 -> a) -> a -> ByteString -> a
BS.foldl' (\Integer
acc Word8
w -> Integer
acc forall a. Num a => a -> a -> a
* Integer
10 forall a. Num a => a -> a -> a
+ Word8 -> Integer
fromWord8 Word8
w) Integer
0

complexByteStringToInteger :: Int -> BS.ByteString -> Integer
complexByteStringToInteger :: Int -> ByteString -> Integer
complexByteStringToInteger Int
len ByteString
bs = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
    MutableArray s Integer
arr <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
len' Integer
0

    if forall a. Integral a => a -> Bool
even Int
len
    then do
        forall s. MutableArray s Integer -> Int -> Int -> ST s Integer
loop MutableArray s Integer
arr Int
0 Int
0
    else do
        forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s Integer
arr Int
0 forall a b. (a -> b) -> a -> b
$! ByteString -> Int -> Integer
indexBS ByteString
bs Int
0
        forall s. MutableArray s Integer -> Int -> Int -> ST s Integer
loop MutableArray s Integer
arr Int
1 Int
1
  where
    len' :: Int
len' = (Int
len forall a. Num a => a -> a -> a
+ Int
1) forall a. Integral a => a -> a -> a
`div` Int
2

    loop :: MutableArray s Integer -> Int -> Int -> ST s Integer
    loop :: forall s. MutableArray s Integer -> Int -> Int -> ST s Integer
loop !MutableArray s Integer
arr !Int
i !Int
o | Int
i forall a. Ord a => a -> a -> Bool
< Int
len = do
        forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s Integer
arr Int
o forall a b. (a -> b) -> a -> b
$! ByteString -> Int -> Integer
indexBS ByteString
bs Int
i forall a. Num a => a -> a -> a
* Integer
10 forall a. Num a => a -> a -> a
+ ByteString -> Int -> Integer
indexBS ByteString
bs (Int
i forall a. Num a => a -> a -> a
+ Int
1)
        forall s. MutableArray s Integer -> Int -> Int -> ST s Integer
loop MutableArray s Integer
arr (Int
i forall a. Num a => a -> a -> a
+ Int
2) (Int
o forall a. Num a => a -> a -> a
+ Int
1)
    loop MutableArray s Integer
arr Int
_ Int
_ = forall s. MutableArray s Integer -> Int -> Integer -> ST s Integer
algorithm MutableArray s Integer
arr Int
len' Integer
100

indexBS :: BS.ByteString -> Int -> Integer
indexBS :: ByteString -> Int -> Integer
indexBS ByteString
bs Int
i = Word8 -> Integer
fromWord8 (HasCallStack => ByteString -> Int -> Word8
BS.index ByteString
bs Int
i)
{-# INLINE indexBS #-}

fromWord8 :: Word8 -> Integer
fromWord8 :: Word8 -> Integer
fromWord8 Word8
w = forall a. Integral a => a -> Integer
toInteger (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
w forall a. Num a => a -> a -> a
- Int
48 :: Int)
{-# INLINE fromWord8 #-}

-------------------------------------------------------------------------------
-- String
-------------------------------------------------------------------------------

-- | Convert 'String' to 'Integer'.
--
-- Semantically same as @List.foldl' (\acc c -> acc * 10 + toInteger c - 48) 0@,
-- but this is more efficient.
--
-- >>> stringToInteger "123456789"
-- 123456789
--
-- For non-decimal inputs some nonsense is calculated
--
-- >>> stringToInteger "foobar"
-- 6098556
--
stringToInteger :: String -> Integer
stringToInteger :: String -> Integer
stringToInteger String
str = String -> Int -> Integer
stringToIntegerWithLen String
str (forall (t :: * -> *) a. Foldable t => t a -> Int
length String
str)

-- | Convert 'String' to 'Integer' when you know the length beforehand.
--
-- >>> stringToIntegerWithLen "123" 3
-- 123
--
-- If the length is wrong, you may get wrong results.
-- (Simple algorithm is used for short strings).
--
-- >>> stringToIntegerWithLen (replicate 40 '0' ++ "123") 45
-- 12300
--
-- >>> stringToIntegerWithLen (replicate 40 '0' ++ "123") 44
-- 1200
--
-- >>> stringToIntegerWithLen (replicate 40 '0' ++ "123") 42
-- 12
--
stringToIntegerWithLen :: String -> Int -> Integer
stringToIntegerWithLen :: String -> Int -> Integer
stringToIntegerWithLen String
str Int
len
    | Int
len forall a. Ord a => a -> a -> Bool
>= Int
40    = Int -> String -> Integer
complexStringToInteger Int
len String
str
    | Bool
otherwise    = String -> Integer
simpleStringToInteger String
str

simpleStringToInteger :: String -> Integer
simpleStringToInteger :: String -> Integer
simpleStringToInteger = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' Integer -> Char -> Integer
step Integer
0 where
  step :: Integer -> Char -> Integer
step Integer
a Char
b = Integer
a forall a. Num a => a -> a -> a
* Integer
10 forall a. Num a => a -> a -> a
+ Char -> Integer
fromChar Char
b

complexStringToInteger :: Int -> String -> Integer
complexStringToInteger :: Int -> String -> Integer
complexStringToInteger Int
len String
str = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
    MutableArray s Integer
arr <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
len' Integer
integer0
    if forall a. Integral a => a -> Bool
even Int
len
    then forall s. MutableArray s Integer -> String -> Int -> ST s Integer
loop MutableArray s Integer
arr String
str     Int
0
    else case String
str of
        []   -> forall (m :: * -> *) a. Monad m => a -> m a
return Integer
integer0 -- cannot happen, length is odd! but could, via stringToIntegerWithLen.
        Char
a:String
bs -> do
            forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s Integer
arr Int
0 forall a b. (a -> b) -> a -> b
$ Char -> Integer
fromChar Char
a
            forall s. MutableArray s Integer -> String -> Int -> ST s Integer
loop MutableArray s Integer
arr String
bs Int
1
  where
    len' :: Int
len' = (Int
len forall a. Num a => a -> a -> a
+ Int
1) forall a. Integral a => a -> a -> a
`div` Int
2

    loop :: MutableArray s Integer -> String -> Int -> ST s Integer
    loop :: forall s. MutableArray s Integer -> String -> Int -> ST s Integer
loop !MutableArray s Integer
arr (Char
a:Char
b:String
cs) !Int
o | Int
o forall a. Ord a => a -> a -> Bool
< Int
len' = do
        forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s Integer
arr Int
o forall a b. (a -> b) -> a -> b
$! Char -> Integer
fromChar Char
a forall a. Num a => a -> a -> a
* Integer
10 forall a. Num a => a -> a -> a
+ Char -> Integer
fromChar Char
b
        forall s. MutableArray s Integer -> String -> Int -> ST s Integer
loop MutableArray s Integer
arr String
cs (Int
o forall a. Num a => a -> a -> a
+ Int
1)
    loop MutableArray s Integer
arr String
_ Int
_ = forall s. MutableArray s Integer -> Int -> Integer -> ST s Integer
algorithm MutableArray s Integer
arr Int
len' Integer
100

-------------------------------------------------------------------------------
-- Algorithm
-------------------------------------------------------------------------------

-- The core of algorithm uses mutable arrays.
-- An alternative (found in e.g. @base@) uses lists.
-- For very big integers (thousands of decimal digits) the difference
-- is small (runtime is dominated by integer multiplication),
-- but for medium sized integers this is slightly faster, as we avoid cons cell allocation.
--
algorithm
    :: forall s. MutableArray s Integer  -- ^ working buffer
    -> Int                               -- ^ buffer size
    -> Integer                           -- ^ base
    -> ST s Integer
algorithm :: forall s. MutableArray s Integer -> Int -> Integer -> ST s Integer
algorithm !MutableArray s Integer
arr !Int
len !Integer
base
    | Int
len forall a. Ord a => a -> a -> Bool
<= Int
40 = Integer -> Int -> ST s Integer
finish Integer
0 Int
0
    | forall a. Integral a => a -> Bool
even Int
len  = Int -> Int -> ST s Integer
loop Int
0 Int
0
    | Bool
otherwise = Int -> Int -> ST s Integer
loop Int
1 Int
1
  where
    loop :: Int -> Int -> ST s Integer
    loop :: Int -> Int -> ST s Integer
loop !Int
i !Int
o | Int
i forall a. Ord a => a -> a -> Bool
< Int
len = do
        -- read at i, i +1
        Integer
a <- forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray s Integer
arr Int
i
        Integer
b <- forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray s Integer
arr (Int
i forall a. Num a => a -> a -> a
+ Int
1)

        -- rewrite with constant to release memory
        forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s Integer
arr Int
i       Integer
integer0
        forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s Integer
arr (Int
i forall a. Num a => a -> a -> a
+ Int
1) Integer
integer0

        -- write at o
        forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s Integer
arr Int
o forall a b. (a -> b) -> a -> b
$! Integer
a forall a. Num a => a -> a -> a
* Integer
base forall a. Num a => a -> a -> a
+ Integer
b

        -- continue
        Int -> Int -> ST s Integer
loop (Int
i forall a. Num a => a -> a -> a
+ Int
2) (Int
o forall a. Num a => a -> a -> a
+ Int
1)

    loop Int
_ Int
_ = forall s. MutableArray s Integer -> Int -> Integer -> ST s Integer
algorithm MutableArray s Integer
arr Int
len' Integer
base'
      where
        !base' :: Integer
base' = Integer
base forall a. Num a => a -> a -> a
* Integer
base
        !len' :: Int
len'  = (Int
len forall a. Num a => a -> a -> a
+ Int
1) forall a. Integral a => a -> a -> a
`div` Int
2

    finish :: Integer -> Int -> ST s Integer
    finish :: Integer -> Int -> ST s Integer
finish !Integer
acc !Int
i | Int
i forall a. Ord a => a -> a -> Bool
< Int
len = do
        Integer
a <- forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray s Integer
arr Int
i
        Integer -> Int -> ST s Integer
finish (Integer
acc forall a. Num a => a -> a -> a
* Integer
base forall a. Num a => a -> a -> a
+ Integer
a) (Int
i forall a. Num a => a -> a -> a
+ Int
1)
    finish !Integer
acc !Int
_ =
        forall (m :: * -> *) a. Monad m => a -> m a
return Integer
acc

-------------------------------------------------------------------------------
-- List variant
-------------------------------------------------------------------------------

{-

-- | A sub-quadratic algorithm implementation using lists.
--
-- Sometimes this is faster, but I fail to quantify when exactly.
--
algorithmL
    :: Integer    -- ^ base
    -> Int        -- ^ length of digits
    -> [Integer]  -- ^ digits
    -> Integer
algorithmL = go
  where
    go :: Integer -> Int -> [Integer] -> Integer
    go _ _ []  = 0
    go _ _ [d] = d
    go b l ds
        | l > 40 = b' `seq` go b' l' (combine b ds')
        | otherwise = finishAlgorithmL b ds
      where
        -- ensure that we have an even number of digits
        -- before we call combine:
        ds' = if even l then ds else 0 : ds
        b' = b * b
        l' = (l + 1) `quot` 2

    combine b (d1 : d2 : ds) = d `seq` (d : combine b ds)
      where
        d = d1 * b + d2
    combine _ []  = []
    combine _ [_] = errorWithoutStackTrace "this should not happen"

-- | The following algorithm is only linear for types whose Num operations
-- are in constant time.
--
-- We export this (mostly) for testing purposes.
--
finishAlgorithmL :: Integer -> [Integer] -> Integer
finishAlgorithmL base = go 0
  where
    go r [] = r
    go r (d : ds) = r' `seq` go r' ds
      where
        r' = r * base + fromIntegral d
-}

-------------------------------------------------------------------------------
-- Misc
-------------------------------------------------------------------------------

integer0 :: Integer
integer0 :: Integer
integer0 = Integer
0