{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP          #-}

-- |Primitives to calculate the encoding size of a value
module Flat.Encoder.Size where

import           Data.Bits
import qualified Data.ByteString                as B
import qualified Data.ByteString.Lazy           as L
import qualified Data.ByteString.Short.Internal as SBS
import           Data.Char
import           Flat.Encoder.Prim         (w7l)
import           Flat.Encoder.Types
import           Flat.Types
import qualified Data.Text                      as T
#ifndef ghcjs_HOST_OS
import qualified Data.Text.Internal             as TI
#endif
import           Data.ZigZag
#include "MachDeps.h"
-- A filler can take anything from 1 to 8 bits
sFillerMax :: NumBits
sFillerMax :: NumBits
sFillerMax = NumBits
8

sBool :: NumBits
sBool :: NumBits
sBool = NumBits
1

sWord8 :: NumBits
sWord8 :: NumBits
sWord8 = NumBits
8

sInt8 :: NumBits
sInt8 :: NumBits
sInt8 = NumBits
8

sFloat :: NumBits
sFloat :: NumBits
sFloat = NumBits
32

sDouble :: NumBits
sDouble :: NumBits
sDouble = NumBits
64

{-# INLINE sChar #-}
sChar :: Char -> NumBits
sChar :: Char -> NumBits
sChar = Word32 -> NumBits
sWord32 (Word32 -> NumBits) -> (Char -> Word32) -> Char -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NumBits -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (NumBits -> Word32) -> (Char -> NumBits) -> Char -> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> NumBits
ord

sCharMax :: NumBits
sCharMax :: NumBits
sCharMax = NumBits
24

{-# INLINE sWord #-}
sWord :: Word -> NumBits
{-# INLINE sInt #-}
sInt :: Int -> NumBits
#if WORD_SIZE_IN_BITS == 64
sWord :: Word -> NumBits
sWord = Word64 -> NumBits
sWord64 (Word64 -> NumBits) -> (Word -> Word64) -> Word -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral

sInt :: NumBits -> NumBits
sInt = Int64 -> NumBits
sInt64 (Int64 -> NumBits) -> (NumBits -> Int64) -> NumBits -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NumBits -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral
#elif WORD_SIZE_IN_BITS == 32
sWord = sWord32 . fromIntegral

sInt = sInt32 . fromIntegral
#else
#error expected WORD_SIZE_IN_BITS to be 32 or 64
#endif
-- TODO: optimize ints sizes
{-# INLINE sInt16 #-}
sInt16 :: Int16 -> NumBits
sInt16 :: Int16 -> NumBits
sInt16 = Word16 -> NumBits
sWord16 (Word16 -> NumBits) -> (Int16 -> Word16) -> Int16 -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int16 -> Word16
forall signed unsigned.
ZigZag signed unsigned =>
signed -> unsigned
zigZag

{-# INLINE sInt32 #-}
sInt32 :: Int32 -> NumBits
sInt32 :: Int32 -> NumBits
sInt32 = Word32 -> NumBits
sWord32 (Word32 -> NumBits) -> (Int32 -> Word32) -> Int32 -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int32 -> Word32
forall signed unsigned.
ZigZag signed unsigned =>
signed -> unsigned
zigZag

{-# INLINE sInt64 #-}
sInt64 :: Int64 -> NumBits
sInt64 :: Int64 -> NumBits
sInt64 = Word64 -> NumBits
sWord64 (Word64 -> NumBits) -> (Int64 -> Word64) -> Int64 -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Word64
forall signed unsigned.
ZigZag signed unsigned =>
signed -> unsigned
zigZag

{-# INLINE sWord16 #-}
sWord16 :: Word16 -> NumBits
sWord16 :: Word16 -> NumBits
sWord16 Word16
w
  | Word16
w Word16 -> Word16 -> Bool
forall a. Ord a => a -> a -> Bool
< Word16
128 = NumBits
8
  | Word16
w Word16 -> Word16 -> Bool
forall a. Ord a => a -> a -> Bool
< Word16
16384 = NumBits
16
  | Bool
otherwise = NumBits
24

{-# INLINE sWord32 #-}
sWord32 :: Word32 -> NumBits
sWord32 :: Word32 -> NumBits
sWord32 Word32
w
  | Word32
w Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
< Word32
128 = NumBits
8
  | Word32
w Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
< Word32
16384 = NumBits
16
  | Word32
w Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
< Word32
2097152 = NumBits
24
  | Word32
w Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
< Word32
268435456 = NumBits
32
  | Bool
otherwise = NumBits
40

{-# INLINE sWord64 #-}
sWord64 :: Word64 -> NumBits
sWord64 :: Word64 -> NumBits
sWord64 Word64
w
  | Word64
w Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
128 = NumBits
8
  | Word64
w Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
16384 = NumBits
16
  | Word64
w Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
2097152 = NumBits
24
  | Word64
w Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
268435456 = NumBits
32
  | Word64
w Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
34359738368 = NumBits
40
  | Word64
w Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
4398046511104 = NumBits
48
  | Word64
w Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
562949953421312 = NumBits
56
  | Word64
w Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
72057594037927936 = NumBits
64
  | Word64
w Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
9223372036854775808 = NumBits
72
  | Bool
otherwise = NumBits
80

{-# INLINE sInteger #-}
sInteger :: Integer -> NumBits
sInteger :: Integer -> NumBits
sInteger = Natural -> NumBits
forall t. (Bits t, Integral t) => t -> NumBits
sIntegral (Natural -> NumBits) -> (Integer -> Natural) -> Integer -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Natural
forall signed unsigned.
ZigZag signed unsigned =>
signed -> unsigned
zigZag

{-# INLINE sNatural #-}
sNatural :: Natural -> NumBits
sNatural :: Natural -> NumBits
sNatural = Integer -> NumBits
forall t. (Bits t, Integral t) => t -> NumBits
sIntegral (Integer -> NumBits) -> (Natural -> Integer) -> Natural -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Natural -> Integer
forall a. Integral a => a -> Integer
toInteger

-- BAD: duplication of work with encoding
{-# INLINE sIntegral #-}
sIntegral :: (Bits t, Integral t) => t -> Int
sIntegral :: t -> NumBits
sIntegral t
t =
  let vs :: [Word8]
vs = t -> [Word8]
forall t. (Bits t, Integral t) => t -> [Word8]
w7l t
t
   in [Word8] -> NumBits
forall (t :: * -> *) a. Foldable t => t a -> NumBits
length [Word8]
vs NumBits -> NumBits -> NumBits
forall a. Num a => a -> a -> a
* NumBits
8

--sUTF8 :: T.Text -> NumBits
--sUTF8 t = fold
-- Wildly pessimistic but fast
{-# INLINE sUTF8Max #-}
sUTF8Max :: Text -> NumBits
sUTF8Max :: Text -> NumBits
sUTF8Max = NumBits -> NumBits
blobBits (NumBits -> NumBits) -> (Text -> NumBits) -> Text -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NumBits
4 NumBits -> NumBits -> NumBits
forall a. Num a => a -> a -> a
*) (NumBits -> NumBits) -> (Text -> NumBits) -> Text -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> NumBits
T.length
#ifndef ghcjs_HOST_OS
{-# INLINE sUTF16 #-}
sUTF16 :: T.Text -> NumBits
sUTF16 :: Text -> NumBits
sUTF16 = NumBits -> NumBits
blobBits (NumBits -> NumBits) -> (Text -> NumBits) -> Text -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> NumBits
textBytes
#endif
{-# INLINE sBytes #-}
sBytes :: B.ByteString -> NumBits
sBytes :: ByteString -> NumBits
sBytes = NumBits -> NumBits
blobBits (NumBits -> NumBits)
-> (ByteString -> NumBits) -> ByteString -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> NumBits
B.length

{-# INLINE sLazyBytes #-}
sLazyBytes :: L.ByteString -> NumBits
sLazyBytes :: ByteString -> NumBits
sLazyBytes ByteString
bs = NumBits
16 NumBits -> NumBits -> NumBits
forall a. Num a => a -> a -> a
+ (ByteString -> NumBits -> NumBits)
-> NumBits -> ByteString -> NumBits
forall a. (ByteString -> a -> a) -> a -> ByteString -> a
L.foldrChunks (\ByteString
b NumBits
l -> ByteString -> NumBits
blkBitsBS ByteString
b NumBits -> NumBits -> NumBits
forall a. Num a => a -> a -> a
+ NumBits
l) NumBits
0 ByteString
bs

{-# INLINE sShortBytes #-}
sShortBytes :: SBS.ShortByteString -> NumBits
sShortBytes :: ShortByteString -> NumBits
sShortBytes = NumBits -> NumBits
blobBits (NumBits -> NumBits)
-> (ShortByteString -> NumBits) -> ShortByteString -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShortByteString -> NumBits
SBS.length

#ifndef ghcjs_HOST_OS
-- We are not interested in the number of unicode chars (returned by T.length, an O(n) operation)
-- just the number of bytes
-- > T.length (T.pack "\x1F600")
-- 1
-- > textBytes (T.pack "\x1F600")
-- 4
{-# INLINE textBytes #-}
textBytes :: T.Text -> Int
textBytes :: Text -> NumBits
textBytes !(TI.Text Array
_ NumBits
_ NumBits
w16Len) = NumBits
w16Len NumBits -> NumBits -> NumBits
forall a. Num a => a -> a -> a
* NumBits
2
#endif

{-# INLINE bitsToBytes #-}
bitsToBytes :: Int -> Int
bitsToBytes :: NumBits -> NumBits
bitsToBytes = NumBits -> NumBits -> NumBits
forall t. Integral t => t -> t -> t
numBlks NumBits
8

{-# INLINE numBlks #-}
numBlks :: Integral t => t -> t -> t
numBlks :: t -> t -> t
numBlks t
blkSize t
bits =
  let (t
d, t
m) = t
bits t -> t -> (t, t)
forall a. Integral a => a -> a -> (a, a)
`divMod` t
blkSize
   in t
d t -> t -> t
forall a. Num a => a -> a -> a
+
      (if t
m t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
0
         then t
0
         else t
1)

{-# INLINE arrayBits #-}
arrayBits :: Int -> NumBits
arrayBits :: NumBits -> NumBits
arrayBits = (NumBits
8 NumBits -> NumBits -> NumBits
forall a. Num a => a -> a -> a
*) (NumBits -> NumBits) -> (NumBits -> NumBits) -> NumBits -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NumBits -> NumBits
arrayChunks

{-# INLINE arrayChunks #-}
arrayChunks :: Int -> NumBits
arrayChunks :: NumBits -> NumBits
arrayChunks = (NumBits
1 NumBits -> NumBits -> NumBits
forall a. Num a => a -> a -> a
+) (NumBits -> NumBits) -> (NumBits -> NumBits) -> NumBits -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NumBits -> NumBits -> NumBits
forall t. Integral t => t -> t -> t
numBlks NumBits
255

{-# INLINE blobBits #-}
blobBits :: Int -> NumBits
blobBits :: NumBits -> NumBits
blobBits NumBits
numBytes =
  NumBits
16 -- initial filler + final 0
   NumBits -> NumBits -> NumBits
forall a. Num a => a -> a -> a
+
  NumBits -> NumBits
blksBits NumBits
numBytes

{-# INLINE blkBitsBS #-}
blkBitsBS :: B.ByteString -> NumBits
blkBitsBS :: ByteString -> NumBits
blkBitsBS = NumBits -> NumBits
blksBits (NumBits -> NumBits)
-> (ByteString -> NumBits) -> ByteString -> NumBits
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> NumBits
B.length

{-# INLINE blksBits #-}
blksBits :: Int -> NumBits
blksBits :: NumBits -> NumBits
blksBits NumBits
numBytes = NumBits
8 NumBits -> NumBits -> NumBits
forall a. Num a => a -> a -> a
* (NumBits
numBytes NumBits -> NumBits -> NumBits
forall a. Num a => a -> a -> a
+ NumBits -> NumBits -> NumBits
forall t. Integral t => t -> t -> t
numBlks NumBits
255 NumBits
numBytes)