content: {-# LANGUAGE CPP, MagicHash #-}
  -- for unboxed shifts
  
  -----------------------------------------------------------------------------
  -- |
  -- Module      : Data.Binary.Get
  -- Copyright   : Lennart Kolmodin
  -- License     : BSD3-style (see LICENSE)
  -- 
  -- Maintainer  : Lennart Kolmodin <kolmodin@dtek.chalmers.se>
  -- Stability   : experimental
  -- Portability : portable to Hugs and GHC.
  --
  -- The Get monad. A monad for efficiently building structures from
  -- encoded lazy ByteStrings
  --
  -----------------------------------------------------------------------------
  
  #if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
  #include "MachDeps.h"
  #endif
  
  module Data.Binary.Get (
  
      -- * The Get type
        Get
      , runGet
      , runGetState
  
      -- * Parsing
      , skip
      , uncheckedSkip
      , lookAhead
      , lookAheadM
      , lookAheadE
      , uncheckedLookAhead
  
      -- * Utility
      , bytesRead
      , getBytes
      , remaining
      , isEmpty
  
      -- * Parsing particular types
      , getWord8
  
      -- ** ByteStrings
      , getByteString
      , getLazyByteString
      , getLazyByteStringNul
      , getRemainingLazyByteString
  
      -- ** Big-endian reads
      , getWord16be
      , getWord32be
      , getWord64be
  
      -- ** Little-endian reads
      , getWord16le
      , getWord32le
      , getWord64le
  
      -- ** Host-endian, unaligned reads
      , getWordhost
      , getWord16host
      , getWord32host
      , getWord64host
  
    ) where
  
  import Control.Monad (when,liftM, ap)
  import Control.Monad.Fix
  import Data.Maybe (isNothing)
  
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Lazy as L
  
  #ifdef BYTESTRING_IN_BASE
  import qualified Data.ByteString.Base as B
  #else
  import qualified Data.ByteString.Internal as B
  import qualified Data.ByteString.Lazy.Internal as L
  #endif
  
  import Control.Applicative (Applicative(..))
  
  import Foreign
  
  -- used by splitAtST
  #if MIN_VERSION_base(4,6,0)
  import Control.Monad.ST.Unsafe(unsafeInterleaveST)
  #else
  import Control.Monad.ST(unsafeInterleaveST)
  #endif
  import Control.Monad.ST(runST)
  import Data.STRef
  
  #if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
  import GHC.Base
  import GHC.Word
  --import GHC.Int
  #endif
  
  -- | The parse state
  data S = S {-# UNPACK #-} !B.ByteString  -- current chunk
             L.ByteString                  -- the rest of the input
             {-# UNPACK #-} !Int64         -- bytes read
  
  -- | The Get monad is just a State monad carrying around the input ByteString
  newtype Get a = Get { unGet :: S -> (a, S) }
  
  instance Functor Get where
      fmap f m = Get (\s -> case unGet m s of
                              (a, s') -> (f a, s'))
      {-# INLINE fmap #-}
  
  instance Applicative Get where
      pure  = return
      (<*>) = ap
  
  instance Monad Get where
      return a  = Get (\s -> (a, s))
      {-# INLINE return #-}
  
      m >>= k   = Get (\s -> case unGet m s of
                               (a, s') -> unGet (k a) s')
      {-# INLINE (>>=) #-}
  
      fail      = failDesc
  
  instance MonadFix Get where
      mfix f = Get (\s -> let (a,s') = unGet (f a) s 
                          in (a,s'))
  
  ------------------------------------------------------------------------
  
  get :: Get S
  get   = Get (\s -> (s, s))
  
  put :: S -> Get ()
  put s = Get (\_ -> ((), s))
  
  ------------------------------------------------------------------------
  --
  -- dons, GHC 6.10: explicit inlining disabled, was killing performance.
  -- Without it, GHC seems to do just fine. And we get similar
  -- performance with 6.8.2 anyway.
  --
  
  initState :: L.ByteString -> S
  initState xs = mkState xs 0
  {- INLINE initState -}
  
  {-
  initState (B.LPS xs) =
      case xs of
        []      -> S B.empty L.empty 0
        (x:xs') -> S x (B.LPS xs') 0
  -}
  
  #ifndef BYTESTRING_IN_BASE
  mkState :: L.ByteString -> Int64 -> S
  mkState l = case l of
      L.Empty      -> S B.empty L.empty
      L.Chunk x xs -> S x xs
  {- INLINE mkState -}
  
  #else
  mkState :: L.ByteString -> Int64 -> S
  mkState (B.LPS xs) =
      case xs of
          [] -> S B.empty L.empty
          (x:xs') -> S x (B.LPS xs')
  #endif
  
  -- | Run the Get monad applies a 'get'-based parser on the input ByteString
  runGet :: Get a -> L.ByteString -> a
  runGet m str = case unGet m (initState str) of (a, _) -> a
  
  -- | Run the Get monad applies a 'get'-based parser on the input
  -- ByteString. Additional to the result of get it returns the number of
  -- consumed bytes and the rest of the input.
  runGetState :: Get a -> L.ByteString -> Int64 -> (a, L.ByteString, Int64)
  runGetState m str off =
      case unGet m (mkState str off) of
        (a, ~(S s ss newOff)) -> (a, s `joinBS` ss, newOff)
  
  ------------------------------------------------------------------------
  
  failDesc :: String -> Get a
  failDesc err = do
      S _ _ bytes <- get
      Get (error (err ++ ". Failed reading at byte position " ++ show bytes))
  
  -- | Skip ahead @n@ bytes. Fails if fewer than @n@ bytes are available.
  skip :: Int -> Get ()
  skip n = readN (fromIntegral n) (const ())
  
  -- | Skip ahead @n@ bytes. No error if there isn't enough bytes.
  uncheckedSkip :: Int64 -> Get ()
  uncheckedSkip n = do
      S s ss bytes <- get
      if fromIntegral (B.length s) >= n
        then put (S (B.drop (fromIntegral n) s) ss (bytes + n))
        else do
          let rest = L.drop (n - fromIntegral (B.length s)) ss
          put $! mkState rest (bytes + n)
  
  -- | Run @ga@, but return without consuming its input.
  -- Fails if @ga@ fails.
  lookAhead :: Get a -> Get a
  lookAhead ga = do
      s <- get
      a <- ga
      put s
      return a
  
  -- | Like 'lookAhead', but consume the input if @gma@ returns 'Just _'.
  -- Fails if @gma@ fails.
  lookAheadM :: Get (Maybe a) -> Get (Maybe a)
  lookAheadM gma = do
      s <- get
      ma <- gma
      when (isNothing ma) $
          put s
      return ma
  
  -- | Like 'lookAhead', but consume the input if @gea@ returns 'Right _'.
  -- Fails if @gea@ fails.
  lookAheadE :: Get (Either a b) -> Get (Either a b)
  lookAheadE gea = do
      s <- get
      ea <- gea
      case ea of
          Left _ -> put s
          _      -> return ()
      return ea
  
  -- | Get the next up to @n@ bytes as a lazy ByteString, without consuming them. 
  uncheckedLookAhead :: Int64 -> Get L.ByteString
  uncheckedLookAhead n = do
      S s ss _ <- get
      if n <= fromIntegral (B.length s)
          then return (L.fromChunks [B.take (fromIntegral n) s])
          else return $ L.take n (s `joinBS` ss)
  
  ------------------------------------------------------------------------
  -- Utility
  
  -- | Get the total number of bytes read to this point.
  bytesRead :: Get Int64
  bytesRead = do
      S _ _ b <- get
      return b
  
  -- | Get the number of remaining unparsed bytes.
  -- Useful for checking whether all input has been consumed.
  -- Note that this forces the rest of the input.
  remaining :: Get Int64
  remaining = do
      S s ss _ <- get
      return (fromIntegral (B.length s) + L.length ss)
  
  -- | Test whether all input has been consumed,
  -- i.e. there are no remaining unparsed bytes.
  isEmpty :: Get Bool
  isEmpty = do
      S s ss _ <- get
      return (B.null s && L.null ss)
  
  ------------------------------------------------------------------------
  -- Utility with ByteStrings
  
  -- | An efficient 'get' method for strict ByteStrings. Fails if fewer
  -- than @n@ bytes are left in the input.
  getByteString :: Int -> Get B.ByteString
  getByteString n = readN n id
  {-# INLINE getByteString #-}
  
  -- | An efficient 'get' method for lazy ByteStrings. Does not fail if fewer than
  -- @n@ bytes are left in the input.
  getLazyByteString :: Int64 -> Get L.ByteString
  getLazyByteString n = do
      S s ss bytes <- get
      let big = s `joinBS` ss
      case splitAtST n big of
        (consume, rest) -> do put $ mkState rest (bytes + n)
                              return consume
  {-# INLINE getLazyByteString #-}
  
  -- | Get a lazy ByteString that is terminated with a NUL byte. Fails
  -- if it reaches the end of input without hitting a NUL.
  getLazyByteStringNul :: Get L.ByteString
  getLazyByteStringNul = do
      S s ss bytes <- get
      let big = s `joinBS` ss
          (consume, t) = L.break (== 0) big
          (h, rest) = L.splitAt 1 t
      if L.null h
        then fail "too few bytes"
        else do
          put $ mkState rest (bytes + L.length consume + 1)
          return consume
  {-# INLINE getLazyByteStringNul #-}
  
  -- | Get the remaining bytes as a lazy ByteString
  getRemainingLazyByteString :: Get L.ByteString
  getRemainingLazyByteString = do
      S s ss _ <- get
      return (s `joinBS` ss)
  
  ------------------------------------------------------------------------
  -- Helpers
  
  -- | Pull @n@ bytes from the input, as a strict ByteString.
  getBytes :: Int -> Get B.ByteString
  getBytes n = do
      S s ss bytes <- get
      if n <= B.length s
          then do let (consume,rest) = B.splitAt n s
                  put $! S rest ss (bytes + fromIntegral n)
                  return $! consume
          else
                case L.splitAt (fromIntegral n) (s `joinBS` ss) of
                  (consuming, rest) ->
                      do let now = B.concat . L.toChunks $ consuming
                         put $! mkState rest (bytes + fromIntegral n)
                         -- forces the next chunk before this one is returned
                         if (B.length now < n)
                           then
                              fail "too few bytes"
                           else
                              return now
  {- INLINE getBytes -}
  -- ^ important
  
  #ifndef BYTESTRING_IN_BASE
  joinBS :: B.ByteString -> L.ByteString -> L.ByteString
  joinBS bb lb
      | B.null bb = lb
      | otherwise = L.Chunk bb lb
  
  #else
  joinBS :: B.ByteString -> L.ByteString -> L.ByteString
  joinBS bb (B.LPS lb)
      | B.null bb = B.LPS lb
      | otherwise = B.LPS (bb:lb)
  #endif
      -- don't use L.append, it's strict in it's second argument :/
  {- INLINE joinBS -}
  
  -- | Split a ByteString. If the first result is consumed before the --
  -- second, this runs in constant heap space.
  --
  -- You must force the returned tuple for that to work, e.g.
  -- 
  -- > case splitAtST n xs of
  -- >    (ys,zs) -> consume ys ... consume zs
  --
  splitAtST :: Int64 -> L.ByteString -> (L.ByteString, L.ByteString)
  splitAtST i ps | i <= 0 = (L.empty, ps)
  #ifndef BYTESTRING_IN_BASE
  splitAtST i ps          = runST (
       do r  <- newSTRef undefined
          xs <- first r i ps
          ys <- unsafeInterleaveST (readSTRef r)
          return (xs, ys))
  
    where
          first r 0 xs@(L.Chunk _ _) = writeSTRef r xs    >> return L.Empty
          first r _ L.Empty          = writeSTRef r L.Empty >> return L.Empty
  
          first r n (L.Chunk x xs)
            | n < l     = do writeSTRef r (L.Chunk (B.drop (fromIntegral n) x) xs)
                             return $ L.Chunk (B.take (fromIntegral n) x) L.Empty
            | otherwise = do writeSTRef r (L.drop (n - l) xs)
                             liftM (L.Chunk x) $ unsafeInterleaveST (first r (n - l) xs)
  
           where l = fromIntegral (B.length x)
  #else
  splitAtST i (B.LPS ps)  = runST (
       do r  <- newSTRef undefined
          xs <- first r i ps
          ys <- unsafeInterleaveST (readSTRef r)
          return (B.LPS xs, B.LPS ys))
  
    where first r 0 xs     = writeSTRef r xs >> return []
          first r _ []     = writeSTRef r [] >> return []
          first r n (x:xs)
            | n < l     = do writeSTRef r (B.drop (fromIntegral n) x : xs)
                             return [B.take (fromIntegral n) x]
            | otherwise = do writeSTRef r (L.toChunks (L.drop (n - l) (B.LPS xs)))
                             fmap (x:) $ unsafeInterleaveST (first r (n - l) xs)
  
           where l = fromIntegral (B.length x)
  #endif
  {- INLINE splitAtST -}
  
  -- Pull n bytes from the input, and apply a parser to those bytes,
  -- yielding a value. If less than @n@ bytes are available, fail with an
  -- error. This wraps @getBytes@.
  readN :: Int -> (B.ByteString -> a) -> Get a
  readN n f = fmap f $ getBytes n
  {- INLINE readN -}
  -- ^ important
  
  ------------------------------------------------------------------------
  -- Primtives
  
  -- helper, get a raw Ptr onto a strict ByteString copied out of the
  -- underlying lazy byteString. So many indirections from the raw parser
  -- state that my head hurts...
  
  getPtr :: Storable a => Int -> Get a
  getPtr n = do
      (fp,o,_) <- readN n B.toForeignPtr
      return . B.inlinePerformIO $ withForeignPtr fp $ \p -> peek (castPtr $ p `plusPtr` o)
  {- INLINE getPtr -}
  
  ------------------------------------------------------------------------
  
  -- | Read a Word8 from the monad state
  getWord8 :: Get Word8
  getWord8 = getPtr (sizeOf (undefined :: Word8))
  {- INLINE getWord8 -}
  
  -- | Read a Word16 in big endian format
  getWord16be :: Get Word16
  getWord16be = do
      s <- readN 2 id
      return $! (fromIntegral (s `B.index` 0) `shiftl_w16` 8) .|.
                (fromIntegral (s `B.index` 1))
  {- INLINE getWord16be -}
  
  -- | Read a Word16 in little endian format
  getWord16le :: Get Word16
  getWord16le = do
      s <- readN 2 id
      return $! (fromIntegral (s `B.index` 1) `shiftl_w16` 8) .|.
                (fromIntegral (s `B.index` 0) )
  {- INLINE getWord16le -}
  
  -- | Read a Word32 in big endian format
  getWord32be :: Get Word32
  getWord32be = do
      s <- readN 4 id
      return $! (fromIntegral (s `B.index` 0) `shiftl_w32` 24) .|.
                (fromIntegral (s `B.index` 1) `shiftl_w32` 16) .|.
                (fromIntegral (s `B.index` 2) `shiftl_w32`  8) .|.
                (fromIntegral (s `B.index` 3) )
  {- INLINE getWord32be -}
  
  -- | Read a Word32 in little endian format
  getWord32le :: Get Word32
  getWord32le = do
      s <- readN 4 id
      return $! (fromIntegral (s `B.index` 3) `shiftl_w32` 24) .|.
                (fromIntegral (s `B.index` 2) `shiftl_w32` 16) .|.
                (fromIntegral (s `B.index` 1) `shiftl_w32`  8) .|.
                (fromIntegral (s `B.index` 0) )
  {- INLINE getWord32le -}
  
  -- | Read a Word64 in big endian format
  getWord64be :: Get Word64
  getWord64be = do
      s <- readN 8 id
      return $! (fromIntegral (s `B.index` 0) `shiftl_w64` 56) .|.
                (fromIntegral (s `B.index` 1) `shiftl_w64` 48) .|.
                (fromIntegral (s `B.index` 2) `shiftl_w64` 40) .|.
                (fromIntegral (s `B.index` 3) `shiftl_w64` 32) .|.
                (fromIntegral (s `B.index` 4) `shiftl_w64` 24) .|.
                (fromIntegral (s `B.index` 5) `shiftl_w64` 16) .|.
                (fromIntegral (s `B.index` 6) `shiftl_w64`  8) .|.
                (fromIntegral (s `B.index` 7) )
  {- INLINE getWord64be -}
  
  -- | Read a Word64 in little endian format
  getWord64le :: Get Word64
  getWord64le = do
      s <- readN 8 id
      return $! (fromIntegral (s `B.index` 7) `shiftl_w64` 56) .|.
                (fromIntegral (s `B.index` 6) `shiftl_w64` 48) .|.
                (fromIntegral (s `B.index` 5) `shiftl_w64` 40) .|.
                (fromIntegral (s `B.index` 4) `shiftl_w64` 32) .|.
                (fromIntegral (s `B.index` 3) `shiftl_w64` 24) .|.
                (fromIntegral (s `B.index` 2) `shiftl_w64` 16) .|.
                (fromIntegral (s `B.index` 1) `shiftl_w64`  8) .|.
                (fromIntegral (s `B.index` 0) )
  {- INLINE getWord64le -}
  
  ------------------------------------------------------------------------
  -- Host-endian reads
  
  -- | /O(1)./ Read a single native machine word. The word is read in
  -- host order, host endian form, for the machine you're on. On a 64 bit
  -- machine the Word is an 8 byte value, on a 32 bit machine, 4 bytes.
  getWordhost :: Get Word
  getWordhost = getPtr (sizeOf (undefined :: Word))
  {- INLINE getWordhost -}
  
  -- | /O(1)./ Read a 2 byte Word16 in native host order and host endianness.
  getWord16host :: Get Word16
  getWord16host = getPtr (sizeOf (undefined :: Word16))
  {- INLINE getWord16host -}
  
  -- | /O(1)./ Read a Word32 in native host order and host endianness.
  getWord32host :: Get Word32
  getWord32host = getPtr  (sizeOf (undefined :: Word32))
  {- INLINE getWord32host -}
  
  -- | /O(1)./ Read a Word64 in native host order and host endianess.
  getWord64host   :: Get Word64
  getWord64host = getPtr  (sizeOf (undefined :: Word64))
  {- INLINE getWord64host -}
  
  ------------------------------------------------------------------------
  -- Unchecked shifts
  
  shiftl_w16 :: Word16 -> Int -> Word16
  shiftl_w32 :: Word32 -> Int -> Word32
  shiftl_w64 :: Word64 -> Int -> Word64
  
  #if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
  shiftl_w16 (W16# w) (I# i) = W16# (w `uncheckedShiftL#`   i)
  shiftl_w32 (W32# w) (I# i) = W32# (w `uncheckedShiftL#`   i)
  
  #if WORD_SIZE_IN_BITS < 64
  shiftl_w64 (W64# w) (I# i) = W64# (w `uncheckedShiftL64#` i)
  
  #if __GLASGOW_HASKELL__ <= 606
  -- Exported by GHC.Word in GHC 6.8 and higher
  foreign import ccall unsafe "stg_uncheckedShiftL64"
      uncheckedShiftL64#     :: Word64# -> Int# -> Word64#
  #endif
  
  #else
  shiftl_w64 (W64# w) (I# i) = W64# (w `uncheckedShiftL#` i)
  #endif
  
  #else
  shiftl_w16 = shiftL
  shiftl_w32 = shiftL
  shiftl_w64 = shiftL
  #endif
  