{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_HADDOCK prune not-home #-}
module Data.Attoparsec.Framer.Testing (
parsesFromFramerOk,
chunksOfN,
linkedSrcAndSink,
linkedSrcAndSink',
) where
import Control.Exception (catch)
import Control.Monad (when)
import qualified Data.Attoparsec.ByteString as A
import Data.Attoparsec.Framer
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.ByteString.Builder (byteStringHex, toLazyByteString)
import qualified Data.ByteString.Lazy.Char8 as C8
import Data.IORef (
IORef,
modifyIORef',
newIORef,
readIORef,
writeIORef,
)
import Data.List (unfoldr)
import Data.Word (Word32)
parsesFromFramerOk :: Eq a => (a -> ByteString) -> A.Parser a -> Word32 -> [a] -> IO Bool
parsesFromFramerOk :: forall a.
Eq a =>
(a -> ByteString) -> Parser a -> Word32 -> [a] -> IO Bool
parsesFromFramerOk a -> ByteString
asBytes Parser a
parser Word32
chunkSize' [a]
wanted = do
IORef (Maybe [ByteString])
chunkStore <- Maybe [ByteString] -> IO (IORef (Maybe [ByteString]))
forall a. a -> IO (IORef a)
newIORef Maybe [ByteString]
forall a. Maybe a
Nothing
IORef [a]
dst <- [a] -> IO (IORef [a])
forall a. a -> IO (IORef a)
newIORef []
let updateDst :: a -> IO ()
updateDst a
x = IORef [a] -> ([a] -> [a]) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef [a]
dst ((:) a
x)
mkChunks :: Int -> [ByteString]
mkChunks Int
n = [[ByteString]] -> [ByteString]
forall a. Monoid a => [a] -> a
mconcat ([[ByteString]] -> [ByteString]) -> [[ByteString]] -> [ByteString]
forall a b. (a -> b) -> a -> b
$ (a -> [ByteString]) -> [a] -> [[ByteString]]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> ByteString -> [ByteString]
chunksOfN Int
n (ByteString -> [ByteString])
-> (a -> ByteString) -> a -> [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ByteString
asBytes) [a]
wanted
src :: Word32 -> IO ByteString
src = (Int -> [ByteString])
-> IORef (Maybe [ByteString]) -> Word32 -> IO ByteString
nextFrom' Int -> [ByteString]
mkChunks IORef (Maybe [ByteString])
chunkStore
frames :: Framer IO a
frames = Word32 -> Framer IO a -> Framer IO a
forall (m :: * -> *) a. Word32 -> Framer m a -> Framer m a
setChunkSize Word32
chunkSize' (Framer IO a -> Framer IO a) -> Framer IO a -> Framer IO a
forall a b. (a -> b) -> a -> b
$ Parser a
-> (a -> IO ()) -> (Word32 -> IO ByteString) -> Framer IO a
forall (m :: * -> *) frame.
MonadThrow m =>
Parser frame -> (frame -> m ()) -> ByteSource m -> Framer m frame
mkFramer Parser a
parser a -> IO ()
updateDst Word32 -> IO ByteString
src
Framer IO a -> IO ()
forall (m :: * -> *) frame. MonadThrow m => Framer m frame -> m ()
runFramer Framer IO a
frames IO () -> (NoMoreInput -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` (\(NoMoreInput
_e :: NoMoreInput) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
[a]
got <- IORef [a] -> IO [a]
forall a. IORef a -> IO a
readIORef IORef [a]
dst
Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ [a]
got [a] -> [a] -> Bool
forall a. Eq a => a -> a -> Bool
== [a] -> [a]
forall a. [a] -> [a]
reverse [a]
wanted
chunksOfN :: Int -> ByteString -> [ByteString]
chunksOfN :: Int -> ByteString -> [ByteString]
chunksOfN Int
x ByteString
b =
let go :: ByteString -> Maybe (ByteString, ByteString)
go ByteString
y =
let taken :: ByteString
taken = Int -> ByteString -> ByteString
BS.take Int
x ByteString
y
in if ByteString -> Bool
BS.null ByteString
taken then Maybe (ByteString, ByteString)
forall a. Maybe a
Nothing else (ByteString, ByteString) -> Maybe (ByteString, ByteString)
forall a. a -> Maybe a
Just (ByteString
taken, Int -> ByteString -> ByteString
BS.drop Int
x ByteString
y)
in (ByteString -> Maybe (ByteString, ByteString))
-> ByteString -> [ByteString]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr ByteString -> Maybe (ByteString, ByteString)
go ByteString
b
nextFrom' ::
(Int -> [ByteString]) -> IORef (Maybe [ByteString]) -> Word32 -> IO ByteString
nextFrom' :: (Int -> [ByteString])
-> IORef (Maybe [ByteString]) -> Word32 -> IO ByteString
nextFrom' Int -> [ByteString]
initChunks IORef (Maybe [ByteString])
chunkStore Word32
chunkSize' = do
IORef (Maybe [ByteString]) -> IO (Maybe [ByteString])
forall a. IORef a -> IO a
readIORef IORef (Maybe [ByteString])
chunkStore IO (Maybe [ByteString])
-> (Maybe [ByteString] -> IO ByteString) -> IO ByteString
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
Maybe [ByteString]
Nothing -> do
IORef (Maybe [ByteString]) -> Maybe [ByteString] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe [ByteString])
chunkStore (Maybe [ByteString] -> IO ()) -> Maybe [ByteString] -> IO ()
forall a b. (a -> b) -> a -> b
$ [ByteString] -> Maybe [ByteString]
forall a. a -> Maybe a
Just ([ByteString] -> Maybe [ByteString])
-> [ByteString] -> Maybe [ByteString]
forall a b. (a -> b) -> a -> b
$ Int -> [ByteString]
initChunks (Int -> [ByteString]) -> Int -> [ByteString]
forall a b. (a -> b) -> a -> b
$ Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
chunkSize'
(Int -> [ByteString])
-> IORef (Maybe [ByteString]) -> Word32 -> IO ByteString
nextFrom' Int -> [ByteString]
initChunks IORef (Maybe [ByteString])
chunkStore Word32
chunkSize'
Just [] -> ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
BS.empty
Just (ByteString
x : [ByteString]
xs) -> do
IORef (Maybe [ByteString]) -> Maybe [ByteString] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe [ByteString])
chunkStore (Maybe [ByteString] -> IO ()) -> Maybe [ByteString] -> IO ()
forall a b. (a -> b) -> a -> b
$ [ByteString] -> Maybe [ByteString]
forall a. a -> Maybe a
Just [ByteString]
xs
ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
x
linkedSrcAndSink :: [ByteString] -> IO (ByteSource IO, (ByteString -> IO ()))
linkedSrcAndSink :: [ByteString] -> IO (Word32 -> IO ByteString, ByteString -> IO ())
linkedSrcAndSink [ByteString]
responses = do
IORef (Maybe ByteString)
refSrc <- Maybe ByteString -> IO (IORef (Maybe ByteString))
forall a. a -> IO (IORef a)
newIORef Maybe ByteString
forall a. Maybe a
Nothing
IORef [ByteString]
refSink <- [ByteString] -> IO (IORef [ByteString])
forall a. a -> IO (IORef a)
newIORef [ByteString]
responses
(Word32 -> IO ByteString, ByteString -> IO ())
-> IO (Word32 -> IO ByteString, ByteString -> IO ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IORef (Maybe ByteString) -> Word32 -> IO ByteString
ioRefByteSource IORef (Maybe ByteString)
refSrc, Bool
-> IORef [ByteString]
-> IORef (Maybe ByteString)
-> ByteString
-> IO ()
ioRefByteSink Bool
False IORef [ByteString]
refSink IORef (Maybe ByteString)
refSrc)
linkedSrcAndSink' :: [ByteString] -> IO (ByteSource IO, (ByteString -> IO ()))
linkedSrcAndSink' :: [ByteString] -> IO (Word32 -> IO ByteString, ByteString -> IO ())
linkedSrcAndSink' [ByteString]
responses = do
IORef (Maybe ByteString)
refSrc <- Maybe ByteString -> IO (IORef (Maybe ByteString))
forall a. a -> IO (IORef a)
newIORef Maybe ByteString
forall a. Maybe a
Nothing
IORef [ByteString]
refSink <- [ByteString] -> IO (IORef [ByteString])
forall a. a -> IO (IORef a)
newIORef [ByteString]
responses
(Word32 -> IO ByteString, ByteString -> IO ())
-> IO (Word32 -> IO ByteString, ByteString -> IO ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IORef (Maybe ByteString) -> Word32 -> IO ByteString
ioRefByteSource IORef (Maybe ByteString)
refSrc, Bool
-> IORef [ByteString]
-> IORef (Maybe ByteString)
-> ByteString
-> IO ()
ioRefByteSink Bool
True IORef [ByteString]
refSink IORef (Maybe ByteString)
refSrc)
ioRefByteSource :: IORef (Maybe ByteString) -> ByteSource IO
ioRefByteSource :: IORef (Maybe ByteString) -> Word32 -> IO ByteString
ioRefByteSource IORef (Maybe ByteString)
refSrc Word32
size = do
IORef (Maybe ByteString) -> IO (Maybe ByteString)
forall a. IORef a -> IO a
readIORef IORef (Maybe ByteString)
refSrc IO (Maybe ByteString)
-> (Maybe ByteString -> IO ByteString) -> IO ByteString
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
Maybe ByteString
Nothing -> ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
BS.empty
Just ByteString
src -> do
let taken :: ByteString
taken = Int -> ByteString -> ByteString
BS.take (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
size) ByteString
src
rest :: ByteString
rest = Int -> ByteString -> ByteString
BS.drop (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
size) ByteString
src
stored :: Maybe ByteString
stored = if ByteString -> Bool
BS.null ByteString
taken then Maybe ByteString
forall a. Maybe a
Nothing else ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
rest
IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe ByteString)
refSrc Maybe ByteString
stored
ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
taken
ioRefByteSink :: Bool -> IORef [ByteString] -> IORef (Maybe ByteString) -> ByteString -> IO ()
ioRefByteSink :: Bool
-> IORef [ByteString]
-> IORef (Maybe ByteString)
-> ByteString
-> IO ()
ioRefByteSink Bool
debug IORef [ByteString]
refResponses IORef (Maybe ByteString)
refSrc ByteString
_ignored = do
let asHex :: ByteString -> ByteString
asHex = Builder -> ByteString
toLazyByteString (Builder -> ByteString)
-> (ByteString -> Builder) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Builder
byteStringHex
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> IO ()
C8.putStrLn (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString
"bytesink got: " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> (ByteString -> ByteString
asHex ByteString
_ignored)
IORef [ByteString] -> IO [ByteString]
forall a. IORef a -> IO a
readIORef IORef [ByteString]
refResponses IO [ByteString] -> ([ByteString] -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
[] -> do
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> IO ()
C8.putStrLn ByteString
"bytesource has nothing"
IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe ByteString)
refSrc Maybe ByteString
forall a. Maybe a
Nothing
(ByteString
x : [ByteString]
xs) -> do
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> IO ()
C8.putStrLn (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString
"bytesink will reply with: " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> (ByteString -> ByteString
asHex ByteString
x)
IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe ByteString)
refSrc (Maybe ByteString -> IO ()) -> Maybe ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
x
IORef [ByteString] -> [ByteString] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef [ByteString]
refResponses ([ByteString] -> IO ()) -> [ByteString] -> IO ()
forall a b. (a -> b) -> a -> b
$ [ByteString]
xs