{-# 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.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))