-----------------------------------------------------------------------------
-- |
-- Module      : Data.ByteString.Parser
-- Copyright   : Lennart Kolmodin, George Giorgidze
-- License     : BSD3
--
-- Maintainer  : George Giorgidze <http://cs.nott.ac.uk/~ggg/>
-- Stability   : experimental
-- Portability : Portable
--
-- A monad for efficiently building structures from
-- encoded lazy ByteStrings.
--
-----------------------------------------------------------------------------

module Codec.ByteString.Parser (

    -- * The Parser type
      Parser
    , runParser
    , runParserState

    -- * Parsing
    , choice
    , expect
    , skip
    , lookAhead
    , lookAheadM
    , lookAheadE

    -- * Utility
    , bytesRead
    , getBytes
    , remaining
    , isEmpty

    -- * Parsing particular types
    , satisfy
    , getString
    , getStringNul
    , string
    , getWord8
    , getInt8
    , word8
    , int8

    -- ** ByteStrings
    , getByteString
    , getLazyByteString
    , getLazyByteStringNul
    , getRemainingLazyByteString

    -- ** Big-endian reads
    , getWord16be
    , word16be
    , getWord24be
    , word24be
    , getWord32be
    , word32be
    , getWord64be
    , word64be

    , getInt16be
    , int16be
    , getInt32be
    , int32be
    , getInt64be
    , int64be

    -- ** Little-endian reads
    , getWord16le
    , word16le
    , getWord24le
    , word24le
    , getWord32le
    , word32le
    , getWord64le
    , word64le

    , getInt16le
    , int16le
    , getInt32le
    , int32le
    , getInt64le
    , int64le

    -- ** Host-endian, unaligned reads
    , getWordHost
    , wordHost
    , getWord16host
    , word16host
    , getWord32host
    , word32host
    , getWord64host
    , word64host

    -- Variable length reads
    , getVarLenBe
    , varLenBe
    , getVarLenLe
    , varLenLe
  ) where

import qualified Data.ByteString               as B
import qualified Data.ByteString.Lazy          as L
import qualified Data.ByteString.Internal      as B
import qualified Data.ByteString.Lazy.Internal as L

import Foreign.Storable        (Storable, peek, sizeOf)
import Foreign.Ptr             (plusPtr, castPtr)
import Foreign.ForeignPtr      (withForeignPtr)
import Control.Monad.ST        (runST)
import Control.Monad.ST.Unsafe (unsafeInterleaveST)

import Control.Monad
import Control.Applicative
import Data.STRef
import Data.Word
import Data.Int
import Data.Bits
import Data.Maybe

-- | 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 Parser a = Parser { unParser :: S -> Either String (a, S) }

instance Functor Parser where
    fmap f m = Parser $ \s -> case unParser m s of
      Left e -> Left e
      Right (a, s') -> Right (f a, s')
    
instance Monad Parser where
    return a  = Parser (\s -> Right (a, s))
    m >>= k   = Parser $ \s -> case (unParser m) s of
      Left e -> Left e
      Right (a, s') -> (unParser (k a)) s'
    fail  err  = Parser $ \(S _ _ bytes) ->
        Left (err ++ ". Failed reading at byte position " ++ show bytes)
instance MonadPlus Parser where
  mzero = Parser $ \_ -> Left []
  mplus p1 p2 = Parser $ \s -> case (unParser p1 s) of
    Left e1 -> case (unParser p2 s) of
      Left e2 -> Left (e1 ++ "\n" ++ e2)
      ok -> ok
    ok -> ok

instance Applicative Parser where
  pure  = return
  (<*>) = ap
  
instance Alternative Parser where
  empty = mzero
  (<|>) = mplus

------------------------------------------------------------------------

get :: Parser S
get   = Parser $ \s -> Right (s, s)

put :: S -> Parser ()
put s = Parser $ \_ -> Right ((), s)

------------------------------------------------------------------------

initState :: L.ByteString -> S
initState xs = mkState xs 0

mkState :: L.ByteString -> Int64 -> S
mkState l = case l of
    L.Empty      -> S B.empty L.empty
    L.Chunk x xs -> S x xs

-- | Run the Get monad applies a 'get'-based parser on the input ByteString
runParser :: Parser a -> L.ByteString -> Either String a
runParser m str = case unParser m (initState str) of
  Left e -> Left e
  Right (a, _) -> Right 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.
runParserState :: Parser a -> L.ByteString -> Int64 -> Either String (a, L.ByteString, Int64)
runParserState m str off =
    case unParser m (mkState str off) of
      Left e -> Left e
      Right (a, ~(S s ss newOff)) -> Right (a, s `bsJoin` ss, newOff)

------------------------------------------------------------------------

choice :: [Parser a] -> Parser a
choice = foldl (<|>) mzero

-- | Skip ahead @n@ bytes. Fails if fewer than @n@ bytes are available.
skip :: Word64 -> Parser ()
skip n = readN (fromIntegral n) (const ())

-- | Run @ga@, but return without consuming its input.
-- Fails if @ga@ fails.
lookAhead :: Parser a -> Parser 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 :: Parser (Maybe a) -> Parser (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 :: Parser (Either a b) -> Parser (Either a b)
lookAheadE gea = do
    s <- get
    ea <- gea
    case ea of
        Left _ -> put s
        _      -> return ()
    return ea

expect :: (Show a, Eq a) => (a -> Bool) -> Parser a -> Parser a
expect f p = do
  v <- p
  when (not $ f v) $ fail $ show v ++ " was not expected."
  return v

getString :: Int -> Parser String
getString l = do
  bs <- getLazyByteString (fromIntegral l)
  return $! map B.w2c (L.unpack bs)

getStringNul :: Parser String
getStringNul = do
  bs <- getLazyByteStringNul
  return $! map B.w2c (L.unpack bs)

string :: String -> Parser String
string s = expect (s ==) (getString $ length s)

-- Utility

-- | Get the total number of bytes read to this point.
bytesRead :: Parser 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 :: Parser 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 :: Parser 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 -> Parser B.ByteString
getByteString n = readN n id

-- | An efficient 'get' method for lazy ByteStrings. Does not fail if fewer than
-- @n@ bytes are left in the input.
getLazyByteString :: Int64 -> Parser L.ByteString
getLazyByteString n = do
    S s ss bytes <- get
    let big = s `bsJoin` ss
    case splitAtST n big of
      (consume, rest) -> do put $ mkState rest (bytes + n)
                            return consume

-- | Get a lazy ByteString that is terminated with a NUL byte. Fails
-- if it reaches the end of input without hitting a NUL.
getLazyByteStringNul :: Parser L.ByteString
getLazyByteStringNul = do
    S s ss bytes <- get
    let big = s `bsJoin` ss
        (consume, t) = L.break (== 0) big
        (h, rest) = L.splitAt 1 t
    when (L.null h) $ fail "too few bytes"
    put $ mkState rest (bytes + L.length consume + 1)
    return consume

-- | Get the remaining bytes as a lazy ByteString
getRemainingLazyByteString :: Parser L.ByteString
getRemainingLazyByteString = do
    S s ss _ <- get
    return $! (s `bsJoin` ss)

------------------------------------------------------------------------
-- Helpers

-- | Pull @n@ bytes from the input, as a strict ByteString.
getBytes :: Int -> Parser 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 `bsJoin` 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
                       when (B.length now < n) $ fail "too few bytes"
                       return now

bsJoin :: B.ByteString -> L.ByteString -> L.ByteString
bsJoin bb lb
    | B.null bb = lb
    | otherwise = L.Chunk bb lb

-- | 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)
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)

-- 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) -> Parser a
readN n f = fmap f $ getBytes n


------------------------------------------------------------------------
-- 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 -> Parser a
getPtr n = do
    (fp,o,_) <- readN n B.toForeignPtr
    return . B.inlinePerformIO $ withForeignPtr fp $ \p -> peek (castPtr $ p `plusPtr` o)

------------------------------------------------------------------------

satisfy :: (Word8 -> Bool) -> Parser Word8
satisfy f = do
  w <- getWord8
  guard (f w)
  return w

-- | Read a Word8 from the monad state
getWord8 :: Parser Word8
getWord8 = getPtr (sizeOf (undefined :: Word8))

word8 :: Word8 -> Parser Word8
word8 w = expect (w ==) getWord8

-- | Read a Word16 in big endian format
getWord16be :: Parser Word16
getWord16be = do
    s <- readN 2 id
    return $! (fromIntegral (s `B.index` 0) `shiftL` 8) .|.
              (fromIntegral (s `B.index` 1))

word16be :: Word16 -> Parser Word16
word16be w = expect (w ==) getWord16be

-- | Read a Word16 in little endian format
getWord16le :: Parser Word16
getWord16le = do
    s <- readN 2 id
    return $! (fromIntegral (s `B.index` 1) `shiftL` 8) .|.
              (fromIntegral (s `B.index` 0) )

word16le :: Word16 -> Parser Word16
word16le w = expect (w ==) getWord16le

-- | Read a 24 bit word into Word32 in big endian format
getWord24be :: Parser Word32
getWord24be = do
    s <- readN 3 id
    return $! (fromIntegral (s `B.index` 0) `shiftL` 16) .|.
              (fromIntegral (s `B.index` 1) `shiftL`  8) .|.
              (fromIntegral (s `B.index` 2) )

word24be :: Word32 -> Parser Word32
word24be w = expect (w ==) getWord24be

getWord24le :: Parser Word32
getWord24le = do
    s <- readN 3 id
    return $! (fromIntegral (s `B.index` 2) `shiftL` 16) .|.
              (fromIntegral (s `B.index` 1) `shiftL`  8) .|.
              (fromIntegral (s `B.index` 0) )

word24le :: Word32 -> Parser Word32
word24le w = expect (w ==) getWord24le

-- | Read a Word32 in big endian format
getWord32be :: Parser Word32
getWord32be = do
    s <- readN 4 id
    return $! (fromIntegral (s `B.index` 0) `shiftL` 24) .|.
              (fromIntegral (s `B.index` 1) `shiftL` 16) .|.
              (fromIntegral (s `B.index` 2) `shiftL`  8) .|.
              (fromIntegral (s `B.index` 3) )

word32be :: Word32 -> Parser Word32
word32be w = expect (w ==) getWord32be

-- | Read a Word32 in little endian format
getWord32le :: Parser Word32
getWord32le = do
    s <- readN 4 id
    return $! (fromIntegral (s `B.index` 3) `shiftL` 24) .|.
              (fromIntegral (s `B.index` 2) `shiftL` 16) .|.
              (fromIntegral (s `B.index` 1) `shiftL`  8) .|.
              (fromIntegral (s `B.index` 0) )

word32le :: Word32 -> Parser Word32
word32le w = expect (w ==) getWord32le


-- | Read a Word64 in big endian format
getWord64be :: Parser Word64
getWord64be = do
    s <- readN 8 id
    return $! (fromIntegral (s `B.index` 0) `shiftL` 56) .|.
              (fromIntegral (s `B.index` 1) `shiftL` 48) .|.
              (fromIntegral (s `B.index` 2) `shiftL` 40) .|.
              (fromIntegral (s `B.index` 3) `shiftL` 32) .|.
              (fromIntegral (s `B.index` 4) `shiftL` 24) .|.
              (fromIntegral (s `B.index` 5) `shiftL` 16) .|.
              (fromIntegral (s `B.index` 6) `shiftL`  8) .|.
              (fromIntegral (s `B.index` 7) )

word64be :: Word64 -> Parser Word64
word64be w = expect (w ==) getWord64be

-- | Read a Word64 in little endian format
getWord64le :: Parser Word64
getWord64le = do
    s <- readN 8 id
    return $! (fromIntegral (s `B.index` 7) `shiftL` 56) .|.
              (fromIntegral (s `B.index` 6) `shiftL` 48) .|.
              (fromIntegral (s `B.index` 5) `shiftL` 40) .|.
              (fromIntegral (s `B.index` 4) `shiftL` 32) .|.
              (fromIntegral (s `B.index` 3) `shiftL` 24) .|.
              (fromIntegral (s `B.index` 2) `shiftL` 16) .|.
              (fromIntegral (s `B.index` 1) `shiftL`  8) .|.
              (fromIntegral (s `B.index` 0) )

word64le :: Word64 -> Parser Word64
word64le w = expect (w ==) getWord64le
------------------------------------------------------------------------
getInt8 :: Parser Int8
getInt8 = getWord8 >>= return . fromIntegral

int8 :: Int8 -> Parser Int8
int8 i = expect (i ==) getInt8

getInt16le :: Parser Int16
getInt16le = getWord16le >>= return . fromIntegral

int16le :: Int16 -> Parser Int16
int16le i = expect (i ==) getInt16le

getInt16be :: Parser Int16
getInt16be = getWord16be >>= return . fromIntegral

int16be :: Int16 -> Parser Int16
int16be i = expect (i ==) getInt16be

getInt32le :: Parser Int32
getInt32le = getWord32le >>= return . fromIntegral

int32le :: Int32 -> Parser Int32
int32le i = expect (i ==) getInt32le

getInt32be :: Parser Int32
getInt32be = getWord32be >>= return . fromIntegral

int32be :: Int32 -> Parser Int32
int32be i = expect (i ==) getInt32be

getInt64le :: Parser Int64
getInt64le = getWord64le >>= return . fromIntegral

int64le :: Int64 -> Parser Int64
int64le i = expect (i ==) getInt64le

getInt64be :: Parser Int64
getInt64be = getWord64be >>= return . fromIntegral

int64be :: Int64 -> Parser Int64
int64be i = expect (i ==) getInt64be

------------------------------------------------------------------------
-- 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 :: Parser Word
getWordHost = getPtr (sizeOf (undefined :: Word))

wordHost :: Word -> Parser Word
wordHost w = expect (w ==) getWordHost

-- | /O(1)./ Read a 2 byte Word16 in native host order and host endianness.
getWord16host :: Parser Word16
getWord16host = getPtr (sizeOf (undefined :: Word16))

word16host :: Word16 -> Parser Word16
word16host w = expect (w ==) getWord16host

-- | /O(1)./ Read a Word32 in native host order and host endianness.
getWord32host :: Parser Word32
getWord32host = getPtr  (sizeOf (undefined :: Word32))

word32host :: Word32 -> Parser Word32
word32host w = expect (w ==) getWord32host

-- | /O(1)./ Read a Word64 in native host order and host endianess.
getWord64host   :: Parser Word64
getWord64host = getPtr  (sizeOf (undefined :: Word64))

word64host :: Word64 -> Parser Word64
word64host w = expect (w ==) getWord64host

-- Variable length numbers

getVarLenBe :: Parser Word64
getVarLenBe = f 0
  where
  f :: Word64 -> Parser Word64
  f acc =  do
    w <- getWord8 >>= return . fromIntegral
    if testBit w 7
      then f      $! (shiftL acc 7) .|. (clearBit w 7)
      else return $! (shiftL acc 7) .|. w

varLenBe :: Word64 -> Parser Word64
varLenBe a = expect (a ==) getVarLenBe

getVarLenLe :: Parser Word64
getVarLenLe = do
  w <- getWord8 >>= return . fromIntegral
  if testBit w 7
    then do
      w' <- getVarLenLe
      return $! (clearBit w 7) .|. (shiftL w' 7)
    else return $! w

varLenLe :: Word64 -> Parser Word64
varLenLe a = expect (a ==) getVarLenLe