{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE DeriveDataTypeable     #-}
{-# LANGUAGE DeriveGeneric          #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE KindSignatures         #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE OverloadedStrings      #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# OPTIONS_HADDOCK not-home #-}

module Servant.API.Stream (
    Stream,
    StreamGet,
    StreamPost,
    StreamBody,
    StreamBody',
    -- * Source
    --
    -- | 'SourceIO' are equivalent to some *source* in streaming libraries.
    SourceIO,
    ToSourceIO (..),
    FromSourceIO (..),
    -- ** Auxiliary classes
    SourceToSourceIO (..),
    -- * Framing
    FramingRender (..),
    FramingUnrender (..),
    -- ** Strategies
    NoFraming,
    NewlineFraming,
    NetstringFraming,
    ) where


import           Control.Applicative
                 ((<|>))
import           Control.Monad.IO.Class
                 (MonadIO (..))
import qualified Data.Attoparsec.ByteString       as A
import qualified Data.Attoparsec.ByteString.Char8 as A8
import qualified Data.ByteString                  as BS
import qualified Data.ByteString.Lazy             as LBS
import qualified Data.ByteString.Lazy.Char8       as LBS8
import           Data.List.NonEmpty
                 (NonEmpty (..))
import           Data.Monoid
                 ((<>))
import           Data.Proxy
                 (Proxy)
import           Data.Typeable
                 (Typeable)
import           GHC.Generics
                 (Generic)
import           GHC.TypeLits
                 (Nat)
import           Network.HTTP.Types.Method
                 (StdMethod (..))
import           Servant.Types.SourceT

-- | A Stream endpoint for a given method emits a stream of encoded values at a
-- given @Content-Type@, delimited by a @framing@ strategy.
-- Type synonyms are provided for standard methods.
--
data Stream (method :: k1) (status :: Nat) (framing :: *) (contentType :: *) (a :: *)
  deriving (Typeable, (forall x.
 Stream method status framing contentType a
 -> Rep (Stream method status framing contentType a) x)
-> (forall x.
    Rep (Stream method status framing contentType a) x
    -> Stream method status framing contentType a)
-> Generic (Stream method status framing contentType a)
forall x.
Rep (Stream method status framing contentType a) x
-> Stream method status framing contentType a
forall x.
Stream method status framing contentType a
-> Rep (Stream method status framing contentType a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall k1 (method :: k1) (status :: Nat) framing contentType a x.
Rep (Stream method status framing contentType a) x
-> Stream method status framing contentType a
forall k1 (method :: k1) (status :: Nat) framing contentType a x.
Stream method status framing contentType a
-> Rep (Stream method status framing contentType a) x
$cto :: forall k1 (method :: k1) (status :: Nat) framing contentType a x.
Rep (Stream method status framing contentType a) x
-> Stream method status framing contentType a
$cfrom :: forall k1 (method :: k1) (status :: Nat) framing contentType a x.
Stream method status framing contentType a
-> Rep (Stream method status framing contentType a) x
Generic)

type StreamGet  = Stream 'GET 200
type StreamPost = Stream 'POST 200

-- | A stream request body.
type StreamBody = StreamBody' '[]

data StreamBody' (mods :: [*]) (framing :: *) (contentType :: *) (a :: *)
  deriving (Typeable, (forall x.
 StreamBody' mods framing contentType a
 -> Rep (StreamBody' mods framing contentType a) x)
-> (forall x.
    Rep (StreamBody' mods framing contentType a) x
    -> StreamBody' mods framing contentType a)
-> Generic (StreamBody' mods framing contentType a)
forall (mods :: [*]) framing contentType a x.
Rep (StreamBody' mods framing contentType a) x
-> StreamBody' mods framing contentType a
forall (mods :: [*]) framing contentType a x.
StreamBody' mods framing contentType a
-> Rep (StreamBody' mods framing contentType a) x
forall x.
Rep (StreamBody' mods framing contentType a) x
-> StreamBody' mods framing contentType a
forall x.
StreamBody' mods framing contentType a
-> Rep (StreamBody' mods framing contentType a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall (mods :: [*]) framing contentType a x.
Rep (StreamBody' mods framing contentType a) x
-> StreamBody' mods framing contentType a
$cfrom :: forall (mods :: [*]) framing contentType a x.
StreamBody' mods framing contentType a
-> Rep (StreamBody' mods framing contentType a) x
Generic)

-------------------------------------------------------------------------------
-- Sources
-------------------------------------------------------------------------------

-- | Stream endpoints may be implemented as producing a @'SourceIO' chunk@.
--
-- Clients reading from streaming endpoints can be implemented as consuming a
-- @'SourceIO' chunk@.
--
type SourceIO = SourceT IO

-- | 'ToSourceIO' is intended to be implemented for types such as Conduit, Pipe,
-- etc. By implementing this class, all such streaming abstractions can be used
-- directly as endpoints.
class ToSourceIO chunk a | a -> chunk where
    toSourceIO :: a -> SourceIO chunk

-- | Auxiliary class for @'ToSourceIO' x ('SourceT' m x)@ instance.
class SourceToSourceIO m where
    sourceToSourceIO :: SourceT m a -> SourceT IO a

instance SourceToSourceIO IO where
    sourceToSourceIO :: SourceT IO a -> SourceT IO a
sourceToSourceIO = SourceT IO a -> SourceT IO a
forall a. a -> a
id

-- | Relax to use auxiliary class, have m
instance SourceToSourceIO m => ToSourceIO chunk (SourceT m chunk) where
    toSourceIO :: SourceT m chunk -> SourceIO chunk
toSourceIO = SourceT m chunk -> SourceIO chunk
forall (m :: * -> *) a.
SourceToSourceIO m =>
SourceT m a -> SourceT IO a
sourceToSourceIO

instance ToSourceIO a (NonEmpty a) where
    toSourceIO :: NonEmpty a -> SourceIO a
toSourceIO (a
x :| [a]
xs) = StepT IO a -> SourceIO a
forall (m :: * -> *) a. StepT m a -> SourceT m a
fromStepT (a -> StepT IO a -> StepT IO a
forall (m :: * -> *) a. a -> StepT m a -> StepT m a
Yield a
x ((a -> StepT IO a -> StepT IO a) -> StepT IO a -> [a] -> StepT IO a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> StepT IO a -> StepT IO a
forall (m :: * -> *) a. a -> StepT m a -> StepT m a
Yield StepT IO a
forall (m :: * -> *) a. StepT m a
Stop [a]
xs))

instance ToSourceIO a [a] where
    toSourceIO :: [a] -> SourceIO a
toSourceIO = [a] -> SourceIO a
forall a (m :: * -> *). [a] -> SourceT m a
source

-- | 'FromSourceIO' is intended to be implemented for types such as Conduit,
-- Pipe, etc. By implementing this class, all such streaming abstractions can
-- be used directly on the client side for talking to streaming endpoints.
class FromSourceIO chunk a | a -> chunk where
    fromSourceIO :: SourceIO chunk -> a

instance MonadIO m => FromSourceIO a (SourceT m a) where
    fromSourceIO :: SourceIO a -> SourceT m a
fromSourceIO = SourceIO a -> SourceT m a
forall (m :: * -> *) a. MonadIO m => SourceIO a -> SourceT m a
sourceFromSourceIO

sourceFromSourceIO :: forall m a. MonadIO m => SourceT IO a -> SourceT m a
sourceFromSourceIO :: SourceT IO a -> SourceT m a
sourceFromSourceIO SourceT IO a
src =
    (forall b. (StepT m a -> m b) -> m b) -> SourceT m a
forall (m :: * -> *) a.
(forall b. (StepT m a -> m b) -> m b) -> SourceT m a
SourceT ((forall b. (StepT m a -> m b) -> m b) -> SourceT m a)
-> (forall b. (StepT m a -> m b) -> m b) -> SourceT m a
forall a b. (a -> b) -> a -> b
$ \StepT m a -> m b
k ->
    StepT m a -> m b
k (StepT m a -> m b) -> StepT m a -> m b
forall a b. (a -> b) -> a -> b
$ m (StepT m a) -> StepT m a
forall (m :: * -> *) a. m (StepT m a) -> StepT m a
Effect (m (StepT m a) -> StepT m a) -> m (StepT m a) -> StepT m a
forall a b. (a -> b) -> a -> b
$ IO (StepT m a) -> m (StepT m a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (StepT m a) -> m (StepT m a))
-> IO (StepT m a) -> m (StepT m a)
forall a b. (a -> b) -> a -> b
$ SourceT IO a -> (StepT IO a -> IO (StepT m a)) -> IO (StepT m a)
forall (m :: * -> *) a.
SourceT m a -> forall b. (StepT m a -> m b) -> m b
unSourceT SourceT IO a
src (StepT m a -> IO (StepT m a)
forall (m :: * -> *) a. Monad m => a -> m a
return (StepT m a -> IO (StepT m a))
-> (StepT IO a -> StepT m a) -> StepT IO a -> IO (StepT m a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StepT IO a -> StepT m a
go)
  where
    go :: StepT IO a -> StepT m a
    go :: StepT IO a -> StepT m a
go StepT IO a
Stop        = StepT m a
forall (m :: * -> *) a. StepT m a
Stop
    go (Error String
err) = String -> StepT m a
forall (m :: * -> *) a. String -> StepT m a
Error String
err
    go (Skip StepT IO a
s)    = StepT m a -> StepT m a
forall (m :: * -> *) a. StepT m a -> StepT m a
Skip (StepT IO a -> StepT m a
go StepT IO a
s)
    go (Effect IO (StepT IO a)
ms) = m (StepT m a) -> StepT m a
forall (m :: * -> *) a. m (StepT m a) -> StepT m a
Effect (IO (StepT m a) -> m (StepT m a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO ((StepT IO a -> StepT m a) -> IO (StepT IO a) -> IO (StepT m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap StepT IO a -> StepT m a
go IO (StepT IO a)
ms))
    go (Yield a
x StepT IO a
s) = a -> StepT m a -> StepT m a
forall (m :: * -> *) a. a -> StepT m a -> StepT m a
Yield a
x (StepT IO a -> StepT m a
go StepT IO a
s)

-- This fires e.g. in Client.lhs
-- {-# OPTIONS_GHC -ddump-simpl -ddump-rule-firings -ddump-to-file #-}
{-# NOINLINE [2] sourceFromSourceIO #-}
{-# RULES "sourceFromSourceIO @IO" sourceFromSourceIO = id :: SourceT IO a -> SourceT IO a #-}

-------------------------------------------------------------------------------
-- Framing
-------------------------------------------------------------------------------

-- | The 'FramingRender' class provides the logic for emitting a framing strategy.
-- The strategy transforms a @'SourceT' m a@ into @'SourceT' m 'LBS.ByteString'@,
-- therefore it can prepend, append and intercalate /framing/ structure
-- around chunks.
--
-- /Note:/ as the @'Monad' m@ is generic, this is pure transformation.
--
class FramingRender strategy where
    framingRender :: Monad m => Proxy strategy -> (a -> LBS.ByteString) -> SourceT m a -> SourceT m LBS.ByteString

-- | The 'FramingUnrender' class provides the logic for parsing a framing
-- strategy.
class FramingUnrender strategy where
    framingUnrender :: Monad m => Proxy strategy -> (LBS.ByteString -> Either String a) -> SourceT m BS.ByteString -> SourceT m a

-------------------------------------------------------------------------------
-- NoFraming
-------------------------------------------------------------------------------

-- | A framing strategy that does not do any framing at all, it just passes the
-- input data This will be used most of the time with binary data, such as
-- files
data NoFraming

instance FramingRender NoFraming where
    framingRender :: Proxy NoFraming
-> (a -> ByteString) -> SourceT m a -> SourceT m ByteString
framingRender Proxy NoFraming
_ = (a -> ByteString) -> SourceT m a -> SourceT m ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap

-- | As 'NoFraming' doesn't have frame separators, we take the chunks
-- as given and try to convert them one by one.
--
-- That works well when @a@ is a 'ByteString'.
instance FramingUnrender NoFraming where
    framingUnrender :: Proxy NoFraming
-> (ByteString -> Either String a)
-> SourceT m ByteString
-> SourceT m a
framingUnrender Proxy NoFraming
_ ByteString -> Either String a
f = (StepT m ByteString -> StepT m a)
-> SourceT m ByteString -> SourceT m a
forall (m :: * -> *) a b.
(StepT m a -> StepT m b) -> SourceT m a -> SourceT m b
mapStepT StepT m ByteString -> StepT m a
forall (m :: * -> *). Functor m => StepT m ByteString -> StepT m a
go
      where
        go :: StepT m ByteString -> StepT m a
go StepT m ByteString
Stop        = StepT m a
forall (m :: * -> *) a. StepT m a
Stop
        go (Error String
err) = String -> StepT m a
forall (m :: * -> *) a. String -> StepT m a
Error String
err
        go (Skip StepT m ByteString
s)    = StepT m a -> StepT m a
forall (m :: * -> *) a. StepT m a -> StepT m a
Skip (StepT m ByteString -> StepT m a
go StepT m ByteString
s)
        go (Effect m (StepT m ByteString)
ms) = m (StepT m a) -> StepT m a
forall (m :: * -> *) a. m (StepT m a) -> StepT m a
Effect ((StepT m ByteString -> StepT m a)
-> m (StepT m ByteString) -> m (StepT m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap StepT m ByteString -> StepT m a
go m (StepT m ByteString)
ms)
        go (Yield ByteString
x StepT m ByteString
s) = case ByteString -> Either String a
f (ByteString -> ByteString
LBS.fromStrict ByteString
x) of
            Right a
y  -> a -> StepT m a -> StepT m a
forall (m :: * -> *) a. a -> StepT m a -> StepT m a
Yield a
y (StepT m ByteString -> StepT m a
go StepT m ByteString
s)
            Left String
err -> String -> StepT m a
forall (m :: * -> *) a. String -> StepT m a
Error String
err

-------------------------------------------------------------------------------
-- NewlineFraming
-------------------------------------------------------------------------------

-- | A simple framing strategy that has no header, and inserts a
-- newline character after each frame.  This assumes that it is used with a
-- Content-Type that encodes without newlines (e.g. JSON).
data NewlineFraming

instance FramingRender NewlineFraming where
    framingRender :: Proxy NewlineFraming
-> (a -> ByteString) -> SourceT m a -> SourceT m ByteString
framingRender Proxy NewlineFraming
_ a -> ByteString
f = (a -> ByteString) -> SourceT m a -> SourceT m ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
x -> a -> ByteString
f a
x ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"\n")

instance FramingUnrender NewlineFraming where
    framingUnrender :: Proxy NewlineFraming
-> (ByteString -> Either String a)
-> SourceT m ByteString
-> SourceT m a
framingUnrender Proxy NewlineFraming
_ ByteString -> Either String a
f = Parser a -> SourceT m ByteString -> SourceT m a
forall (m :: * -> *) a.
Monad m =>
Parser a -> SourceT m ByteString -> SourceT m a
transformWithAtto (Parser a -> SourceT m ByteString -> SourceT m a)
-> Parser a -> SourceT m ByteString -> SourceT m a
forall a b. (a -> b) -> a -> b
$ do
        ByteString
bs <- (Word8 -> Bool) -> Parser ByteString
A.takeWhile (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
10)
        () () -> Parser ByteString Word8 -> Parser ByteString ()
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Word8 -> Parser ByteString Word8
A.word8 Word8
10 Parser ByteString ()
-> Parser ByteString () -> Parser ByteString ()
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Parser ByteString ()
forall t. Chunk t => Parser t ()
A.endOfInput
        (String -> Parser a)
-> (a -> Parser a) -> Either String a -> Parser a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either String -> Parser a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Either String a
f (ByteString -> ByteString
LBS.fromStrict ByteString
bs))

-------------------------------------------------------------------------------
-- NetstringFraming
-------------------------------------------------------------------------------

-- | The netstring framing strategy as defined by djb:
-- <http://cr.yp.to/proto/netstrings.txt>
--
-- Any string of 8-bit bytes may be encoded as @[len]":"[string]","@.  Here
-- @[string]@ is the string and @[len]@ is a nonempty sequence of ASCII digits
-- giving the length of @[string]@ in decimal. The ASCII digits are @<30>@ for
-- 0, @<31>@ for 1, and so on up through @<39>@ for 9. Extra zeros at the front
-- of @[len]@ are prohibited: @[len]@ begins with @<30>@ exactly when
-- @[string]@ is empty.
--
-- For example, the string @"hello world!"@ is encoded as
-- @<31 32 3a 68 65 6c 6c 6f 20 77 6f 72 6c 64 21 2c>@,
-- i.e., @"12:hello world!,"@.
-- The empty string is encoded as @"0:,"@.
--
data NetstringFraming

instance FramingRender NetstringFraming where
    framingRender :: Proxy NetstringFraming
-> (a -> ByteString) -> SourceT m a -> SourceT m ByteString
framingRender Proxy NetstringFraming
_ a -> ByteString
f = (a -> ByteString) -> SourceT m a -> SourceT m ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> ByteString) -> SourceT m a -> SourceT m ByteString)
-> (a -> ByteString) -> SourceT m a -> SourceT m ByteString
forall a b. (a -> b) -> a -> b
$ \a
x ->
        let bs :: ByteString
bs = a -> ByteString
f a
x
        in String -> ByteString
LBS8.pack (Int64 -> String
forall a. Show a => a -> String
show (ByteString -> Int64
LBS8.length ByteString
bs)) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
":" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
bs ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
","

instance FramingUnrender NetstringFraming where
    framingUnrender :: Proxy NetstringFraming
-> (ByteString -> Either String a)
-> SourceT m ByteString
-> SourceT m a
framingUnrender Proxy NetstringFraming
_ ByteString -> Either String a
f = Parser a -> SourceT m ByteString -> SourceT m a
forall (m :: * -> *) a.
Monad m =>
Parser a -> SourceT m ByteString -> SourceT m a
transformWithAtto (Parser a -> SourceT m ByteString -> SourceT m a)
-> Parser a -> SourceT m ByteString -> SourceT m a
forall a b. (a -> b) -> a -> b
$ do
        Int
len <- Parser Int
forall a. Integral a => Parser a
A8.decimal
        Char
_ <- Char -> Parser Char
A8.char Char
':'
        ByteString
bs <- Int -> Parser ByteString
A.take Int
len
        Char
_ <- Char -> Parser Char
A8.char Char
','
        (String -> Parser a)
-> (a -> Parser a) -> Either String a -> Parser a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either String -> Parser a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Either String a
f (ByteString -> ByteString
LBS.fromStrict ByteString
bs))