@@ -19,8 +19,8 @@
1919module Streamly.Internal.Network.Socket
2020 (
2121 -- * Use a socket
22− withSocket
23− , withSocketS
22+ withSocketM
23+ , withSocket
2424
2525 -- * Read from connection
2626 , read
@@ -28,15 +28,19 @@ module Streamly.Internal.Network.Socket
2828 -- , readLines
2929 -- , readFrames
3030 -- , readByChunks
31+ , toStream
3132
3233 -- -- * Array Read
3334 -- , readArrayUpto
3435 -- , readArrayOf
3536
3637 -- , readArraysUpto
37− -- , readArraysOf
38+ , readArraysOf
3839 , readArrays
3940
41+ , toStreamArraysOf
42+ , toStreamArrays
43+
4044 -- * Write to connection
4145 , write
4246 -- , writeUtf8
@@ -72,6 +76,7 @@ import Prelude hiding (read)
7276
7377import qualified Network.Socket as Net
7478
79+import Streamly.Internal.Data.Unfold.Types (Unfold(..))
7580import Streamly.Internal.Memory.Array.Types (Array(..))
7681import Streamly.Streams.Serial (SerialT)
7782import Streamly.Streams.StreamK.Type (IsStream, mkStream)
@@ -80,21 +85,23 @@ import Streamly.Data.Fold (Fold)
8085
8186import qualified Streamly.Data.Fold as FL
8287import qualified Streamly.Internal.Data.Fold.Types as FL
88+import qualified Streamly.Internal.Data.Unfold as UF
8389import qualified Streamly.Memory.Array as A
8490import qualified Streamly.Internal.Memory.ArrayStream as AS
8591import qualified Streamly.Internal.Memory.Array.Types as A
8692import qualified Streamly.Prelude as S
93+import qualified Streamly.Streams.StreamD.Type as D
8794
88−-- | @'withSocket' socket act@ runs the monadic computation @act@ passing the
89−-- socket handle to it. The handle will be closed on exit from 'withSocket',
95+-- | @'withSocketM' socket act@ runs the monadic computation @act@ passing the
96+-- socket handle to it. The handle will be closed on exit from 'withSocketM',
9097-- whether by normal termination or by raising an exception. If closing the
9198-- handle raises an exception, then this exception will be raised by
92−-- 'withSocket' rather than any exception raised by 'act'.
99+-- 'withSocketM' rather than any exception raised by 'act'.
93100--
94101-- @since 0.7.0
95−{-# INLINE withSocket #-}
96−withSocket :: (MonadCatch m, MonadIO m) => Socket -> (Socket -> m ()) -> m ()
97−withSocket sk f = do
102+{-# INLINE withSocketM #-}
103+withSocketM :: (MonadCatch m, MonadIO m) => Socket -> (Socket -> m ()) -> m ()
104+withSocketM sk f = do
98105 f sk `onException` liftIO (Net.close sk)
99106 liftIO (Net.close sk)
100107
@@ -105,14 +112,14 @@ withSocket sk f = do
105112-- withSocketS sk = A.flattenArrays $
106113-- S.finally (liftIO (Net.close sk)) (readArrays sk)
107114
108−-- | Like 'withSocket' but runs a streaming computation instead of a monadic
115+-- | Like 'withSocketM' but runs a streaming computation instead of a monadic
109116-- computation.
110117--
111118-- @since 0.7.0
112−{-# INLINE withSocketS #-}
113−withSocketS :: (IsStream t, MonadCatch m, MonadIO m)
119+{-# INLINE withSocket #-}
120+withSocket :: (IsStream t, MonadCatch m, MonadIO m)
114121 => Socket -> (Socket -> t m a) -> t m a
115−withSocketS sk = S.bracket (return sk) (liftIO . Net.close)
122+withSocket sk = S.bracket (return sk) (liftIO . Net.close)
116123
117124-------------------------------------------------------------------------------
118125-- Array IO (Input)
@@ -205,24 +212,49 @@ readArraysUptoWith f size h = go
205212 then stp
206213 else yld arr go
207214
208−-- | @readArraysOf size h@ reads a stream of arrays from file handle @h@.
215+-- | @toStreamArraysOf size h@ reads a stream of arrays from file handle @h@.
209216-- The maximum size of a single array is limited to @size@.
210217-- 'fromHandleArraysUpto' ignores the prevailing 'TextEncoding' and 'NewlineMode'
211218-- on the 'Handle'.
212−{-# INLINABLE readArraysOf #-}
213−readArraysOf :: (IsStream t, MonadIO m)
219+{-# INLINABLE toStreamArraysOf #-}
220+toStreamArraysOf :: (IsStream t, MonadIO m)
214221 => Int -> Socket -> t m (Array Word8)
215−readArraysOf = readArraysUptoWith readArrayOf
222+toStreamArraysOf = readArraysUptoWith readArrayOf
216223
217224-- XXX read 'Array a' instead of Word8
218225--
219−-- | @readArrays h@ reads a stream of arrays from socket handle @h@.
226+-- | @toStreamArrays h@ reads a stream of arrays from socket handle @h@.
220227-- The maximum size of a single array is limited to @defaultChunkSize@.
221228--
222229-- @since 0.7.0
230+{-# INLINE toStreamArrays #-}
231+toStreamArrays :: (IsStream t, MonadIO m) => Socket -> t m (Array Word8)
232+toStreamArrays = toStreamArraysOf A.defaultChunkSize
233+
234+-- | Unfold the tuple @(size, socket)@ into a stream of 'Word8' arrays. The
235+-- stream consists of arrays representing chunks of data read from the socket.
236+-- The maximum size of a single array is limited to @size@.
237+--
238+-- @since 0.7.0
239+{-# INLINE_NORMAL readArraysOf #-}
240+readArraysOf :: MonadIO m => Unfold m (Int, Socket) (Array Word8)
241+readArraysOf = Unfold step return
242+ where
243+ {-# INLINE_LATE step #-}
244+ step (size, h) = do
245+ arr <- liftIO $ readArrayOf size h
246+ return $
247+ case A.length arr of
248+ 0 -> D.Stop
249+ _ -> D.Yield arr (size, h)
250+
251+-- | Unfolds a socket into a stream of 'Word8' arrays. The maximum size of a
252+-- single array is limited to @defaultChunkSize@.
253+--
254+-- @since 0.7.0
223255{-# INLINE readArrays #-}
224−readArrays :: (IsStream t, MonadIO m) => Socket -> t m (Array Word8)
225−readArrays = readArraysOf A.defaultChunkSize
256+readArrays :: MonadIO m => Unfold m Socket (Array Word8)
257+readArrays = UF.first readArraysOf A.defaultChunkSize
226258
227259-------------------------------------------------------------------------------
228260-- Read File to Stream
@@ -250,9 +282,24 @@ readInChunksOf chunkSize h = A.flattenArrays $ readArraysUpto chunkSize h
250282-- stream ends when EOF is encountered.
251283--
252284-- @since 0.7.0
285+{-# INLINE toStream #-}
286+toStream :: (IsStream t, MonadIO m) => Socket -> t m Word8
287+toStream = AS.concat . toStreamArrays
288+
289+-- | @readInChunksOf@ unfolds @(bufsize, socket)@ into a byte stream, reads
290+-- are performed in buffers of up to @bufsize@.
291+--
292+-- @since 0.7.0
293+{-# INLINE readInChunksOf #-}
294+readInChunksOf :: MonadIO m => Unfold m (Int, Socket) Word8
295+readInChunksOf = UF.concat readArraysOf A.read
296+
297+-- | Unfolds a 'Socket' into a byte stream.
298+--
299+-- @since 0.7.0
253300{-# INLINE read #-}
254−read :: (IsStream t, MonadIO m) => Socket -> t m Word8
255−read = AS.concat . readArrays
301+read :: MonadIO m => Unfold m Socket Word8
302+read = UF.first readInChunksOf A.defaultChunkSize
256303
257304-------------------------------------------------------------------------------
258305-- Writing