{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DeriveFunctor       #-}
{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Codec.CBOR.Write
  ( toBuilder          
  , toLazyByteString   
  , toStrictByteString 
  ) where
#include "cbor.h"
import           Data.Bits
import           Data.Int
#if ! MIN_VERSION_base(4,11,0)
import           Data.Monoid
#endif
import           Data.Word
import           Foreign.Ptr
import qualified Data.ByteString                       as S
import qualified Data.ByteString.Builder               as B
import qualified Data.ByteString.Builder.Internal      as BI
import           Data.ByteString.Builder.Prim          (condB, (>$<), (>*<))
import qualified Data.ByteString.Builder.Prim          as P
import qualified Data.ByteString.Builder.Prim.Internal as PI
import qualified Data.ByteString.Lazy                  as L
import qualified Data.Text                             as T
import qualified Data.Text.Encoding                    as T
#if defined(OPTIMIZE_GMP)
import           Control.Exception.Base                (assert)
import           GHC.Exts
import qualified GHC.Integer.GMP.Internals             as Gmp
#if __GLASGOW_HASKELL__ < 710
import           GHC.Word
#endif
#endif
import qualified Codec.CBOR.ByteArray.Sliced           as BAS
import           Codec.CBOR.Encoding
import           Codec.CBOR.Magic
toLazyByteString :: Encoding     
                 -> L.ByteString 
toLazyByteString = B.toLazyByteString . toBuilder
toStrictByteString :: Encoding     
                   -> S.ByteString 
toStrictByteString = L.toStrict . B.toLazyByteString . toBuilder
toBuilder :: Encoding  
          -> B.Builder 
toBuilder =
    \(Encoding vs0) -> BI.builder (buildStep (vs0 TkEnd))
buildStep :: Tokens
          -> (BI.BufferRange -> IO (BI.BuildSignal a))
          -> BI.BufferRange
          -> IO (BI.BuildSignal a)
buildStep vs1 k (BI.BufferRange op0 ope0) =
    go vs1 op0
  where
    go vs !op
      | op `plusPtr` bound <= ope0 = case vs of
          TkWord     x vs' -> PI.runB wordMP     x op >>= go vs'
          TkWord64   x vs' -> PI.runB word64MP   x op >>= go vs'
          TkInt      x vs' -> PI.runB intMP      x op >>= go vs'
          TkInt64    x vs' -> PI.runB int64MP    x op >>= go vs'
          TkBytes        x vs' -> BI.runBuilderWith
                                    (bytesMP  x) (buildStep vs' k)
                                    (BI.BufferRange op ope0)
          TkByteArray    x vs' -> BI.runBuilderWith
                                    (byteArrayMP x) (buildStep vs' k)
                                    (BI.BufferRange op ope0)
          TkUtf8ByteArray x vs' -> BI.runBuilderWith
                                     (utf8ByteArrayMP x) (buildStep vs' k)
                                     (BI.BufferRange op ope0)
          TkString        x vs' -> BI.runBuilderWith
                                     (stringMP x) (buildStep vs' k)
                                     (BI.BufferRange op ope0)
          TkBytesBegin vs' -> PI.runB bytesBeginMP  () op >>= go vs'
          TkStringBegin vs'-> PI.runB stringBeginMP () op >>= go vs'
          TkListLen  x vs' -> PI.runB arrayLenMP     x op >>= go vs'
          TkListBegin  vs' -> PI.runB arrayBeginMP  () op >>= go vs'
          TkMapLen   x vs' -> PI.runB mapLenMP       x op >>= go vs'
          TkMapBegin   vs' -> PI.runB mapBeginMP    () op >>= go vs'
          TkTag      x vs' -> PI.runB tagMP          x op >>= go vs'
          TkTag64    x vs' -> PI.runB tag64MP        x op >>= go vs'
#if defined(OPTIMIZE_GMP)
          
          
          
          TkInteger (Gmp.S# i) vs' -> PI.runB intMP (I# i) op >>= go vs'
          
          TkInteger integer@(Gmp.Jp# bigNat) vs'
            | integer <= fromIntegral (maxBound :: Word64) ->
                PI.runB word64MP (fromIntegral integer) op >>= go vs'
            | otherwise ->
               let buffer = BI.BufferRange op ope0
               in BI.runBuilderWith
                    (bigNatMP bigNat) (buildStep vs' k) buffer
          
          TkInteger integer@(Gmp.Jn# bigNat) vs'
            | integer >= -1 - fromIntegral (maxBound :: Word64) ->
                PI.runB negInt64MP (fromIntegral (-1 - integer)) op >>= go vs'
            | otherwise ->
                let buffer = BI.BufferRange op ope0
                in BI.runBuilderWith
                     (negBigNatMP bigNat) (buildStep vs' k) buffer
#else
          TkInteger  x vs'
            | x >= 0
            , x <= fromIntegral (maxBound :: Word64)
                            -> PI.runB word64MP (fromIntegral x) op >>= go vs'
            | x <  0
            , x >= -1 - fromIntegral (maxBound :: Word64)
                            -> PI.runB negInt64MP (fromIntegral (-1 - x)) op >>= go vs'
            | otherwise     -> BI.runBuilderWith
                                 (integerMP x) (buildStep vs' k)
                                 (BI.BufferRange op ope0)
#endif
          TkBool False vs' -> PI.runB falseMP   () op >>= go vs'
          TkBool True  vs' -> PI.runB trueMP    () op >>= go vs'
          TkNull       vs' -> PI.runB nullMP    () op >>= go vs'
          TkUndef      vs' -> PI.runB undefMP   () op >>= go vs'
          TkSimple   w vs' -> PI.runB simpleMP   w op >>= go vs'
          TkFloat16  f vs' -> PI.runB halfMP     f op >>= go vs'
          TkFloat32  f vs' -> PI.runB floatMP    f op >>= go vs'
          TkFloat64  f vs' -> PI.runB doubleMP   f op >>= go vs'
          TkBreak      vs' -> PI.runB breakMP   () op >>= go vs'
          TkEncoded  x vs' -> BI.runBuilderWith
                                (B.byteString x) (buildStep vs' k)
                                (BI.BufferRange op ope0)
          TkEnd            -> k (BI.BufferRange op ope0)
      | otherwise = return $ BI.bufferFull bound op (buildStep vs k)
    
    bound :: Int
    bound = 9
header :: P.BoundedPrim Word8
header = P.liftFixedToBounded P.word8
constHeader :: Word8 -> P.BoundedPrim ()
constHeader h = P.liftFixedToBounded (const h >$< P.word8)
withHeader :: P.FixedPrim a -> P.BoundedPrim (Word8, a)
withHeader p = P.liftFixedToBounded (P.word8 >*< p)
withConstHeader :: Word8 -> P.FixedPrim a -> P.BoundedPrim a
withConstHeader h p = P.liftFixedToBounded ((,) h >$< (P.word8 >*< p))
{-# INLINE wordMP #-}
wordMP :: P.BoundedPrim Word
wordMP =
    condB (<= 0x17)       (fromIntegral >$< header) $
    condB (<= 0xff)       (fromIntegral >$< withConstHeader 24 P.word8) $
    condB (<= 0xffff)     (fromIntegral >$< withConstHeader 25 P.word16BE) $
#if defined(ARCH_64bit)
    condB (<= 0xffffffff) (fromIntegral >$< withConstHeader 26 P.word32BE) $
                          (fromIntegral >$< withConstHeader 27 P.word64BE)
#else
                          (fromIntegral >$< withConstHeader 26 P.word32BE)
#endif
{-# INLINE word64MP #-}
word64MP :: P.BoundedPrim Word64
word64MP =
    condB (<= 0x17)       (fromIntegral >$< header) $
    condB (<= 0xff)       (fromIntegral >$< withConstHeader 24 P.word8) $
    condB (<= 0xffff)     (fromIntegral >$< withConstHeader 25 P.word16BE) $
    condB (<= 0xffffffff) (fromIntegral >$< withConstHeader 26 P.word32BE) $
                          (fromIntegral >$< withConstHeader 27 P.word64BE)
negInt64MP :: P.BoundedPrim Word64
negInt64MP =
    condB (<= 0x17)       (fromIntegral . (0x20 +) >$< header) $
    condB (<= 0xff)       (fromIntegral >$< withConstHeader 0x38 P.word8) $
    condB (<= 0xffff)     (fromIntegral >$< withConstHeader 0x39 P.word16BE) $
    condB (<= 0xffffffff) (fromIntegral >$< withConstHeader 0x3a P.word32BE) $
                          (fromIntegral >$< withConstHeader 0x3b P.word64BE)
{-# INLINE intMP #-}
intMP :: P.BoundedPrim Int
intMP =
    prep >$< (
      condB ((<= 0x17)       . snd) (encIntSmall >$< header) $
      condB ((<= 0xff)       . snd) (encInt8  >$< withHeader P.word8) $
      condB ((<= 0xffff)     . snd) (encInt16 >$< withHeader P.word16BE) $
#if defined(ARCH_64bit)
      condB ((<= 0xffffffff) . snd) (encInt32 >$< withHeader P.word32BE)
                                    (encInt64 >$< withHeader P.word64BE)
#else
                                    (encInt32 >$< withHeader P.word32BE)
#endif
    )
  where
    prep :: Int -> (Word8, Word)
    prep n = (mt, ui)
      where
        sign :: Word     
        sign = fromIntegral (n `unsafeShiftR` intBits)
#if MIN_VERSION_base(4,7,0)
        intBits = finiteBitSize (undefined :: Int) - 1
#else
        intBits = bitSize (undefined :: Int) - 1
#endif
        mt   :: Word8    
        mt   = fromIntegral (sign .&. 0x20)
        ui   :: Word     
        ui   = fromIntegral n `xor` sign
    encIntSmall :: (Word8, Word) -> Word8
    encIntSmall (mt, ui) =  mt + fromIntegral ui
    encInt8     (mt, ui) = (mt + 24, fromIntegral ui)
    encInt16    (mt, ui) = (mt + 25, fromIntegral ui)
    encInt32    (mt, ui) = (mt + 26, fromIntegral ui)
#if defined(ARCH_64bit)
    encInt64    (mt, ui) = (mt + 27, fromIntegral ui)
#endif
{-# INLINE int64MP #-}
int64MP :: P.BoundedPrim Int64
int64MP =
    prep >$< (
      condB ((<= 0x17)       . snd) (encIntSmall >$< header) $
      condB ((<= 0xff)       . snd) (encInt8  >$< withHeader P.word8) $
      condB ((<= 0xffff)     . snd) (encInt16 >$< withHeader P.word16BE) $
      condB ((<= 0xffffffff) . snd) (encInt32 >$< withHeader P.word32BE)
                                    (encInt64 >$< withHeader P.word64BE)
    )
  where
    prep :: Int64 -> (Word8, Word64)
    prep n = (mt, ui)
      where
        sign :: Word64   
        sign = fromIntegral (n `unsafeShiftR` intBits)
#if MIN_VERSION_base(4,7,0)
        intBits = finiteBitSize (undefined :: Int64) - 1
#else
        intBits = bitSize (undefined :: Int64) - 1
#endif
        mt   :: Word8    
        mt   = fromIntegral (sign .&. 0x20)
        ui   :: Word64   
        ui   = fromIntegral n `xor` sign
    encIntSmall (mt, ui) =  mt + fromIntegral ui
    encInt8     (mt, ui) = (mt + 24, fromIntegral ui)
    encInt16    (mt, ui) = (mt + 25, fromIntegral ui)
    encInt32    (mt, ui) = (mt + 26, fromIntegral ui)
    encInt64    (mt, ui) = (mt + 27, fromIntegral ui)
bytesMP :: S.ByteString -> B.Builder
bytesMP bs =
    P.primBounded bytesLenMP (fromIntegral $ S.length bs) <> B.byteString bs
bytesLenMP :: P.BoundedPrim Word
bytesLenMP =
    condB (<= 0x17)       (fromIntegral . (0x40 +) >$< header) $
    condB (<= 0xff)       (fromIntegral >$< withConstHeader 0x58 P.word8) $
    condB (<= 0xffff)     (fromIntegral >$< withConstHeader 0x59 P.word16BE) $
    condB (<= 0xffffffff) (fromIntegral >$< withConstHeader 0x5a P.word32BE) $
                          (fromIntegral >$< withConstHeader 0x5b P.word64BE)
byteArrayMP :: BAS.SlicedByteArray -> B.Builder
byteArrayMP ba =
    P.primBounded bytesLenMP n <> BAS.toBuilder ba
  where n = fromIntegral $ BAS.sizeofSlicedByteArray ba
bytesBeginMP :: P.BoundedPrim ()
bytesBeginMP = constHeader 0x5f
stringMP :: T.Text -> B.Builder
stringMP t =
    P.primBounded stringLenMP (fromIntegral $ S.length bs) <> B.byteString bs
  where
    bs  = T.encodeUtf8 t
stringLenMP :: P.BoundedPrim Word
stringLenMP =
    condB (<= 0x17)       (fromIntegral . (0x60 +) >$< header) $
    condB (<= 0xff)       (fromIntegral >$< withConstHeader 0x78 P.word8) $
    condB (<= 0xffff)     (fromIntegral >$< withConstHeader 0x79 P.word16BE) $
    condB (<= 0xffffffff) (fromIntegral >$< withConstHeader 0x7a P.word32BE) $
                          (fromIntegral >$< withConstHeader 0x7b P.word64BE)
stringBeginMP :: P.BoundedPrim ()
stringBeginMP = constHeader 0x7f
utf8ByteArrayMP :: BAS.SlicedByteArray -> B.Builder
utf8ByteArrayMP t =
    P.primBounded stringLenMP n <> BAS.toBuilder t
  where
    n = fromIntegral $ BAS.sizeofSlicedByteArray t
arrayLenMP :: P.BoundedPrim Word
arrayLenMP =
    condB (<= 0x17)       (fromIntegral . (0x80 +) >$< header) $
    condB (<= 0xff)       (fromIntegral >$< withConstHeader 0x98 P.word8) $
    condB (<= 0xffff)     (fromIntegral >$< withConstHeader 0x99 P.word16BE) $
    condB (<= 0xffffffff) (fromIntegral >$< withConstHeader 0x9a P.word32BE) $
                          (fromIntegral >$< withConstHeader 0x9b P.word64BE)
arrayBeginMP :: P.BoundedPrim ()
arrayBeginMP = constHeader 0x9f
mapLenMP :: P.BoundedPrim Word
mapLenMP =
    condB (<= 0x17)       (fromIntegral . (0xa0 +) >$< header) $
    condB (<= 0xff)       (fromIntegral >$< withConstHeader 0xb8 P.word8) $
    condB (<= 0xffff)     (fromIntegral >$< withConstHeader 0xb9 P.word16BE) $
    condB (<= 0xffffffff) (fromIntegral >$< withConstHeader 0xba P.word32BE) $
                          (fromIntegral >$< withConstHeader 0xbb P.word64BE)
mapBeginMP :: P.BoundedPrim ()
mapBeginMP = constHeader 0xbf
tagMP :: P.BoundedPrim Word
tagMP =
    condB (<= 0x17)       (fromIntegral . (0xc0 +) >$< header) $
    condB (<= 0xff)       (fromIntegral >$< withConstHeader 0xd8 P.word8) $
    condB (<= 0xffff)     (fromIntegral >$< withConstHeader 0xd9 P.word16BE) $
#if defined(ARCH_64bit)
    condB (<= 0xffffffff) (fromIntegral >$< withConstHeader 0xda P.word32BE) $
                          (fromIntegral >$< withConstHeader 0xdb P.word64BE)
#else
                          (fromIntegral >$< withConstHeader 0xda P.word32BE)
#endif
tag64MP :: P.BoundedPrim Word64
tag64MP =
    condB (<= 0x17)       (fromIntegral . (0xc0 +) >$< header) $
    condB (<= 0xff)       (fromIntegral >$< withConstHeader 0xd8 P.word8) $
    condB (<= 0xffff)     (fromIntegral >$< withConstHeader 0xd9 P.word16BE) $
    condB (<= 0xffffffff) (fromIntegral >$< withConstHeader 0xda P.word32BE) $
                          (fromIntegral >$< withConstHeader 0xdb P.word64BE)
simpleMP :: P.BoundedPrim Word8
simpleMP =
    condB (<= 0x17) ((0xe0 +) >$< header) $
                    (withConstHeader 0xf8 P.word8)
falseMP :: P.BoundedPrim ()
falseMP = constHeader 0xf4
trueMP :: P.BoundedPrim ()
trueMP = constHeader 0xf5
nullMP :: P.BoundedPrim ()
nullMP = constHeader 0xf6
undefMP :: P.BoundedPrim ()
undefMP = constHeader 0xf7
canonicalNaN :: PI.BoundedPrim a
canonicalNaN = P.liftFixedToBounded $ const (0xf9, (0x7e, 0x00))
                                   >$< P.word8 >*< P.word8 >*< P.word8
halfMP :: P.BoundedPrim Float
halfMP = condB isNaN canonicalNaN
                     (floatToWord16 >$< withConstHeader 0xf9 P.word16BE)
floatMP :: P.BoundedPrim Float
floatMP = condB isNaN canonicalNaN
                      (withConstHeader 0xfa P.floatBE)
doubleMP :: P.BoundedPrim Double
doubleMP = condB isNaN canonicalNaN
                       (withConstHeader 0xfb P.doubleBE)
breakMP :: P.BoundedPrim ()
breakMP = constHeader 0xff
#if defined(OPTIMIZE_GMP)
bigNatMP :: Gmp.BigNat -> B.Builder
bigNatMP n = P.primBounded header 0xc2 <> bigNatToBuilder n
negBigNatMP :: Gmp.BigNat -> B.Builder
negBigNatMP n =
  
  
  
     P.primBounded header 0xc3
  <> bigNatToBuilder (Gmp.minusBigNatWord n (int2Word# 1#))
bigNatToBuilder :: Gmp.BigNat -> B.Builder
bigNatToBuilder = bigNatBuilder
  where
    bigNatBuilder :: Gmp.BigNat -> B.Builder
    bigNatBuilder bigNat =
        let sizeW# = Gmp.sizeInBaseBigNat bigNat 256#
            bounded = PI.boudedPrim (I# (word2Int# sizeW#)) (dumpBigNat sizeW#)
        in P.primBounded bytesLenMP (W# sizeW#) <> P.primBounded bounded bigNat
    dumpBigNat :: Word# -> Gmp.BigNat -> Ptr a -> IO (Ptr a)
    dumpBigNat sizeW# bigNat ptr@(Ptr addr#) = do
        
        
        (W# written#) <- Gmp.exportBigNatToAddr bigNat addr# 1#
        let !newPtr = ptr `plusPtr` (I# (word2Int# written#))
            sanity = isTrue# (sizeW# `eqWord#` written#)
        return $ assert sanity newPtr
#else
integerMP :: Integer -> B.Builder
integerMP n
  | n >= 0    = P.primBounded header 0xc2 <> integerToBuilder n
  | otherwise = P.primBounded header 0xc3 <> integerToBuilder (-1 - n)
integerToBuilder :: Integer -> B.Builder
integerToBuilder n = bytesMP (integerToBytes n)
integerToBytes :: Integer -> S.ByteString
integerToBytes n0
  | n0 == 0   = S.pack [0]
  | otherwise = S.pack (reverse (go n0))
  where
    go n | n == 0    = []
         | otherwise = narrow n : go (n `shiftR` 8)
    narrow :: Integer -> Word8
    narrow = fromIntegral
#endif