IncrementalGet.hs

{-# LANGUAGE CPP #-}
{-# OPTIONS_GHC -fglasgow-exts #-}
-- for unboxed shifts

-----------------------------------------------------------------------------
-- |
-- Module      : Data.Binary.Strict.IncrementalGet
-- Copyright   : Lennart Kolmodin
-- License     : BSD3-style (see LICENSE)
--
-- Maintainer  : Adam Langley <agl@imperialviolet.org>
-- Stability   : experimental
-- Portability : portable to Hugs and GHC.
--
-- This is a version of the Get monad for incremental parsing. The parser is
-- written as if a single, huge, strict ByteString was to be parsed. It
-- produces results as it parses by calling yield.
--
-- However, if the parser runs out of data, rather than failing the caller sees
-- a Partial result, which includes the list of yielded values so far and a
-- continuation. By calling the continuation with more data, the parser
-- continues, none the wiser.
--
-- Take the following example
--
-- > testParse = do
-- >   a <- getWord16be
-- >   b <- getWord16be
-- >   return $ a + b
-- >
-- > test = runGet testParse $ B.pack [1,0,0]
--
-- Here @testParse@ needs to read 4 bytes in order to complete, so test is
-- a Partial, which includes the continuation function, so which you can pass
-- more data until it completes
--
-- The lookahead functions have been removed from this parser because of their
-- incompatibility with the incremental monad at the moment.
-----------------------------------------------------------------------------

#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
#include "MachDeps.h"
#endif

#include "Common.h"

module Data.Binary.Strict.IncrementalGet (
    -- * The Get type
      Get
    , Result(..)
    , runGet

    -- * Utility
    , skip
    , bytesRead
    , remaining
    , isEmpty

    -- * Parsing particular types
    , getWord8

    -- ** ByteStrings
    , getByteString

    -- ** Big-endian reads
    , getWord16be
    , getWord32be
    , getWord64be

    -- ** Little-endian reads
    , getWord16le
    , getWord32le
    , getWord64le

    -- ** Host-endian, unaligned reads
    , getWordhost
    , getWord16host
    , getWord32host
    , getWord64host
) where

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

import Foreign

#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
import GHC.Base
import GHC.Word
#endif

#ifndef __HADDOCK__
-- | The parse state
data S = S {-# UNPACK #-} !B.ByteString  -- ^ input
           {-# UNPACK #-} !Int  -- ^ bytes read
#endif

-- | The result of a partial parse
data Result a = Failed String
                -- ^ the parse failed with the given error message
              | Finished B.ByteString a
                -- ^ the parse finished and produced the given list of
                --   results doing so. Any unparsed data is returned.
              | Partial (B.ByteString -> Result a)
                -- ^ the parse ran out of data before finishing, but produced
                --   the given list of results before doing so. To continue the
                --   parse pass more data to the given continuation

instance (Show a) => Show (Result a) where
  show (Failed err) = "Failed " ++ err
  show (Finished _ rs) = "Finished " ++ show rs
  show (Partial _) = "Partial"

newtype Get r a = Get { unGet :: S -> (a -> S -> Result r) -> Result r }

instance Functor (Get r) where
    fmap f m = Get (\s -> \cont -> unGet m s (cont . f))

instance Monad (Get r) where
  return a = Get (\s -> \k -> k a s)
  m >>= k = Get (\s -> \cont -> unGet m s (\a -> \s' -> unGet (k a) s' cont))
  fail err = Get (const $ const $ Failed err)

get :: Get r S
get = Get (\s -> \k -> k s s)

initState :: B.ByteString -> S
initState input = S input 0
{-# INLINE initState #-}

-- | Start a parser and return the first Result.
runGet :: Get r r -> B.ByteString -> Result r
runGet m input =
  unGet m (initState input) (\v -> (\(S s _) -> Finished s v))

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

-- | Get the total number of bytes read to this point.
bytesRead :: Get r Int
bytesRead = do
  S _ b <- get
  return b

-- | Get the number of remaining unparsed bytes.
-- Useful for checking whether all input has been consumed.
remaining :: Get r Int
remaining = do
  S s _ <- get
  return (fromIntegral (B.length s))

-- | Test whether all input has been consumed,
-- i.e. there are no remaining unparsed bytes.
isEmpty :: Get r Bool
isEmpty = do
  S s _ <- get
  return $ B.null s

------------------------------------------------------------------------
-- Utility with ByteStrings

-- | An efficient 'get' method for strict ByteStrings. Fails if fewer
-- than @n@ bytes are left in the input.
getByteString :: Int -> Get r B.ByteString
getByteString n = readN n id
{-# INLINE getByteString #-}

-- | Pull @n@ bytes from the input, as a strict ByteString.
getBytes :: Int -> Get r B.ByteString
getBytes n = Get $ \(S s offset) -> \cont ->
  if n <= B.length s
     then let (consume, rest) = B.splitAt n s
           in cont consume $ S rest (offset + fromIntegral n)
     else Partial (\s' -> unGet (getBytes n) (S (B.append s s') offset) cont)
{-# INLINE getBytes #-}

-- 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 r a
readN n f = fmap f $ getBytes n
{-# INLINE readN #-}

getPtr :: Storable a => Int -> Get r a
getPtr n = do
    (fp, o, _) <- readN n B.toForeignPtr
    return . B.inlinePerformIO $ withForeignPtr fp $ \p -> peek (castPtr $ p `plusPtr` o)
{-# INLINE getPtr #-}

GETWORDS(Get r, getBytes)
GETHOSTWORDS(Get r)

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)
#else
shiftl_w64 (W64# w) (I# i) = W64# (w `uncheckedShiftL#` i)
#endif

#else
shiftl_w16 = shiftL
shiftl_w32 = shiftL
shiftl_w64 = shiftL
#endif