{-# LANGUAGE FlexibleInstances,FlexibleContexts,MultiParamTypeClasses,CPP #-}
module Data.Encoding.ByteSource where

import Data.Encoding.Exception

import Data.Bits
import Data.Binary.Get
import Data.Char
import Data.Maybe
import Data.Word
import Control.Applicative as A
import Control.Monad (when)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.State (StateT (..), get, gets, put)
import Control.Monad.Identity (Identity)
import Control.Monad.Reader (ReaderT, ask)
import Control.Exception.Extensible
import Control.Throws
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import System.IO

class (Monad m,Throws DecodingException m) => ByteSource m where
    sourceEmpty :: m Bool
    fetchWord8 :: m Word8
    -- 'fetchAhead act' should return the same thing 'act' does, but should
    -- only consume input if 'act' returns a 'Just' value
    fetchAhead :: m (Maybe a) -> m (Maybe a)
    fetchWord16be :: m Word16
    fetchWord16be = do
      w1 <- m Word8
forall (m :: * -> *). ByteSource m => m Word8
fetchWord8
      w2 <- fetchWord8
      return $ ((fromIntegral w1) `shiftL` 8)
                 .|. (fromIntegral w2)
    fetchWord16le :: m Word16
    fetchWord16le = do
      w1 <- m Word8
forall (m :: * -> *). ByteSource m => m Word8
fetchWord8
      w2 <- fetchWord8
      return $ ((fromIntegral w2) `shiftL` 8)
                 .|. (fromIntegral w1)
    fetchWord32be :: m Word32
    fetchWord32be = do
      w1 <- m Word8
forall (m :: * -> *). ByteSource m => m Word8
fetchWord8
      w2 <- fetchWord8
      w3 <- fetchWord8
      w4 <- fetchWord8
      return $ ((fromIntegral w1) `shiftL` 24)
                 .|. ((fromIntegral w2) `shiftL` 16)
                 .|. ((fromIntegral w3) `shiftL`  8)
                 .|. (fromIntegral w4)
    fetchWord32le :: m Word32
    fetchWord32le = do
      w1 <- m Word8
forall (m :: * -> *). ByteSource m => m Word8
fetchWord8
      w2 <- fetchWord8
      w3 <- fetchWord8
      w4 <- fetchWord8
      return $ ((fromIntegral w4) `shiftL` 24)
                 .|. ((fromIntegral w3) `shiftL` 16)
                 .|. ((fromIntegral w2) `shiftL`  8)
                 .|. (fromIntegral w1)
    fetchWord64be :: m Word64
    fetchWord64be = do
      w1 <- m Word8
forall (m :: * -> *). ByteSource m => m Word8
fetchWord8
      w2 <- fetchWord8
      w3 <- fetchWord8
      w4 <- fetchWord8
      w5 <- fetchWord8
      w6 <- fetchWord8
      w7 <- fetchWord8
      w8 <- fetchWord8
      return $ ((fromIntegral w1) `shiftL` 56)
                 .|. ((fromIntegral w2) `shiftL` 48)
                 .|. ((fromIntegral w3) `shiftL` 40)
                 .|. ((fromIntegral w4) `shiftL` 32)
                 .|. ((fromIntegral w5) `shiftL` 24)
                 .|. ((fromIntegral w6) `shiftL` 16)
                 .|. ((fromIntegral w7) `shiftL`  8)
                 .|. (fromIntegral w8)
    fetchWord64le :: m Word64
    fetchWord64le = do
      w1 <- m Word8
forall (m :: * -> *). ByteSource m => m Word8
fetchWord8
      w2 <- fetchWord8
      w3 <- fetchWord8
      w4 <- fetchWord8
      w5 <- fetchWord8
      w6 <- fetchWord8
      w7 <- fetchWord8
      w8 <- fetchWord8
      return $ ((fromIntegral w8) `shiftL` 56)
                 .|. ((fromIntegral w7) `shiftL` 48)
                 .|. ((fromIntegral w6) `shiftL` 40)
                 .|. ((fromIntegral w5) `shiftL` 32)
                 .|. ((fromIntegral w4) `shiftL` 24)
                 .|. ((fromIntegral w3) `shiftL` 16)
                 .|. ((fromIntegral w2) `shiftL`  8)
                 .|. (fromIntegral w1)

instance Throws DecodingException Get where
    throwException :: forall a. DecodingException -> Get a
throwException = DecodingException -> Get a
forall a e. (HasCallStack, Exception e) => e -> a
throw

instance ByteSource Get where
    sourceEmpty :: Get Bool
sourceEmpty = Get Bool
isEmpty
    fetchWord8 :: Get Word8
fetchWord8 = Get Word8
getWord8
#if MIN_VERSION_binary(0,6,0)
    fetchAhead :: forall a. Get (Maybe a) -> Get (Maybe a)
fetchAhead Get (Maybe a)
act = (do
        res <- Get (Maybe a)
act
        case res of
            Maybe a
Nothing -> Get (Maybe a)
forall a. Get a
forall (f :: * -> *) a. Alternative f => f a
A.empty
            Just a
a  -> Maybe a -> Get (Maybe a)
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
res
        ) Get (Maybe a) -> Get (Maybe a) -> Get (Maybe a)
forall a. Get a -> Get a -> Get a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe a -> Get (Maybe a)
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
#else
    fetchAhead act = do
        res <- lookAhead act
        case res of
            Nothing -> return Nothing
            Just a  -> act
#endif
    fetchWord16be :: Get Word16
fetchWord16be = Get Word16
getWord16be
    fetchWord16le :: Get Word16
fetchWord16le = Get Word16
getWord16le
    fetchWord32be :: Get Word32
fetchWord32be = Get Word32
getWord32be
    fetchWord32le :: Get Word32
fetchWord32le = Get Word32
getWord32le
    fetchWord64be :: Get Word64
fetchWord64be = Get Word64
getWord64be
    fetchWord64le :: Get Word64
fetchWord64le = Get Word64
getWord64le

fetchAheadState :: m (Maybe a) -> m (Maybe a)
fetchAheadState m (Maybe a)
act = do
    chs <- m s
forall s (m :: * -> *). MonadState s m => m s
get
    res <- act
    when (isNothing res) (put chs)
    return res

instance ByteSource (StateT [Char] Identity) where
    sourceEmpty :: StateT [Char] Identity Bool
sourceEmpty = ([Char] -> Bool) -> StateT [Char] Identity Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets [Char] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null
    fetchWord8 :: StateT [Char] Identity Word8
fetchWord8 = do
      chs <- StateT [Char] Identity [Char]
forall s (m :: * -> *). MonadState s m => m s
get
      case chs of
        [] -> DecodingException -> StateT [Char] Identity Word8
forall a. DecodingException -> StateT [Char] Identity a
forall e (m :: * -> *) a. Throws e m => e -> m a
throwException DecodingException
UnexpectedEnd
        Char
c:[Char]
cs -> do
          [Char] -> StateT [Char] Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put [Char]
cs
          Word8 -> StateT [Char] Identity Word8
forall a. a -> StateT [Char] Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> Int -> Word8
forall a b. (a -> b) -> a -> b
$ Char -> Int
ord Char
c)
    fetchAhead :: forall a.
StateT [Char] Identity (Maybe a)
-> StateT [Char] Identity (Maybe a)
fetchAhead = StateT [Char] Identity (Maybe a)
-> StateT [Char] Identity (Maybe a)
forall {m :: * -> *} {s} {a}.
MonadState s m =>
m (Maybe a) -> m (Maybe a)
fetchAheadState

#if MIN_VERSION_base(4,3,0)
#else
instance Monad (Either DecodingException) where
    return = Right
    (Left err) >>= g = Left err
    (Right x) >>= g = g x
#endif

instance ByteSource (StateT [Char] (Either DecodingException)) where
    sourceEmpty :: StateT [Char] (Either DecodingException) Bool
sourceEmpty = ([Char] -> Bool) -> StateT [Char] (Either DecodingException) Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets [Char] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null
    fetchWord8 :: StateT [Char] (Either DecodingException) Word8
fetchWord8 = do
      chs <- StateT [Char] (Either DecodingException) [Char]
forall s (m :: * -> *). MonadState s m => m s
get
      case chs of
        [] -> DecodingException -> StateT [Char] (Either DecodingException) Word8
forall a.
DecodingException -> StateT [Char] (Either DecodingException) a
forall e (m :: * -> *) a. Throws e m => e -> m a
throwException DecodingException
UnexpectedEnd
        Char
c:[Char]
cs -> do
          [Char] -> StateT [Char] (Either DecodingException) ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put [Char]
cs
          Word8 -> StateT [Char] (Either DecodingException) Word8
forall a. a -> StateT [Char] (Either DecodingException) a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> Int -> Word8
forall a b. (a -> b) -> a -> b
$ Char -> Int
ord Char
c)
    fetchAhead :: forall a.
StateT [Char] (Either DecodingException) (Maybe a)
-> StateT [Char] (Either DecodingException) (Maybe a)
fetchAhead = StateT [Char] (Either DecodingException) (Maybe a)
-> StateT [Char] (Either DecodingException) (Maybe a)
forall {m :: * -> *} {s} {a}.
MonadState s m =>
m (Maybe a) -> m (Maybe a)
fetchAheadState

instance (Monad m,Throws DecodingException m) => ByteSource (StateT BS.ByteString m) where
    sourceEmpty :: StateT ByteString m Bool
sourceEmpty = (ByteString -> Bool) -> StateT ByteString m Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ByteString -> Bool
BS.null
    fetchWord8 :: StateT ByteString m Word8
fetchWord8 = (ByteString -> m (Word8, ByteString)) -> StateT ByteString m Word8
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT (\ByteString
str -> case ByteString -> Maybe (Word8, ByteString)
BS.uncons ByteString
str of
                                  Maybe (Word8, ByteString)
Nothing -> DecodingException -> m (Word8, ByteString)
forall a. DecodingException -> m a
forall e (m :: * -> *) a. Throws e m => e -> m a
throwException DecodingException
UnexpectedEnd
                                  Just (Word8
c,ByteString
cs) -> (Word8, ByteString) -> m (Word8, ByteString)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Word8
c,ByteString
cs))
    fetchAhead :: forall a.
StateT ByteString m (Maybe a) -> StateT ByteString m (Maybe a)
fetchAhead = StateT ByteString m (Maybe a) -> StateT ByteString m (Maybe a)
forall {m :: * -> *} {s} {a}.
MonadState s m =>
m (Maybe a) -> m (Maybe a)
fetchAheadState

instance ByteSource (StateT LBS.ByteString (Either DecodingException)) where
    sourceEmpty :: StateT ByteString (Either DecodingException) Bool
sourceEmpty = (ByteString -> Bool)
-> StateT ByteString (Either DecodingException) Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ByteString -> Bool
LBS.null
    fetchWord8 :: StateT ByteString (Either DecodingException) Word8
fetchWord8 = (ByteString -> Either DecodingException (Word8, ByteString))
-> StateT ByteString (Either DecodingException) Word8
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT (\ByteString
str -> case ByteString -> Maybe (Word8, ByteString)
LBS.uncons ByteString
str of
                                  Maybe (Word8, ByteString)
Nothing -> DecodingException -> Either DecodingException (Word8, ByteString)
forall a b. a -> Either a b
Left DecodingException
UnexpectedEnd
                                  Just (Word8, ByteString)
ns -> (Word8, ByteString) -> Either DecodingException (Word8, ByteString)
forall a b. b -> Either a b
Right (Word8, ByteString)
ns)
    fetchAhead :: forall a.
StateT ByteString (Either DecodingException) (Maybe a)
-> StateT ByteString (Either DecodingException) (Maybe a)
fetchAhead = StateT ByteString (Either DecodingException) (Maybe a)
-> StateT ByteString (Either DecodingException) (Maybe a)
forall {m :: * -> *} {s} {a}.
MonadState s m =>
m (Maybe a) -> m (Maybe a)
fetchAheadState

instance ByteSource (ReaderT Handle IO) where
    sourceEmpty :: ReaderT Handle IO Bool
sourceEmpty = do
      h <- ReaderT Handle IO Handle
forall r (m :: * -> *). MonadReader r m => m r
ask
      liftIO (hIsEOF h)
    fetchWord8 :: ReaderT Handle IO Word8
fetchWord8 = do
      h <- ReaderT Handle IO Handle
forall r (m :: * -> *). MonadReader r m => m r
ask
      liftIO $ do
        ch <- hGetChar h
        return (fromIntegral $ ord ch)
    fetchAhead :: forall a.
ReaderT Handle IO (Maybe a) -> ReaderT Handle IO (Maybe a)
fetchAhead ReaderT Handle IO (Maybe a)
act = do
      h <- ReaderT Handle IO Handle
forall r (m :: * -> *). MonadReader r m => m r
ask
      pos <- liftIO $ hGetPosn h
      res <- act
      when (isNothing res) (liftIO $ hSetPosn pos)
      return res