summaryrefslogtreecommitdiff
path: root/src/Data/Torrent/Bitfield.hs
diff options
context:
space:
mode:
Diffstat (limited to 'src/Data/Torrent/Bitfield.hs')
-rw-r--r--src/Data/Torrent/Bitfield.hs359
1 files changed, 359 insertions, 0 deletions
diff --git a/src/Data/Torrent/Bitfield.hs b/src/Data/Torrent/Bitfield.hs
new file mode 100644
index 00000000..f7e036d5
--- /dev/null
+++ b/src/Data/Torrent/Bitfield.hs
@@ -0,0 +1,359 @@
1-- |
2-- Copyright : (c) Sam T. 2013
3-- License : MIT
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 #-}
30module Data.Torrent.Bitfield
31 ( PieceIx, PieceCount, Bitfield
32
33 -- * Construction
34 , haveAll, haveNone, have, singleton
35 , interval
36 , adjustSize
37
38 -- * Query
39 , Data.Torrent.Bitfield.null
40 , haveCount, totalCount, completeness
41
42 , member, notMember
43 , findMin, findMax
44
45 , isSubsetOf
46
47 , Frequency, frequencies, rarest
48
49 -- * Combine
50 , union
51 , intersection
52 , difference
53
54 -- * Serialization
55 , fromBitmap, toBitmap
56 , toList
57
58 -- * Selection
59 , Selector
60 , selector, strategyClass
61
62 , strictFirst, strictLast
63 , rarestFirst, randomFirst, endGame
64
65#if defined (TESTING)
66 -- * Debug
67 , mkBitfield
68#endif
69 ) where
70
71import Control.Monad
72import Control.Monad.ST
73import Data.ByteString (ByteString)
74import qualified Data.ByteString as B
75import qualified Data.ByteString.Lazy as Lazy
76import Data.Vector.Unboxed (Vector)
77import qualified Data.Vector.Unboxed as V
78import qualified Data.Vector.Unboxed.Mutable as VM
79import Data.IntervalSet (IntSet)
80import qualified Data.IntervalSet as S
81import qualified Data.IntervalSet.ByteString as S
82import Data.List (foldl')
83import Data.Monoid
84import Data.Ratio
85
86
87-- | Pieces indexed from zero up to 'PieceCount' value.
88type PieceIx = Int
89
90-- | Used to represent max set bound. Min set bound is always set to
91-- zero.
92type PieceCount = Int
93
94-- TODO cache some operations
95
96-- | Bitfields are represented just as integer sets but with
97-- restriction: the each set should be within given interval (or
98-- subset of the specified interval). Size is used to specify
99-- interval, so bitfield of size 10 might contain only indices in
100-- interval [0..9].
101--
102data Bitfield = Bitfield {
103 bfSize :: !PieceCount
104 , bfSet :: !IntSet
105 } deriving (Show, Read, Eq)
106
107-- Invariants: all elements of bfSet lie in [0..bfSize - 1];
108
109instance Monoid Bitfield where
110 {-# SPECIALIZE instance Monoid Bitfield #-}
111 mempty = haveNone 0
112 mappend = union
113 mconcat = unions
114
115{-----------------------------------------------------------------------
116 Construction
117-----------------------------------------------------------------------}
118
119-- | The empty bitfield of the given size.
120haveNone :: PieceCount -> Bitfield
121haveNone s = Bitfield s S.empty
122
123-- | The full bitfield containing all piece indices for the given size.
124haveAll :: PieceCount -> Bitfield
125haveAll s = Bitfield s (S.interval 0 (s - 1))
126
127-- | Insert the index in the set ignoring out of range indices.
128have :: PieceIx -> Bitfield -> Bitfield
129have ix Bitfield {..}
130 | 0 <= ix && ix < bfSize = Bitfield bfSize (S.insert ix bfSet)
131 | otherwise = Bitfield bfSize bfSet
132
133singleton :: PieceIx -> PieceCount -> Bitfield
134singleton ix pc = have ix (haveNone pc)
135
136-- | Assign new size to bitfield. FIXME Normally, size should be only
137-- decreased, otherwise exception raised.
138adjustSize :: PieceCount -> Bitfield -> Bitfield
139adjustSize s Bitfield {..} = Bitfield s bfSet
140
141-- | NOTE: for internal use only
142interval :: PieceCount -> PieceIx -> PieceIx -> Bitfield
143interval pc a b = Bitfield pc (S.interval a b)
144
145{-----------------------------------------------------------------------
146 Query
147-----------------------------------------------------------------------}
148
149-- | Test if bitifield have no one index: peer do not have anything.
150null :: Bitfield -> Bool
151null Bitfield {..} = S.null bfSet
152
153-- | Count of peer have pieces.
154haveCount :: Bitfield -> PieceCount
155haveCount = S.size . bfSet
156
157-- | Total count of pieces and its indices.
158totalCount :: Bitfield -> PieceCount
159totalCount = bfSize
160
161-- | Ratio of /have/ piece count to the /total/ piece count.
162--
163-- > forall bf. 0 <= completeness bf <= 1
164--
165completeness :: Bitfield -> Ratio PieceCount
166completeness b = haveCount b % totalCount b
167
168inRange :: PieceIx -> Bitfield -> Bool
169inRange ix Bitfield {..} = 0 <= ix && ix < bfSize
170
171member :: PieceIx -> Bitfield -> Bool
172member ix bf @ Bitfield {..}
173 | ix `inRange` bf = ix `S.member` bfSet
174 | otherwise = False
175
176notMember :: PieceIx -> Bitfield -> Bool
177notMember ix bf @ Bitfield {..}
178 | ix `inRange` bf = ix `S.notMember` bfSet
179 | otherwise = True
180
181-- | Find first available piece index.
182findMin :: Bitfield -> PieceIx
183findMin = S.findMin . bfSet
184{-# INLINE findMin #-}
185
186-- | Find last available piece index.
187findMax :: Bitfield -> PieceIx
188findMax = S.findMax . bfSet
189{-# INLINE findMax #-}
190
191isSubsetOf :: Bitfield -> Bitfield -> Bool
192isSubsetOf a b = bfSet a `S.isSubsetOf` bfSet b
193
194-- | Frequencies are needed in piece selection startegies which use
195-- availability quantity to find out the optimal next piece index to
196-- download.
197type Frequency = Int
198
199-- | How many times each piece index occur in the given bitfield set.
200frequencies :: [Bitfield] -> Vector Frequency
201frequencies [] = V.fromList []
202frequencies xs = runST $ do
203 v <- VM.new size
204 VM.set v 0
205 forM_ xs $ \ Bitfield {..} -> do
206 forM_ (S.toList bfSet) $ \ x -> do
207 fr <- VM.read v x
208 VM.write v x (succ fr)
209 V.unsafeFreeze v
210 where
211 size = maximum (map bfSize xs)
212
213-- TODO it seems like this operation is veeery slow
214
215-- | Find least available piece index. If no piece available return
216-- 'Nothing'.
217rarest :: [Bitfield] -> Maybe PieceIx
218rarest xs
219 | V.null freqMap = Nothing
220 | otherwise
221 = Just $ fst $ V.ifoldr' minIx (0, freqMap V.! 0) freqMap
222 where
223 freqMap = frequencies xs
224
225 minIx :: PieceIx -> Frequency
226 -> (PieceIx, Frequency)
227 -> (PieceIx, Frequency)
228 minIx ix fr acc@(_, fra)
229 | fr < fra && fr > 0 = (ix, fr)
230 | otherwise = acc
231
232
233{-----------------------------------------------------------------------
234 Combine
235-----------------------------------------------------------------------}
236
237-- | Find indices at least one peer have.
238union :: Bitfield -> Bitfield -> Bitfield
239union a b = {-# SCC union #-} Bitfield {
240 bfSize = bfSize a `max` bfSize b
241 , bfSet = bfSet a `S.union` bfSet b
242 }
243
244-- | Find indices both peers have.
245intersection :: Bitfield -> Bitfield -> Bitfield
246intersection a b = {-# SCC intersection #-} Bitfield {
247 bfSize = bfSize a `min` bfSize b
248 , bfSet = bfSet a `S.intersection` bfSet b
249 }
250
251-- | Find indices which have first peer but do not have the second peer.
252difference :: Bitfield -> Bitfield -> Bitfield
253difference a b = {-# SCC difference #-} Bitfield {
254 bfSize = bfSize a -- FIXME is it reasonable?
255 , bfSet = bfSet a `S.difference` bfSet b
256 }
257
258-- | Find indices the any of the peers have.
259unions :: [Bitfield] -> Bitfield
260unions = {-# SCC unions #-} foldl' union (haveNone 0)
261
262{-----------------------------------------------------------------------
263 Serialization
264-----------------------------------------------------------------------}
265
266-- | List all have indexes.
267toList :: Bitfield -> [PieceIx]
268toList Bitfield {..} = S.toList bfSet
269
270-- | Unpack 'Bitfield' from tightly packed bit array. Note resulting
271-- size might be more than real bitfield size, use 'adjustSize'.
272fromBitmap :: ByteString -> Bitfield
273fromBitmap bs = {-# SCC fromBitmap #-} Bitfield {
274 bfSize = B.length bs * 8
275 , bfSet = S.fromByteString bs
276 }
277{-# INLINE fromBitmap #-}
278
279-- | Pack a 'Bitfield' to tightly packed bit array.
280toBitmap :: Bitfield -> Lazy.ByteString
281toBitmap Bitfield {..} = {-# SCC toBitmap #-} Lazy.fromChunks [intsetBM, alignment]
282 where
283 byteSize = bfSize `div` 8 + if bfSize `mod` 8 == 0 then 0 else 1
284 alignment = B.replicate (byteSize - B.length intsetBM) 0
285 intsetBM = S.toByteString bfSet
286
287{-----------------------------------------------------------------------
288 Debug
289-----------------------------------------------------------------------}
290
291-- | For internal use only.
292mkBitfield :: PieceCount -> [PieceIx] -> Bitfield
293mkBitfield s ixs = Bitfield {
294 bfSize = s
295 , bfSet = S.splitGT (-1) $ S.splitLT s $ S.fromList ixs
296 }
297
298{-----------------------------------------------------------------------
299 Selection
300-----------------------------------------------------------------------}
301
302type Selector = Bitfield -- ^ Indices of client /have/ pieces.
303 -> Bitfield -- ^ Indices of peer /have/ pieces.
304 -> [Bitfield] -- ^ Indices of other peers /have/ pieces.
305 -> Maybe PieceIx -- ^ Zero-based index of piece to request
306 -- to, if any.
307
308selector :: Selector -- ^ Selector to use at the start.
309 -> Ratio PieceCount
310 -> Selector -- ^ Selector to use after the client have
311 -- the C pieces.
312 -> Selector -- ^ Selector that changes behaviour based
313 -- on completeness.
314selector start pt ready h a xs =
315 case strategyClass pt h of
316 SCBeginning -> start h a xs
317 SCReady -> ready h a xs
318 SCEnd -> endGame h a xs
319
320data StartegyClass
321 = SCBeginning
322 | SCReady
323 | SCEnd
324 deriving (Show, Eq, Ord, Enum, Bounded)
325
326
327strategyClass :: Ratio PieceCount -> Bitfield -> StartegyClass
328strategyClass threshold = classify . completeness
329 where
330 classify c
331 | c < threshold = SCBeginning
332 | c + 1 % numerator c < 1 = SCReady
333 -- FIXME numerator have is not total count
334 | otherwise = SCEnd
335
336
337-- | Select the first available piece.
338strictFirst :: Selector
339strictFirst h a _ = Just $ findMin (difference a h)
340
341-- | Select the last available piece.
342strictLast :: Selector
343strictLast h a _ = Just $ findMax (difference a h)
344
345-- |
346rarestFirst :: Selector
347rarestFirst h a xs = rarest (map (intersection want) xs)
348 where
349 want = difference h a
350
351-- | In average random first is faster than rarest first strategy but
352-- only if all pieces are available.
353randomFirst :: Selector
354randomFirst = do
355-- randomIO
356 error "randomFirst"
357
358endGame :: Selector
359endGame = strictLast