{-# LANGUAGE BangPatterns #-}
module Data.Bencode.Util
  ( arrayFromRevListN
  , readKnownNaturalAsInt
  , readKnownNaturalAsInt64
  , readKnownNaturalAsWord
  , readKnownNaturalAsWord64
  ) where

import Data.Bits
import Data.Int
import Data.Word
import qualified Data.ByteString as B
import qualified Data.ByteString.Unsafe as B
import qualified Data.Primitive.Array as A

-- | Create an array from a list in reverse order. The list length must be n.
arrayFromRevListN :: Int -> [a] -> A.Array a
arrayFromRevListN :: forall a. Int -> [a] -> Array a
arrayFromRevListN Int
n [a]
xs = forall a.
Int -> a -> (forall s. MutableArray s a -> ST s ()) -> Array a
A.createArray Int
n forall a. a
errorElement forall a b. (a -> b) -> a -> b
$ \MutableArray s a
a ->
  let f :: a -> (Int -> ST s ()) -> Int -> ST s ()
f a
x Int -> ST s ()
k = \Int
i ->
        if Int
i forall a. Eq a => a -> a -> Bool
== -Int
1
        then forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        else forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
A.writeArray MutableArray s a
a Int
i a
x forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> ST s ()
k (Int
iforall a. Num a => a -> a -> a
-Int
1)
  in forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> (Int -> ST s ()) -> Int -> ST s ()
f (\ !Int
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) [a]
xs (Int
nforall a. Num a => a -> a -> a
-Int
1)
{-# INLINE arrayFromRevListN #-}

errorElement :: a
errorElement :: forall a. a
errorElement = forall a. HasCallStack => [Char] -> a
error [Char]
"errorElement"

-- | The input string must be an unsigned decimal integer with no extraneous
-- leading zeros. Returns Nothing if the value is outside the bounds of an
-- @Int@.
readKnownNaturalAsInt :: Bool -> B.ByteString -> Maybe Int
readKnownNaturalAsInt :: Bool -> ByteString -> Maybe Int
readKnownNaturalAsInt = forall a.
(Bounded a, Integral a) =>
Int -> Bool -> ByteString -> Maybe a
readInt Int
maxIntLen
  where
    maxIntLen :: Int
maxIntLen = case forall b. FiniteBits b => b -> Int
finiteBitSize (Int
0 :: Int) of
      Int
32 -> Int
10
      Int
64 -> Int
19
      Int
_  -> forall a. HasCallStack => [Char] -> a
error [Char]
"unsupported word size"
{-# INLINE readKnownNaturalAsInt #-}

-- | Similar to 'readKnownNaturalAsInt', for 'Int64'.
readKnownNaturalAsInt64 :: Bool -> B.ByteString -> Maybe Int64
readKnownNaturalAsInt64 :: Bool -> ByteString -> Maybe Int64
readKnownNaturalAsInt64 = forall a.
(Bounded a, Integral a) =>
Int -> Bool -> ByteString -> Maybe a
readInt Int
19
{-# INLINE readKnownNaturalAsInt64 #-}

readInt :: (Bounded a, Integral a) => Int -> Bool -> B.ByteString -> Maybe a
readInt :: forall a.
(Bounded a, Integral a) =>
Int -> Bool -> ByteString -> Maybe a
readInt Int
maxLen Bool
neg ByteString
s = if Bool
neg then forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Num a => a -> a
negate Maybe a
n else Maybe a
n
  where
    -- last digit of maxBound = 7, minBound = 8
    n :: Maybe a
n = forall a.
(Bounded a, Integral a) =>
Int -> a -> a -> ByteString -> Maybe a
readWord Int
maxLen (forall a. Bounded a => a
maxBound forall a. Integral a => a -> a -> a
`div` a
10) (a
7 forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Enum a => a -> Int
fromEnum Bool
neg)) ByteString
s
{-# INLINE readInt #-}

-- | The input string must be an unsigned decimal integer with no extraneous
-- leading zeros. Returns Nothing if the value is outside the bounds of a
-- @Word@.
readKnownNaturalAsWord :: B.ByteString -> Maybe Word
readKnownNaturalAsWord :: ByteString -> Maybe Word
readKnownNaturalAsWord = forall a.
(Bounded a, Integral a) =>
Int -> a -> a -> ByteString -> Maybe a
readWord Int
maxWordLen (forall a. Bounded a => a
maxBound forall a. Integral a => a -> a -> a
`div` Word
10) Word
5
  where
    -- last digit of maxBound = 5
    maxWordLen :: Int
maxWordLen = case forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) of
      Int
32 -> Int
10
      Int
64 -> Int
20
      Int
_  -> forall a. HasCallStack => [Char] -> a
error [Char]
"unsupported word size"
{-# INLINE readKnownNaturalAsWord #-}

-- | Similar to 'readKnownNaturalAsWord', for 'Word64'.
readKnownNaturalAsWord64 :: B.ByteString -> Maybe Word64
readKnownNaturalAsWord64 :: ByteString -> Maybe Word64
readKnownNaturalAsWord64 = forall a.
(Bounded a, Integral a) =>
Int -> a -> a -> ByteString -> Maybe a
readWord Int
20 (forall a. Bounded a => a
maxBound forall a. Integral a => a -> a -> a
`div` Word64
10) Word64
5
{-# INLINE readKnownNaturalAsWord64 #-}

-- maxLen must be > 0!
readWord :: (Bounded a, Integral a) => Int -> a -> a -> B.ByteString -> Maybe a
readWord :: forall a.
(Bounded a, Integral a) =>
Int -> a -> a -> ByteString -> Maybe a
readWord Int
maxLen a
maxValueDiv10 a
maxValueMod10 ByteString
s =
  case forall a. Ord a => a -> a -> Ordering
compare (ByteString -> Int
B.length ByteString
s) Int
maxLen of
    Ordering
LT -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$! forall a. Integral a => ByteString -> a
readFull' ByteString
s
    Ordering
EQ ->
      let n :: a
n = forall a. Integral a => ByteString -> a
readFull (ByteString -> ByteString
B.unsafeInit ByteString
s)
          d :: a
d = forall a. Integral a => Word8 -> a
digitToI (ByteString -> Word8
B.unsafeLast ByteString
s)
      in if a
n forall a. Ord a => a -> a -> Bool
< a
maxValueDiv10 Bool -> Bool -> Bool
|| a
n forall a. Eq a => a -> a -> Bool
== a
maxValueDiv10 Bool -> Bool -> Bool
&& a
d forall a. Ord a => a -> a -> Bool
<= a
maxValueMod10
         then forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$! a
nforall a. Num a => a -> a -> a
*a
10 forall a. Num a => a -> a -> a
+ a
d
         else forall a. Maybe a
Nothing
    Ordering
GT -> forall a. Maybe a
Nothing
{-# INLINE readWord #-}

readFull :: Integral a => B.ByteString -> a
readFull :: forall a. Integral a => ByteString -> a
readFull = forall a. (a -> Word8 -> a) -> a -> ByteString -> a
B.foldl' (\a
acc Word8
c -> a
acc forall a. Num a => a -> a -> a
* a
10 forall a. Num a => a -> a -> a
+ forall a. Integral a => Word8 -> a
digitToI Word8
c) a
0
{-# INLINE readFull #-}

-- Same as readFull but avoids
-- https://gitlab.haskell.org/ghc/ghc/-/issues/24203
readFull' :: Integral a => B.ByteString -> a
readFull' :: forall a. Integral a => ByteString -> a
readFull' ByteString
s = case ByteString -> Maybe (ByteString, Word8)
B.unsnoc ByteString
s of
  Maybe (ByteString, Word8)
Nothing     -> a
0
  Just (ByteString
s',Word8
c) -> forall a. Integral a => ByteString -> a
readFull ByteString
s' forall a. Num a => a -> a -> a
* a
10 forall a. Num a => a -> a -> a
+ forall a. Integral a => Word8 -> a
digitToI Word8
c
{-# INLINE readFull' #-}

digitToI :: Integral a => Word8 -> a
digitToI :: forall a. Integral a => Word8 -> a
digitToI Word8
c = forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
c forall a. Num a => a -> a -> a
- a
48
{-# INLINE digitToI #-}