1{- Original: Network/Socket/ByteString/* -}
2
3-- | Support module for the POSIX writev system call.
4module Network.Sendfile.IOVec (
5    IOVec(..)
6  , SfHdtr(..)
7  , withSfHdtr
8  , remainingChunks
9  ) where
10
11import Control.Monad (zipWithM_)
12import Data.ByteString (ByteString)
13import qualified Data.ByteString as BS
14import Data.ByteString.Unsafe (unsafeUseAsCStringLen)
15import Foreign.C.Types (CChar, CInt, CSize)
16import Foreign.Marshal (alloca)
17import Foreign.Marshal.Array (allocaArray)
18import Foreign.Ptr (Ptr, nullPtr, plusPtr)
19import Foreign.Storable (Storable(..))
20
21#include <sys/uio.h>
22#include <sys/socket.h>
23
24----------------------------------------------------------------
25
26data IOVec = IOVec {
27    iovBase :: Ptr CChar
28  , iovLen  :: CSize
29  }
30
31instance Storable IOVec where
32  sizeOf _    = (#const sizeof(struct iovec))
33  alignment _ = alignment (undefined :: CInt)
34
35  peek p = do
36      base <- (#peek struct iovec, iov_base) p
37      len  <- (#peek struct iovec, iov_len)  p
38      return $ IOVec base len
39
40  poke p iov = do
41      (#poke struct iovec, iov_base) p (iovBase iov)
42      (#poke struct iovec, iov_len)  p (iovLen  iov)
43
44----------------------------------------------------------------
45
46data SfHdtr = SfHdtr {
47    sfhdtrHdr    :: Ptr IOVec
48  , sfhdtrHdrLen :: CInt
49  }
50
51instance Storable SfHdtr where
52  sizeOf _    = (#const sizeof(struct sf_hdtr))
53  alignment _ = alignment (undefined :: CInt)
54
55  peek p = do
56      hdr  <- (#peek struct sf_hdtr, headers) p
57      hlen <- (#peek struct sf_hdtr, hdr_cnt)  p
58      return $ SfHdtr hdr hlen
59
60  poke p sfhdtr = do
61      (#poke struct sf_hdtr, headers)  p (sfhdtrHdr sfhdtr)
62      (#poke struct sf_hdtr, hdr_cnt)  p (sfhdtrHdrLen sfhdtr)
63      (#poke struct sf_hdtr, trailers) p nullPtr
64      (#poke struct sf_hdtr, trl_cnt)  p (0 :: CInt)
65
66----------------------------------------------------------------
67
68withIOVec :: [ByteString] -> ((Ptr IOVec, Int) -> IO a) -> IO a
69withIOVec cs f =
70    allocaArray csLen $ \aPtr -> do
71        zipWithM_ pokeIov (ptrs aPtr) cs
72        f (aPtr, csLen)
73  where
74    csLen = length cs
75    ptrs = iterate (`plusPtr` sizeOf (undefined :: IOVec))
76    pokeIov ptr s =
77        unsafeUseAsCStringLen s $ \(sPtr, sLen) ->
78        poke ptr $ IOVec sPtr (fromIntegral sLen)
79
80withSfHdtr :: [ByteString] -> (Ptr SfHdtr -> IO a) -> IO a
81withSfHdtr cs f = withIOVec cs $ \(iovecp,len) ->
82    alloca $ \sfhdtrp -> do
83        poke sfhdtrp $ SfHdtr iovecp (fromIntegral len)
84        f sfhdtrp
85
86----------------------------------------------------------------
87
88remainingChunks :: Int -> [ByteString] -> [ByteString]
89remainingChunks _ [] = []
90remainingChunks i (x:xs)
91    | i < len        = BS.drop i x : xs
92    | otherwise      = let i' = i - len in i' `seq` remainingChunks i' xs
93  where
94    len = BS.length x
95