-- |
-- Module      : Data.ASN1.Get
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : experimental
-- Portability : unknown
--
-- Simple get module with really simple accessor for ASN1.
--
-- Original code is pulled from the Get module from cereal
-- which is covered by:
-- Copyright   : Lennart Kolmodin, Galois Inc. 2009
-- License     : BSD3-style (see LICENSE)
--
-- The original code has been tailored and reduced to only cover the useful
-- case for asn1 and augmented by a position.
--
{-# LANGUAGE Rank2Types #-}
module Data.ASN1.Get
    ( Result(..)
    , Input
    , Get
    , runGetPos
    , runGet
    , getBytes
    , getBytesCopy
    , getWord8
    ) where

import Control.Applicative (Applicative(..),Alternative(..))
import Control.Monad (ap,MonadPlus(..))
import Data.Maybe (fromMaybe)
import Foreign

import qualified Data.ByteString          as B
import qualified Data.ByteString.Unsafe   as B

-- | The result of a parse.
data Result r = Fail String
              -- ^ The parse failed. The 'String' is the
              --   message describing the error, if any.
              | Partial (B.ByteString -> Result r)
              -- ^ Supply this continuation with more input so that
              --   the parser can resume. To indicate that no more
              --   input is available, use an 'B.empty' string.
              | Done r Position B.ByteString
              -- ^ The parse succeeded.  The 'B.ByteString' is the
              --   input that had not yet been consumed (if any) when
              --   the parse succeeded.

instance Show r => Show (Result r) where
    show (Fail msg)  = "Fail " ++ show msg
    show (Partial _) = "Partial _"
    show (Done r pos bs) = "Done " ++ show r ++ " " ++ show pos ++ " " ++ show bs

instance Functor Result where
    fmap _ (Fail msg)  = Fail msg
    fmap f (Partial k) = Partial (fmap f . k)
    fmap f (Done r p bs) = Done (f r) p bs

type Input  = B.ByteString
type Buffer = Maybe B.ByteString

type Failure   r = Input -> Buffer -> More -> Position -> String -> Result r
type Success a r = Input -> Buffer -> More -> Position -> a      -> Result r
type Position    = Word64

-- | Have we read all available input?
data More = Complete
          | Incomplete (Maybe Int)
          deriving (Eq)

-- | The Get monad is an Exception and State monad.
newtype Get a = Get
	{ unGet :: forall r. Input -> Buffer -> More -> Position -> Failure r -> Success a r -> Result r }

append :: Buffer -> Buffer -> Buffer
append l r = B.append `fmap` l <*> r
{-# INLINE append #-}

bufferBytes :: Buffer -> B.ByteString
bufferBytes  = fromMaybe B.empty
{-# INLINE bufferBytes #-}

instance Functor Get where
    fmap p m =
      Get $ \s0 b0 m0 p0 kf ks ->
        let ks' s1 b1 m1 p1 a = ks s1 b1 m1 p1 (p a)
         in unGet m s0 b0 m0 p0 kf ks'

instance Applicative Get where
    pure  = return
    (<*>) = ap

instance Alternative Get where
    empty = failDesc "empty"
    (<|>) = mplus

-- Definition directly from Control.Monad.State.Strict
instance Monad Get where
    return a = Get $ \ s0 b0 m0 p0 _ ks -> ks s0 b0 m0 p0 a

    m >>= g  = Get $ \s0 b0 m0 p0 kf ks ->
        let ks' s1 b1 m1 p1 a = unGet (g a) s1 b1 m1 p1 kf ks
         in unGet m s0 b0 m0 p0 kf ks'

    fail     = failDesc

instance MonadPlus Get where
    mzero     = failDesc "mzero"
    mplus a b =
      Get $ \s0 b0 m0 p0 kf ks ->
        let kf' _ b1 m1 p1 _ = unGet b (s0 `B.append` bufferBytes b1)
                                       (b0 `append` b1) m1 p1 kf ks
         in unGet a s0 (Just B.empty) m0 p0 kf' ks

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

put :: Position -> B.ByteString -> Get ()
put pos s = Get (\_ b0 m p0 _ k -> k s b0 m (p0+pos) ())
{-# INLINE put #-}

finalK :: B.ByteString -> t -> t1 -> Position -> r -> Result r
finalK s _ _ p a = Done a p s

failK :: Failure a
failK _ _ _ p s = Fail (show p ++ ":" ++ s)

-- | Run the Get monad applies a 'get'-based parser on the input ByteString
runGetPos :: Position -> Get a -> B.ByteString -> Result a
runGetPos pos m str = unGet m str Nothing (Incomplete Nothing) pos failK finalK
{-# INLINE runGetPos #-}

runGet :: Get a -> B.ByteString -> Result a
runGet = runGetPos 0
{-# INLINE runGet #-}

-- | If at least @n@ bytes of input are available, return the current
--   input, otherwise fail.
ensure :: Int -> Get B.ByteString
ensure n = n `seq` Get $ \ s0 b0 m0 p0 kf ks ->
    if B.length s0 >= n
    then ks s0 b0 m0 p0 s0
    else unGet (demandInput >> ensureRec n) s0 b0 m0 p0 kf ks
{-# INLINE ensure #-}

-- | If at least @n@ bytes of input are available, return the current
--   input, otherwise fail.
ensureRec :: Int -> Get B.ByteString
ensureRec n = Get $ \s0 b0 m0 p0 kf ks ->
    if B.length s0 >= n
    then ks s0 b0 m0 p0 s0
    else unGet (demandInput >> ensureRec n) s0 b0 m0 p0 kf ks

-- | Immediately demand more input via a 'Partial' continuation
--   result.
demandInput :: Get ()
demandInput = Get $ \s0 b0 m0 p0 kf ks ->
  case m0 of
    Complete      -> kf s0 b0 m0 p0 "too few bytes"
    Incomplete mb -> Partial $ \s ->
      if B.null s
      then kf s0 b0 m0 p0 "too few bytes"
      else let update l = l - B.length s
               s1 = s0 `B.append` s
               b1 = b0 `append` Just s
            in ks s1 b1 (Incomplete (update `fmap` mb)) p0 ()

failDesc :: String -> Get a
failDesc err = Get (\s0 b0 m0 p0 kf _ -> kf s0 b0 m0 p0 ("Failed reading: " ++ err))

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

-- | An efficient 'get' method for strict ByteStrings. Fails if fewer
-- than @n@ bytes are left in the input. This function creates a fresh
-- copy of the underlying bytes.
getBytesCopy :: Int -> Get B.ByteString
getBytesCopy n = do
  bs <- getBytes n
  return $! B.copy bs

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

-- | Pull @n@ bytes from the input, as a strict ByteString.
getBytes :: Int -> Get B.ByteString
getBytes n = do
     s <- ensure n
     put (fromIntegral n) $ B.unsafeDrop n s
     return $ B.unsafeTake n s

getWord8 :: Get Word8
getWord8 = do
     s <- ensure 1
     put 1 $ B.unsafeTail s
     return $ B.unsafeHead s