1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
|
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS -fno-warn-orphans #-}
module System.Torrent.FileMap
( FileMap
-- * Construction
, Mode (..)
, def
, mmapFiles
, unmapFiles
-- * Query
, System.Torrent.FileMap.size
-- * Modification
, readBytes
, writeBytes
, unsafeReadBytes
-- * Unsafe conversions
, fromLazyByteString
, toLazyByteString
) where
import Control.Applicative
import Control.Monad as L
import Data.ByteString as BS
import Data.ByteString.Internal as BS
import Data.ByteString.Lazy as BL
import Data.ByteString.Lazy.Internal as BL
import Data.Default
import Data.Vector as V -- TODO use unboxed vector
import Foreign
import System.IO.MMap
import Data.Torrent
data FileEntry = FileEntry
{ filePosition :: {-# UNPACK #-} !FileOffset
, fileBytes :: {-# UNPACK #-} !BS.ByteString -- XXX: mutable buffer (see 'writeBytes').
} deriving (Show, Eq)
type FileMap = Vector FileEntry
instance Default Mode where
def = ReadWriteEx
mmapFiles :: Mode -> FileLayout FileSize -> IO FileMap
mmapFiles mode layout = V.fromList <$> L.mapM mkEntry (accumPositions layout)
where
mkEntry (path, (pos, expectedSize)) = do
let esize = fromIntegral expectedSize -- FIXME does this safe?
(fptr, moff, msize) <- mmapFileForeignPtr path mode $ Just (0, esize)
if msize /= esize
then error "mmapFiles" -- TODO unmap mapped files on exception
else return $ FileEntry pos (PS fptr moff msize)
unmapFiles :: FileMap -> IO ()
unmapFiles = V.mapM_ unmapEntry
where
unmapEntry (FileEntry _ (PS fptr _ _)) = finalizeForeignPtr fptr
-- Unsafe: FileMap 'writeBytes' will modify supplied bytestrings in place.
fromLazyByteString :: BL.ByteString -> FileMap
fromLazyByteString lbs = V.unfoldr f (0, lbs)
where
f (_, Empty ) = Nothing
f (pos, Chunk x xs) = Just (FileEntry pos x, ((pos + chunkSize), xs))
where chunkSize = fromIntegral $ BS.length x
-- | /O(n)/.
--
-- Unsafe: mutable buffers are returned without copy.
toLazyByteString :: FileMap -> BL.ByteString
toLazyByteString = V.foldr f Empty
where
f FileEntry {..} bs = Chunk fileBytes bs
-- | /O(1)/.
size :: FileMap -> FileOffset
size m
| V.null m = 0
| FileEntry {..} <- V.unsafeLast m
= filePosition + fromIntegral (BS.length fileBytes)
-- | Find the file number for a particular byte offset within a torrent.
bsearch :: FileOffset -> FileMap -> Maybe Int
bsearch x m
| V.null m = Nothing
| otherwise = branch (V.length m `div` 2)
where
branch c @ ((m !) -> FileEntry {..})
| x < filePosition = bsearch x (V.take c m)
| x >= filePosition + fileSize = do
ix <- bsearch x (V.drop (succ c) m)
return $ succ c + ix
| otherwise = Just c
where
fileSize = fromIntegral (BS.length fileBytes)
-- | /O(log n)/.
drop :: FileOffset -> FileMap -> (FileSize, FileMap)
drop off m
| Just ix <- bsearch off m
, FileEntry {..} <- m ! ix = (off - filePosition, V.drop ix m)
| otherwise = (0 , V.empty)
-- | /O(log n)/.
take :: FileSize -> FileMap -> (FileMap, FileSize)
take len m
| len >= s = (m , 0)
| Just ix <- bsearch (pred len) m = let m' = V.take (succ ix) m
in (m', System.Torrent.FileMap.size m' - len)
| otherwise = (V.empty , 0)
where
s = System.Torrent.FileMap.size m
-- | /O(log n + m)/. Do not use this function with 'unmapFiles'.
--
-- The returned bytestring points directly into an area memory mapped from a
-- file.
unsafeReadBytes :: FileOffset -> FileSize -> FileMap -> BL.ByteString
unsafeReadBytes off s m
| (l , m') <- System.Torrent.FileMap.drop off m
, (m'', _ ) <- System.Torrent.FileMap.take (off + s) m'
= BL.take (fromIntegral s) $ BL.drop (fromIntegral l) $ toLazyByteString m''
-- The returned bytestring is copied and safe to use after the file is
-- unmapped.
readBytes :: FileOffset -> FileSize -> FileMap -> IO BL.ByteString
readBytes off s m = do
let bs_copy = BL.copy $ unsafeReadBytes off s m
forceLBS bs_copy
return bs_copy
where
forceLBS Empty = return ()
forceLBS (Chunk _ x) = forceLBS x
-- UNSAFE: Uses the first byte string as a pointer to mutable data and writes
-- the contents of the second bytestring there.
bscpy :: BL.ByteString -> BL.ByteString -> IO ()
bscpy (PS _ _ 0 `Chunk` dest_rest) src = bscpy dest_rest src
bscpy dest (PS _ _ 0 `Chunk` src_rest) = bscpy dest src_rest
bscpy (PS dest_fptr dest_off dest_size `Chunk` dest_rest)
(PS src_fptr src_off src_size `Chunk` src_rest)
= do let csize = min dest_size src_size
withForeignPtr dest_fptr $ \dest_ptr ->
withForeignPtr src_fptr $ \src_ptr ->
memcpy (dest_ptr `advancePtr` dest_off)
(src_ptr `advancePtr` src_off)
(fromIntegral csize) -- TODO memmove?
bscpy (PS dest_fptr (dest_off + csize) (dest_size - csize) `Chunk` dest_rest)
(PS src_fptr (src_off + csize) (src_size - csize) `Chunk` src_rest)
bscpy _ _ = return ()
-- UNSAFE: Mutates bytestring contents within the provided FileMap.
writeBytes :: FileOffset -> BL.ByteString -> FileMap -> IO ()
writeBytes off lbs m = bscpy dest src
where
src = BL.take (fromIntegral (BL.length dest)) lbs
dest = unsafeReadBytes off (fromIntegral (BL.length lbs)) m
|
