diff options
Diffstat (limited to 'src/Data/Torrent/Bitfield.hs')
-rw-r--r-- | src/Data/Torrent/Bitfield.hs | 324 |
1 files changed, 0 insertions, 324 deletions
diff --git a/src/Data/Torrent/Bitfield.hs b/src/Data/Torrent/Bitfield.hs deleted file mode 100644 index b65f058b..00000000 --- a/src/Data/Torrent/Bitfield.hs +++ /dev/null | |||
@@ -1,324 +0,0 @@ | |||
1 | -- | | ||
2 | -- Copyright : (c) Sam Truzjan 2013 | ||
3 | -- License : BSD3 | ||
4 | -- Maintainer : pxqr.sta@gmail.com | ||
5 | -- Stability : experimental | ||
6 | -- Portability : portable | ||
7 | -- | ||
8 | -- This modules provides all necessary machinery to work with | ||
9 | -- bitfields. Bitfields are used to keep track indices of complete | ||
10 | -- pieces either peer have or client have. | ||
11 | -- | ||
12 | -- There are also commonly used piece seletion algorithms | ||
13 | -- which used to find out which one next piece to download. | ||
14 | -- Selectors considered to be used in the following order: | ||
15 | -- | ||
16 | -- * Random first - at the start. | ||
17 | -- | ||
18 | -- * Rarest first selection - performed to avoid situation when | ||
19 | -- rarest piece is unaccessible. | ||
20 | -- | ||
21 | -- * /End game/ seletion - performed after a peer has requested all | ||
22 | -- the subpieces of the content. | ||
23 | -- | ||
24 | -- Note that BitTorrent applies the strict priority policy for | ||
25 | -- /subpiece/ or /blocks/ selection. | ||
26 | -- | ||
27 | {-# LANGUAGE CPP #-} | ||
28 | {-# LANGUAGE BangPatterns #-} | ||
29 | {-# LANGUAGE RecordWildCards #-} | ||
30 | module Data.Torrent.Bitfield | ||
31 | ( -- * Bitfield | ||
32 | PieceIx | ||
33 | , PieceCount | ||
34 | , Bitfield | ||
35 | |||
36 | -- * Construction | ||
37 | , haveAll | ||
38 | , haveNone | ||
39 | , have | ||
40 | , singleton | ||
41 | , interval | ||
42 | , adjustSize | ||
43 | |||
44 | -- * Query | ||
45 | -- ** Cardinality | ||
46 | , Data.Torrent.Bitfield.null | ||
47 | , Data.Torrent.Bitfield.full | ||
48 | , haveCount | ||
49 | , totalCount | ||
50 | , completeness | ||
51 | |||
52 | -- ** Membership | ||
53 | , member | ||
54 | , notMember | ||
55 | , findMin | ||
56 | , findMax | ||
57 | , isSubsetOf | ||
58 | |||
59 | -- ** Availability | ||
60 | , complement | ||
61 | , Frequency | ||
62 | , frequencies | ||
63 | , rarest | ||
64 | |||
65 | -- * Combine | ||
66 | , insert | ||
67 | , union | ||
68 | , intersection | ||
69 | , difference | ||
70 | |||
71 | -- * Conversion | ||
72 | , toList | ||
73 | , fromList | ||
74 | |||
75 | -- * Serialization | ||
76 | , fromBitmap | ||
77 | , toBitmap | ||
78 | ) where | ||
79 | |||
80 | import Control.Monad | ||
81 | import Control.Monad.ST | ||
82 | import Data.ByteString (ByteString) | ||
83 | import qualified Data.ByteString as B | ||
84 | import qualified Data.ByteString.Lazy as Lazy | ||
85 | import Data.Vector.Unboxed (Vector) | ||
86 | import qualified Data.Vector.Unboxed as V | ||
87 | import qualified Data.Vector.Unboxed.Mutable as VM | ||
88 | import Data.IntervalSet (IntSet) | ||
89 | import qualified Data.IntervalSet as S | ||
90 | import qualified Data.IntervalSet.ByteString as S | ||
91 | import Data.List (foldl') | ||
92 | import Data.Monoid | ||
93 | import Data.Ratio | ||
94 | |||
95 | import Data.Torrent.Piece | ||
96 | |||
97 | -- TODO cache some operations | ||
98 | |||
99 | -- | Bitfields are represented just as integer sets but with | ||
100 | -- restriction: the each set should be within given interval (or | ||
101 | -- subset of the specified interval). Size is used to specify | ||
102 | -- interval, so bitfield of size 10 might contain only indices in | ||
103 | -- interval [0..9]. | ||
104 | -- | ||
105 | data Bitfield = Bitfield { | ||
106 | bfSize :: !PieceCount | ||
107 | , bfSet :: !IntSet | ||
108 | } deriving (Show, Read, Eq) | ||
109 | |||
110 | -- Invariants: all elements of bfSet lie in [0..bfSize - 1]; | ||
111 | |||
112 | instance Monoid Bitfield where | ||
113 | {-# SPECIALIZE instance Monoid Bitfield #-} | ||
114 | mempty = haveNone 0 | ||
115 | mappend = union | ||
116 | mconcat = unions | ||
117 | |||
118 | {----------------------------------------------------------------------- | ||
119 | Construction | ||
120 | -----------------------------------------------------------------------} | ||
121 | |||
122 | -- | The empty bitfield of the given size. | ||
123 | haveNone :: PieceCount -> Bitfield | ||
124 | haveNone s = Bitfield s S.empty | ||
125 | |||
126 | -- | The full bitfield containing all piece indices for the given size. | ||
127 | haveAll :: PieceCount -> Bitfield | ||
128 | haveAll s = Bitfield s (S.interval 0 (s - 1)) | ||
129 | |||
130 | -- | Insert the index in the set ignoring out of range indices. | ||
131 | have :: PieceIx -> Bitfield -> Bitfield | ||
132 | have ix Bitfield {..} | ||
133 | | 0 <= ix && ix < bfSize = Bitfield bfSize (S.insert ix bfSet) | ||
134 | | otherwise = Bitfield bfSize bfSet | ||
135 | |||
136 | singleton :: PieceIx -> PieceCount -> Bitfield | ||
137 | singleton ix pc = have ix (haveNone pc) | ||
138 | |||
139 | -- | Assign new size to bitfield. FIXME Normally, size should be only | ||
140 | -- decreased, otherwise exception raised. | ||
141 | adjustSize :: PieceCount -> Bitfield -> Bitfield | ||
142 | adjustSize s Bitfield {..} = Bitfield s bfSet | ||
143 | |||
144 | -- | NOTE: for internal use only | ||
145 | interval :: PieceCount -> PieceIx -> PieceIx -> Bitfield | ||
146 | interval pc a b = Bitfield pc (S.interval a b) | ||
147 | |||
148 | {----------------------------------------------------------------------- | ||
149 | Query | ||
150 | -----------------------------------------------------------------------} | ||
151 | |||
152 | -- | Test if bitifield have no one index: peer do not have anything. | ||
153 | null :: Bitfield -> Bool | ||
154 | null Bitfield {..} = S.null bfSet | ||
155 | |||
156 | -- | Test if bitfield have all pieces. | ||
157 | full :: Bitfield -> Bool | ||
158 | full Bitfield {..} = S.size bfSet == bfSize | ||
159 | |||
160 | -- | Count of peer have pieces. | ||
161 | haveCount :: Bitfield -> PieceCount | ||
162 | haveCount = S.size . bfSet | ||
163 | |||
164 | -- | Total count of pieces and its indices. | ||
165 | totalCount :: Bitfield -> PieceCount | ||
166 | totalCount = bfSize | ||
167 | |||
168 | -- | Ratio of /have/ piece count to the /total/ piece count. | ||
169 | -- | ||
170 | -- > forall bf. 0 <= completeness bf <= 1 | ||
171 | -- | ||
172 | completeness :: Bitfield -> Ratio PieceCount | ||
173 | completeness b = haveCount b % totalCount b | ||
174 | |||
175 | inRange :: PieceIx -> Bitfield -> Bool | ||
176 | inRange ix Bitfield {..} = 0 <= ix && ix < bfSize | ||
177 | |||
178 | member :: PieceIx -> Bitfield -> Bool | ||
179 | member ix bf @ Bitfield {..} | ||
180 | | ix `inRange` bf = ix `S.member` bfSet | ||
181 | | otherwise = False | ||
182 | |||
183 | notMember :: PieceIx -> Bitfield -> Bool | ||
184 | notMember ix bf @ Bitfield {..} | ||
185 | | ix `inRange` bf = ix `S.notMember` bfSet | ||
186 | | otherwise = True | ||
187 | |||
188 | -- | Find first available piece index. | ||
189 | findMin :: Bitfield -> PieceIx | ||
190 | findMin = S.findMin . bfSet | ||
191 | {-# INLINE findMin #-} | ||
192 | |||
193 | -- | Find last available piece index. | ||
194 | findMax :: Bitfield -> PieceIx | ||
195 | findMax = S.findMax . bfSet | ||
196 | {-# INLINE findMax #-} | ||
197 | |||
198 | -- | Check if all pieces from first bitfield present if the second bitfield | ||
199 | isSubsetOf :: Bitfield -> Bitfield -> Bool | ||
200 | isSubsetOf a b = bfSet a `S.isSubsetOf` bfSet b | ||
201 | {-# INLINE isSubsetOf #-} | ||
202 | |||
203 | -- | Resulting bitfield includes only missing pieces. | ||
204 | complement :: Bitfield -> Bitfield | ||
205 | complement Bitfield {..} = Bitfield | ||
206 | { bfSet = uni `S.difference` bfSet | ||
207 | , bfSize = bfSize | ||
208 | } | ||
209 | where | ||
210 | Bitfield _ uni = haveAll bfSize | ||
211 | {-# INLINE complement #-} | ||
212 | |||
213 | {----------------------------------------------------------------------- | ||
214 | -- Availability | ||
215 | -----------------------------------------------------------------------} | ||
216 | |||
217 | -- | Frequencies are needed in piece selection startegies which use | ||
218 | -- availability quantity to find out the optimal next piece index to | ||
219 | -- download. | ||
220 | type Frequency = Int | ||
221 | |||
222 | -- TODO rename to availability | ||
223 | -- | How many times each piece index occur in the given bitfield set. | ||
224 | frequencies :: [Bitfield] -> Vector Frequency | ||
225 | frequencies [] = V.fromList [] | ||
226 | frequencies xs = runST $ do | ||
227 | v <- VM.new size | ||
228 | VM.set v 0 | ||
229 | forM_ xs $ \ Bitfield {..} -> do | ||
230 | forM_ (S.toList bfSet) $ \ x -> do | ||
231 | fr <- VM.read v x | ||
232 | VM.write v x (succ fr) | ||
233 | V.unsafeFreeze v | ||
234 | where | ||
235 | size = maximum (map bfSize xs) | ||
236 | |||
237 | -- TODO it seems like this operation is veeery slow | ||
238 | |||
239 | -- | Find least available piece index. If no piece available return | ||
240 | -- 'Nothing'. | ||
241 | rarest :: [Bitfield] -> Maybe PieceIx | ||
242 | rarest xs | ||
243 | | V.null freqMap = Nothing | ||
244 | | otherwise | ||
245 | = Just $ fst $ V.ifoldr' minIx (0, freqMap V.! 0) freqMap | ||
246 | where | ||
247 | freqMap = frequencies xs | ||
248 | |||
249 | minIx :: PieceIx -> Frequency | ||
250 | -> (PieceIx, Frequency) | ||
251 | -> (PieceIx, Frequency) | ||
252 | minIx ix fr acc@(_, fra) | ||
253 | | fr < fra && fr > 0 = (ix, fr) | ||
254 | | otherwise = acc | ||
255 | |||
256 | |||
257 | {----------------------------------------------------------------------- | ||
258 | Combine | ||
259 | -----------------------------------------------------------------------} | ||
260 | |||
261 | insert :: PieceIx -> Bitfield -> Bitfield | ||
262 | insert pix bf @ Bitfield {..} | ||
263 | | 0 <= pix && pix < bfSize = Bitfield | ||
264 | { bfSet = S.insert pix bfSet | ||
265 | , bfSize = bfSize | ||
266 | } | ||
267 | | otherwise = bf | ||
268 | |||
269 | -- | Find indices at least one peer have. | ||
270 | union :: Bitfield -> Bitfield -> Bitfield | ||
271 | union a b = {-# SCC union #-} Bitfield { | ||
272 | bfSize = bfSize a `max` bfSize b | ||
273 | , bfSet = bfSet a `S.union` bfSet b | ||
274 | } | ||
275 | |||
276 | -- | Find indices both peers have. | ||
277 | intersection :: Bitfield -> Bitfield -> Bitfield | ||
278 | intersection a b = {-# SCC intersection #-} Bitfield { | ||
279 | bfSize = bfSize a `min` bfSize b | ||
280 | , bfSet = bfSet a `S.intersection` bfSet b | ||
281 | } | ||
282 | |||
283 | -- | Find indices which have first peer but do not have the second peer. | ||
284 | difference :: Bitfield -> Bitfield -> Bitfield | ||
285 | difference a b = {-# SCC difference #-} Bitfield { | ||
286 | bfSize = bfSize a -- FIXME is it reasonable? | ||
287 | , bfSet = bfSet a `S.difference` bfSet b | ||
288 | } | ||
289 | |||
290 | -- | Find indices the any of the peers have. | ||
291 | unions :: [Bitfield] -> Bitfield | ||
292 | unions = {-# SCC unions #-} foldl' union (haveNone 0) | ||
293 | |||
294 | {----------------------------------------------------------------------- | ||
295 | Serialization | ||
296 | -----------------------------------------------------------------------} | ||
297 | |||
298 | -- | List all /have/ indexes. | ||
299 | toList :: Bitfield -> [PieceIx] | ||
300 | toList Bitfield {..} = S.toList bfSet | ||
301 | |||
302 | -- | Make bitfield from list of /have/ indexes. | ||
303 | fromList :: PieceCount -> [PieceIx] -> Bitfield | ||
304 | fromList s ixs = Bitfield { | ||
305 | bfSize = s | ||
306 | , bfSet = S.splitGT (-1) $ S.splitLT s $ S.fromList ixs | ||
307 | } | ||
308 | |||
309 | -- | Unpack 'Bitfield' from tightly packed bit array. Note resulting | ||
310 | -- size might be more than real bitfield size, use 'adjustSize'. | ||
311 | fromBitmap :: ByteString -> Bitfield | ||
312 | fromBitmap bs = {-# SCC fromBitmap #-} Bitfield { | ||
313 | bfSize = B.length bs * 8 | ||
314 | , bfSet = S.fromByteString bs | ||
315 | } | ||
316 | {-# INLINE fromBitmap #-} | ||
317 | |||
318 | -- | Pack a 'Bitfield' to tightly packed bit array. | ||
319 | toBitmap :: Bitfield -> Lazy.ByteString | ||
320 | toBitmap Bitfield {..} = {-# SCC toBitmap #-} Lazy.fromChunks [intsetBM, alignment] | ||
321 | where | ||
322 | byteSize = bfSize `div` 8 + if bfSize `mod` 8 == 0 then 0 else 1 | ||
323 | alignment = B.replicate (byteSize - B.length intsetBM) 0 | ||
324 | intsetBM = S.toByteString bfSet | ||