@@ -21,8 +21,8 @@ module Streamly.Internal.Network.Socket
2121 (
2222 SockSpec (..)
2323 -- * Use a socket
24− , useSocketM
25− , useSocket
24+ , handleWithM
25+ , handleWith
2626
2727 -- * Accept connections
2828 , accept
@@ -60,7 +60,7 @@ module Streamly.Internal.Network.Socket
6060 , fromBytes
6161
6262 -- -- * Array Write
63− , writeArray
63+ , writeChunk
6464 , writeChunks
6565 , writeStrings
6666
@@ -101,7 +101,6 @@ import Streamly.Data.Fold (Fold)
101101-- import Streamly.String (encodeUtf8, decodeUtf8, foldLines)
102102
103103import qualified Streamly.Data.Fold as FL
104−import qualified Streamly.Data.Unicode.Stream as U
105104import qualified Streamly.Internal.Data.Fold.Types as FL
106105import qualified Streamly.Internal.Data.Unfold as UF
107106import qualified Streamly.Internal.Memory.Array as IA
@@ -111,25 +110,25 @@ import qualified Streamly.Internal.Memory.Array.Types as A
111110import qualified Streamly.Prelude as S
112111import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
113112
114−-- | @'useSocketM' socket act@ runs the monadic computation @act@ passing the
115−-- socket handle to it. The handle will be closed on exit from 'useSocketM',
113+-- | @'handleWithM' socket act@ runs the monadic computation @act@ passing the
114+-- socket handle to it. The handle will be closed on exit from 'handleWithM',
116115-- whether by normal termination or by raising an exception. If closing the
117116-- handle raises an exception, then this exception will be raised by
118−-- 'useSocketM' rather than any exception raised by 'act'.
117+-- 'handleWithM' rather than any exception raised by 'act'.
119118--
120119-- @since 0.7.0
121−{-# INLINE useSocketM #-}
122−useSocketM :: (MonadMask m, MonadIO m) => Socket -> (Socket -> m ()) -> m ()
123−useSocketM sk f = finally (f sk) (liftIO (Net.close sk))
120+{-# INLINE handleWithM #-}
121+handleWithM :: (MonadMask m, MonadIO m) => (Socket -> m ()) -> Socket -> m ()
122+handleWithM f sk = finally (f sk) (liftIO (Net.close sk))
124123
125−-- | Like 'useSocketM' but runs a streaming computation instead of a monadic
124+-- | Like 'handleWithM' but runs a streaming computation instead of a monadic
126125-- computation.
127126--
128127-- @since 0.7.0
129−{-# INLINE useSocket #-}
130−useSocket :: (IsStream t, MonadCatch m, MonadIO m)
128+{-# INLINE handleWith #-}
129+handleWith :: (IsStream t, MonadCatch m, MonadIO m)
131130 => Socket -> (Socket -> t m a) -> t m a
132−useSocket sk f = S.finally (liftIO $ Net.close sk) (f sk)
131+handleWith sk f = S.finally (liftIO $ Net.close sk) (f sk)
133132
134133-------------------------------------------------------------------------------
135134-- Accept (Unfolds)
@@ -282,19 +281,19 @@ writeArrayWith f h Array{..} = withForeignPtr aStart $ \p ->
282281-- | Write an Array to a file handle.
283282--
284283-- @since 0.7.0
285−{-# INLINABLE writeArray #-}
286−writeArray :: Storable a => Socket -> Array a -> IO ()
287−writeArray = writeArrayWith sendAll
284+{-# INLINABLE writeChunk #-}
285+writeChunk :: Storable a => Socket -> Array a -> IO ()
286+writeChunk = writeArrayWith sendAll
288287
289288-------------------------------------------------------------------------------
290289-- Stream of Arrays IO
291290-------------------------------------------------------------------------------
292291
293−{-# INLINABLE readChunksUptoWith #-}
294−readChunksUptoWith :: (IsStream t, MonadIO m)
292+{-# INLINABLE _readChunksUptoWith #-}
293+_readChunksUptoWith :: (IsStream t, MonadIO m)
295294 => (Int -> h -> IO (Array Word8))
296295 -> Int -> h -> t m (Array Word8)
297−readChunksUptoWith f size h = go
296+_readChunksUptoWith f size h = go
298297 where
299298 -- XXX use cons/nil instead
300299 go = mkStream $ \_ yld _ stp -> do
@@ -307,10 +306,19 @@ readChunksUptoWith f size h = go
307306-- The maximum size of a single array is limited to @size@.
308307-- 'fromHandleArraysUpto' ignores the prevailing 'TextEncoding' and 'NewlineMode'
309308-- on the 'Handle'.
310−{-# INLINABLE toChunksWithBufferOf #-}
309+{-# INLINE_NORMAL toChunksWithBufferOf #-}
311310toChunksWithBufferOf :: (IsStream t, MonadIO m)
312311 => Int -> Socket -> t m (Array Word8)
313−toChunksWithBufferOf = readChunksUptoWith readArrayOf
312+-- toChunksWithBufferOf = _readChunksUptoWith readArrayOf
313+toChunksWithBufferOf size h = D.fromStreamD (D.Stream step ())
314+ where
315+ {-# INLINE_LATE step #-}
316+ step _ _ = do
317+ arr <- liftIO $ readArrayOf size h
318+ return $
319+ case A.length arr of
320+ 0 -> D.Stop
321+ _ -> D.Yield arr ()
314322
315323-- XXX read 'Array a' instead of Word8
316324--
@@ -408,7 +416,7 @@ read = UF.supplyFirst readWithBufferOf A.defaultChunkSize
408416{-# INLINE fromChunks #-}
409417fromChunks :: (MonadIO m, Storable a)
410418 => Socket -> SerialT m (Array a) -> m ()
411−fromChunks h m = S.mapM_ (liftIO . writeArray h) m
419+fromChunks h m = S.mapM_ (liftIO . writeChunk h) m
412420
413421-- | Write a stream of arrays to a socket. Each array in the stream is written
414422-- to the socket as a separate IO request.
@@ -416,16 +424,18 @@ fromChunks h m = S.mapM_ (liftIO . writeArray h) m
416424-- @since 0.7.0
417425{-# INLINE writeChunks #-}
418426writeChunks :: (MonadIO m, Storable a) => Socket -> Fold m (Array a) ()
419−writeChunks h = FL.drainBy (liftIO . writeArray h)
427+writeChunks h = FL.drainBy (liftIO . writeChunk h)
420428
421−-- | Write a stream of strings to a socket in Latin1 encoding.
429+-- | Write a stream of strings to a socket in Latin1 encoding. Output is
430+-- flushed to the socket for each string.
422431--
423432-- /Internal/
424433--
425434{-# INLINE writeStrings #-}
426−writeStrings :: MonadIO m => Socket -> Fold m String ()
427−writeStrings h =
428− FL.lmapM (IA.fromStream . U.encodeLatin1 . S.fromList) (writeChunks h)
435+writeStrings :: MonadIO m
436+ => (SerialT m Char -> SerialT m Word8) -> Socket -> Fold m String ()
437+writeStrings encode h =
438+ FL.lmapM (IA.fromStream . encode . S.fromList) (writeChunks h)
429439
430440-- GHC buffer size dEFAULT_FD_BUFFER_SIZE=8192 bytes.
431441--