summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--bench/Main.hs69
-rw-r--r--network-bittorrent.cabal3
-rw-r--r--src/Data/Bitfield.hs443
-rw-r--r--src/Data/Bitfield/Mutable.hs177
-rw-r--r--src/Network/BitTorrent/PeerWire/Message.hs18
-rw-r--r--src/Network/BitTorrent/PeerWire/Selection.hs22
-rw-r--r--tests/Encoding.hs4
-rw-r--r--tests/Main.hs118
8 files changed, 175 insertions, 679 deletions
diff --git a/bench/Main.hs b/bench/Main.hs
index 6e8a0ce3..120f5b04 100644
--- a/bench/Main.hs
+++ b/bench/Main.hs
@@ -22,8 +22,7 @@ instance NFData BlockIx where
22instance NFData Block where 22instance NFData Block where
23 rnf (Block a b c) = a `deepseq` b `deepseq` rnf c 23 rnf (Block a b c) = a `deepseq` b `deepseq` rnf c
24 24
25instance NFData Bitfield where 25instance NFData Bitfield
26 rnf = rnf . bfBits
27 26
28instance NFData Message where 27instance NFData Message where
29 rnf (Have i) = rnf i 28 rnf (Have i) = rnf i
@@ -40,69 +39,5 @@ encodeMessages xs = runPut (mapM_ put xs)
40decodeMessages :: ByteString -> Either String [Message] 39decodeMessages :: ByteString -> Either String [Message]
41decodeMessages = runGet (many get) 40decodeMessages = runGet (many get)
42 41
43bitfieldMin :: Int -> Maybe Int
44bitfieldMin n = findMin (BT.empty n)
45
46bitfieldMax :: Int -> Maybe Int
47bitfieldMax n = findMax (BT.empty n)
48
49bitfieldDiff :: Int -> Bitfield
50bitfieldDiff n = BT.empty n `difference` BT.empty n
51
52bitfieldInter :: Int -> Bitfield
53bitfieldInter n = BT.empty n `intersection` BT.empty n
54
55bitfieldUnion :: Int -> Bitfield
56bitfieldUnion n = BT.empty n `union` BT.empty n
57
58bitfieldHaveCount :: Int -> Int
59bitfieldHaveCount n = haveCount (BT.full n)
60
61selectionStrictFirst :: Int -> Maybe Int
62selectionStrictFirst n = strictFirst (BT.empty n) (BT.empty n) []
63
64selectionStrictLast :: Int -> Maybe Int
65selectionStrictLast n = strictLast (BT.empty n) (BT.empty n) []
66
67selectionRarestFirst :: Int -> Maybe Int
68selectionRarestFirst n = rarestFirst (BT.empty n) (BT.empty n)
69 (replicate 10 (BT.empty n))
70
71selectionEndGame :: Int -> Maybe Int
72selectionEndGame n = endGame (BT.empty n) (BT.empty n) []
73
74main :: IO () 42main :: IO ()
75main = do 43main = defaultMain []
76 let blockixs = replicate 5000 (Request (BlockIx 0 0 0))
77 let bitfields = replicate 5000 (Bitfield (MkBitfield (B.replicate 1000 0)))
78 let chokes = replicate 5000 Choke
79 let havenones = replicate 5000 HaveNone
80
81 let m = 1024 * 1024
82
83 defaultMain $
84 concatMap (uncurry mkMsgBench)
85 [ ("blockIx", blockixs)
86 , ("bitfield", bitfields)
87 , ("choke", chokes)
88 , ("havenone", havenones)
89 ]
90 ++ -- 256KiB * 10M = 2.5TiB bitfield for 10 ms
91 [ bench "bitfield/min" $ nf bitfieldMin (10 * m)
92 , bench "bitfield/max" $ nf bitfieldMax (10 * m)
93 , bench "bitfield/difference" $ nf bitfieldDiff (10 * m)
94 , bench "bitfield/intersection" $ nf bitfieldInter (10 * m)
95 , bench "bitfield/union" $ nf bitfieldUnion (10 * m)
96 , bench "bitfield/haveCount" $ nf bitfieldHaveCount (10 * m)
97
98 , bench "selection/strictFirst" $ nf selectionStrictFirst (10 * m)
99 , bench "selection/strictLast" $ nf selectionStrictLast (10 * m)
100 , bench "selection/rarestFirst" $ nf selectionRarestFirst (10 * m)
101 , bench "selection/endGame" $ nf selectionEndGame (10 * m)
102 ]
103 where
104 mkMsgBench name msgs =
105 [ msgs `deepseq` bench ("message/" ++ name ++ "/encode") $ nf encodeMessages msgs
106 , let binary = encodeMessages msgs in
107 binary `deepseq` bench ("message/" ++ name ++ "/decode") $ nf decodeMessages binary
108 ]
diff --git a/network-bittorrent.cabal b/network-bittorrent.cabal
index 0a89603a..11d371e6 100644
--- a/network-bittorrent.cabal
+++ b/network-bittorrent.cabal
@@ -24,7 +24,6 @@ library
24 exposed-modules: Data.Torrent 24 exposed-modules: Data.Torrent
25 , Data.Torrent.InfoHash 25 , Data.Torrent.InfoHash
26 , Data.Bitfield 26 , Data.Bitfield
27 , Data.Bitfield.Mutable
28 27
29 , Network.BitTorrent 28 , Network.BitTorrent
30 , Network.BitTorrent.Extension 29 , Network.BitTorrent.Extension
@@ -56,7 +55,9 @@ library
56 , array >= 0.4 55 , array >= 0.4
57 , bytestring >= 0.10.2 56 , bytestring >= 0.10.2
58 , containers >= 0.4 57 , containers >= 0.4
58 , intset >= 0.1
59 , text >= 0.11.0 59 , text >= 0.11.0
60 , vector
60 61
61 -- encoding/serialization packages 62 -- encoding/serialization packages
62 , bencoding >= 0.1 63 , bencoding >= 0.1
diff --git a/src/Data/Bitfield.hs b/src/Data/Bitfield.hs
index f7240c8a..546c68e9 100644
--- a/src/Data/Bitfield.hs
+++ b/src/Data/Bitfield.hs
@@ -5,341 +5,196 @@
5-- Stability : experimental 5-- Stability : experimental
6-- Portability : portable 6-- Portability : portable
7-- 7--
8-- 8-- This modules provides all necessary machinery to work with
9-- This module provides Bitfield datatype used to represent sets of 9-- bitfields. Bitfields are used to keep track indices of complete
10-- piece indexes any peer have. All associated operations should be 10-- pieces either peer have or client have.
11-- defined here as well.
12-- 11--
13{-# LANGUAGE BangPatterns #-} 12{-# LANGUAGE BangPatterns #-}
14{-# LANGUAGE RankNTypes #-} 13{-# LANGUAGE RecordWildCards #-}
15module Data.Bitfield 14module Data.Bitfield
16-- TODO: move to Data.Bitfield 15 ( Bitfield, PieceCount
17 ( Bitfield(..)
18 16
19 -- * Construction 17 -- * Construction
20 , empty, full 18 , empty
19 , insert
20 , haveAll, haveNone, have
21 21
22 -- * Query 22 -- * Query
23 , bitfieldByteCount, bitfieldBitCount 23 , haveCount, totalCount, completeness
24
25 , haveCount, completeness
26 , findMin, findMax 24 , findMin, findMax
27 , union, intersection, difference, combine 25 , frequencies, rarest
28 , frequencies 26
27 -- * Combine
28 , union
29 , intersection
30 , difference
29 31
30 -- * Serialization 32 -- * Serialization
31 , getBitfield, putBitfield 33 , getBitfield, putBitfield
34 , bitfieldByteCount
32 35
33 -- * Conversion 36 , -- * Debug
34 , toList 37 mkBitfield
35 , fromByteString, toByteString
36
37 -- * Debug
38 , aligned, alignLow, alignedZip
39 ) where 38 ) where
40 39
41import Control.Applicative hiding (empty) 40import Control.Monad
42import Data.Bits 41import Control.Monad.ST
43import Data.ByteString (ByteString) 42import Data.Vector.Unboxed (Vector)
44import qualified Data.ByteString as B 43import qualified Data.Vector.Unboxed as V
45import qualified Data.ByteString.Internal as B 44import qualified Data.Vector.Unboxed.Mutable as VM
46import Data.List as L hiding (union) 45import Data.IntervalSet (IntSet)
47import Data.Serialize 46import qualified Data.IntervalSet as S
48import Data.Word 47import Data.List (foldl')
48import Data.Monoid
49import Data.Ratio
50import Data.Serialize
51import Network.BitTorrent.PeerWire.Block
49 52
50import Foreign
51 53
52--import Network.BitTorrent.PeerWire.Block 54type PieceCount = Int
53import Data.Torrent
54 55
55-- TODO: one good idea is to aggregate frequently used stats in reducer 56-- TODO cache some operations
56-- it should give a big boost
57newtype Bitfield = MkBitfield {
58 bfBits :: ByteString
59-- , bfSize :: Int
60 } deriving (Show, Eq, Ord)
61 57
58-- | Bitfields are represented just as integer sets but with
59-- restriction: the each set should be within given interval (or
60-- subset of). Size is used to specify interval, so bitfield of size
61-- 10 might contain only indices in interval [0..9].
62--
63data Bitfield = Bitfield {
64 bfSize :: !PieceCount
65 , bfSet :: !IntSet
66 } deriving (Show, Read, Eq)
62 67
63empty :: Int -> Bitfield 68-- Invariants: all elements of bfSet lie in [0..bfSize - 1];
64empty n = MkBitfield $ B.replicate (sizeInBase n 8) 0
65{-# INLINE empty #-}
66 69
67full :: Int -> Bitfield 70instance Monoid Bitfield where
68full n = MkBitfield $ B.replicate (sizeInBase n 8) (complement 0) 71 {-# SPECIALIZE instance Monoid Bitfield #-}
69{-# INLINE full #-} 72 mempty = empty 0
73 mappend = union
74 mconcat = unions
70 75
71toList :: Bitfield -> [Bool] 76-- TODO documentation
72toList (MkBitfield bs) = concatMap unpkg (B.unpack bs) 77{-----------------------------------------------------------------------
73 where 78 Construction
74 unpkg :: Word8 -> [Bool] 79-----------------------------------------------------------------------}
75 unpkg byte = L.map (testBit byte) [0..bitSize (undefined :: Word8) - 1]
76{-# INLINE toList #-}
77 80
78fromByteString :: ByteString -> Bitfield 81empty :: PieceCount -> Bitfield
79fromByteString = MkBitfield 82empty s = Bitfield s S.empty
80{-# INLINE fromByteString #-}
81 83
82toByteString :: Bitfield -> ByteString 84insert :: PieceIx -> Bitfield -> Bitfield
83toByteString = bfBits 85insert ix Bitfield {..}
84{-# INLINE toByteString #-} 86 | 0 <= ix && ix < bfSize = Bitfield bfSize (S.insert ix bfSet)
87 | otherwise = Bitfield bfSize bfSet
85 88
86getBitfield :: Int -> Get Bitfield 89haveNone :: PieceCount -> Bitfield
87getBitfield n = MkBitfield <$> getBytes n 90haveNone = empty
88{-# INLINE getBitfield #-}
89 91
90putBitfield :: Bitfield -> Put 92haveAll :: PieceCount -> Bitfield
91putBitfield = putByteString . bfBits 93haveAll s = Bitfield s (S.interval 0 (s - 1))
92{-# INLINE putBitfield #-}
93 94
94bitfieldByteCount :: Bitfield -> Int 95have :: PieceIx -> Bitfield -> Bitfield
95bitfieldByteCount = B.length . bfBits 96have = insert
96{-# INLINE bitfieldByteCount #-}
97 97
98-- WARN 98{-----------------------------------------------------------------------
99-- TODO 99 Query
100bitfieldBitCount :: Bitfield -> Int 100-----------------------------------------------------------------------}
101bitfieldBitCount bf = bitSize (undefined :: Word8) * bitfieldByteCount bf
102{-# INLINE bitfieldBitCount #-}
103 101
104align :: Storable a => Ptr a -> (Ptr a, Int) 102haveCount :: Bitfield -> PieceCount
105align p = tie (alignPtr p) undefined 103haveCount = S.size . bfSet
106 where 104
107 tie :: Storable a => (Int -> Ptr a) -> a -> (Ptr a, Int) 105totalCount :: Bitfield -> PieceCount
108 tie f a = (f (alignment a), (alignment a)) 106totalCount = bfSize
109
110alignLow :: Ptr Word8 -> Ptr Word
111alignLow ptr =
112 let alg = alignment (undefined :: Word)
113 aptr = alignPtr (castPtr ptr) alg :: Ptr Word
114 in
115 if ptr == castPtr aptr
116 then aptr
117 else castPtr ((castPtr aptr :: Ptr Word8) `advancePtr` negate alg)
118
119isAlignedBy :: Storable a => Ptr a -> Int -> Bool
120isAlignedBy ptr alg = alignPtr ptr alg == ptr
121
122type Mem a = (Ptr a, Int)
123
124aligned :: Storable a => Mem Word8 -> (Mem Word8, Mem a, Mem Word8)
125aligned (ptr, len) =
126 let lowPtr = ptr
127 lowLen = midPtr `minusPtr` ptr
128 midOff = lowLen
129 (midPtr, alg) = align (castPtr ptr)
130 midLen = alg * div (len - midOff) alg
131 midLenA = midLen `div` alg
132 hghOff = midOff + midLen
133 hghPtr = ptr `advancePtr` hghOff
134 hghLen = len - hghOff
135 in
136 ((lowPtr, lowLen), (midPtr, midLenA), (hghPtr, hghLen))
137 where
138{-# INLINE aligned #-}
139
140type Mem3 a = (Ptr a, Ptr a, Ptr a, Int)
141
142emptyMem3 :: Mem3 a
143emptyMem3 = (nullPtr, nullPtr, nullPtr, 0)
144
145-- assume resulting memory is aligned
146alignedZip :: Ptr Word8 -> Ptr Word8 -> Ptr Word8 -> Int
147 -> (Mem3 Word, Mem3 Word8)
148alignedZip aptr bptr cptr size =
149 let alg = alignment (undefined :: Word) in
150 if (aptr `isAlignedBy` alg) && (bptr `isAlignedBy` alg)
151 then
152 let asize = alignLow (nullPtr `plusPtr` size) `minusPtr` nullPtr
153 in
154 ( (castPtr aptr, castPtr bptr, castPtr cptr, asize `div` alg)
155 , ( aptr `advancePtr` asize
156 , bptr `advancePtr` asize
157 , cptr `advancePtr` asize
158 , (size - asize)
159 )
160 )
161 else (emptyMem3, (aptr, bptr, cptr, size))
162
163-- force specialization
164zipWithBS :: (Word -> Word -> Word)
165 -> (Word8 -> Word8 -> Word8)
166 -> ByteString -> ByteString -> ByteString
167zipWithBS f g a b =
168 let (afptr, aoff, asize) = B.toForeignPtr a
169 (bfptr, boff, bsize) = B.toForeignPtr b
170 size = min asize bsize in
171 B.unsafeCreate size $ \rptr -> do
172 withForeignPtr afptr $ \_aptr -> do
173 withForeignPtr bfptr $ \_bptr -> do
174 let aptr = _aptr `advancePtr` aoff
175 let bptr = _bptr `advancePtr` boff
176
177 let (mid, hgh) = alignedZip aptr bptr rptr size
178 zipWords mid
179 zipBytes hgh
180 where
181 zipBytes :: (Ptr Word8, Ptr Word8, Ptr Word8, Int) -> IO ()
182 zipBytes (aptr, bptr, rptr, n) = go 0
183 where
184 go :: Int -> IO ()
185 go i | i < n = do -- TODO unfold
186 av <- peekElemOff aptr i
187 bv <- peekElemOff bptr i
188 pokeElemOff rptr i (g av bv)
189 go (succ i)
190 | otherwise = return ()
191
192 zipWords :: (Ptr Word, Ptr Word, Ptr Word, Int) -> IO ()
193 zipWords (aptr, bptr, rptr, n) = go 0
194 where
195 go :: Int -> IO ()
196 go i | i < n = do -- TODO unfold
197 av <- peekElemOff aptr i
198 bv <- peekElemOff bptr i
199 pokeElemOff rptr i (f av bv)
200 go (succ i)
201 | otherwise = return ()
202
203
204
205zipWithBF :: (forall a. Bits a => a -> a -> a) -> Bitfield -> Bitfield -> Bitfield
206zipWithBF f a b = MkBitfield $ zipWithBS f f (bfBits a) (bfBits b)
207{-# INLINE zipWithBF #-}
208
209findSet :: ByteString -> Maybe Int
210findSet b =
211 let (fptr, off, len) = B.toForeignPtr b in
212 B.inlinePerformIO $ withForeignPtr fptr $ \_ptr -> do
213 let ptr = _ptr `advancePtr` off
214
215 let (low, mid, hgh) = aligned (ptr, len)
216 let lowOff = fst low `minusPtr` ptr
217 let midOff = fst mid `minusPtr` ptr
218 let hghOff = fst hgh `minusPtr` ptr
219
220 let resL = (lowOff +) <$> goFind low
221 let resM = (midOff +) <$> goFind (mid :: Mem Word) -- tune size here
222 -- TODO: with Word8
223 -- bytestring findIndex works 2
224 -- times faster.
225 let resH = (hghOff +) <$> goFind hgh
226
227 let res = resL <|> resM <|> resH
228
229 -- computation of res should not escape withForeignPtr
230 case res of
231 Nothing -> return ()
232 Just _ -> return ()
233
234 return res
235 107
108-- |
109--
110-- > forall bf. 0 <= completeness bf <= 1
111--
112completeness :: Bitfield -> Ratio PieceCount
113completeness b = haveCount b % totalCount b
114
115findMin :: Bitfield -> Maybe PieceIx
116findMin Bitfield {..}
117 | S.null bfSet = Nothing
118 | otherwise = Just (S.findMin bfSet)
119
120findMax :: Bitfield -> Maybe PieceIx
121findMax Bitfield {..}
122 | S.null bfSet = Nothing
123 | otherwise = Just (S.findMax bfSet)
124
125type Frequency = Int
126
127frequencies :: [Bitfield] -> Vector Frequency
128frequencies [] = V.fromList []
129frequencies xs = runST $ do
130 v <- VM.new size
131 VM.set v 0
132 forM_ xs $ \ Bitfield {..} -> do
133 forM_ (S.toList bfSet) $ \ x -> do
134 fr <- VM.read v x
135 VM.write v x (succ fr)
136 V.unsafeFreeze v
236 where 137 where
237 goFind :: (Storable a, Eq a, Num a) => Mem a -> Maybe Int 138 size = maximum (map bfSize xs)
238 goFind (ptr, n) = go 0
239 where
240 go :: Int -> Maybe Int
241 go i | i < n =
242 let v = B.inlinePerformIO (peekElemOff ptr i) in
243 if v /= 0
244 then Just i
245 else go (succ i)
246 | otherwise = Nothing
247
248foldBS :: (Word8 -> Int -> Int) -> (Word -> Int -> Int) -> Int -> ByteString -> Int
249foldBS f g acc b =
250 let (fptr, off, len) = B.toForeignPtr b in
251 B.inlinePerformIO $ withForeignPtr fptr $ \_ptr -> do
252 let ptr = _ptr `advancePtr` off
253
254 let (low, mid, hgh) = aligned (ptr, len)
255 let resL = goFold low acc
256 let resM = goFoldW (mid :: Mem Word) resL
257 let resH = goFold hgh resM
258
259 -- computation of res should not escape withForeignPtr
260 case resH of
261 0 -> return ()
262 _ -> return ()
263
264 return resH
265 139
140rarest :: [Bitfield] -> Maybe PieceIx
141rarest xs
142 | V.null freqMap = Nothing
143 | otherwise = Just $ fst $ V.ifoldr minIx (0, freqMap V.! 0) freqMap
266 where 144 where
267 goFold :: Mem Word8 -> Int -> Int 145 freqMap = frequencies xs
268 goFold (ptr, n) = go 0 146
269 where 147 minIx :: PieceIx -> Frequency -> (PieceIx, Frequency) -> (PieceIx, Frequency)
270 go :: Int -> Int -> Int 148 minIx ix fr acc@(_, fra)
271 go i !a 149 | fr < fra && fr > 0 = (ix, fr)
272 | i < n = 150 | otherwise = acc
273 let v = B.inlinePerformIO (peekElemOff ptr i) 151
274 in go (succ i) (f v a) 152
275 | otherwise = a 153
276 154{-----------------------------------------------------------------------
277 goFoldW :: Mem Word -> Int -> Int 155 Combine
278 goFoldW (ptr, n) = go 0 156-----------------------------------------------------------------------}
279 where
280 go :: Int -> Int -> Int
281 go i !a
282 | i < n =
283 let v = B.inlinePerformIO (peekElemOff ptr i)
284 in go (succ i) (g v a)
285 | otherwise = a
286 157
287union :: Bitfield -> Bitfield -> Bitfield 158union :: Bitfield -> Bitfield -> Bitfield
288union = zipWithBF (.|.) 159union a b = Bitfield {
289{-# INLINE union #-} 160 bfSize = bfSize a `max` bfSize b
161 , bfSet = bfSet a `S.union` bfSet b
162 }
290 163
291intersection :: Bitfield -> Bitfield -> Bitfield 164intersection :: Bitfield -> Bitfield -> Bitfield
292intersection = zipWithBF (.&.) 165intersection a b = Bitfield {
293{-# INLINE intersection #-} 166 bfSize = bfSize a `min` bfSize b
167 , bfSet = bfSet a `S.intersection` bfSet b
168 }
294 169
295difference :: Bitfield -> Bitfield -> Bitfield 170difference :: Bitfield -> Bitfield -> Bitfield
296difference = zipWithBF diffWord8 171difference a b = Bitfield {
297 where 172 bfSize = bfSize a -- FIXME is it more reasonable?
298 diffWord8 :: Bits a => a -> a -> a 173 , bfSet = bfSet a `S.difference` bfSet b
299 diffWord8 a b = a .&. (a `xor` b) 174 }
300 {-# INLINE diffWord8 #-}
301{-# INLINE difference #-}
302 175
303combine :: [Bitfield] -> Maybe Bitfield 176unions :: [Bitfield] -> Bitfield
304combine [] = Nothing 177unions = foldl' union (empty 0)
305combine as = return $ foldr1 intersection as
306 178
307haveCount :: Bitfield -> Int 179{-----------------------------------------------------------------------
308haveCount (MkBitfield b) = foldBS f f 0 b 180 Serialization
309 where 181-----------------------------------------------------------------------}
310 f byte count = popCount byte + count
311 182
312completeness :: Bitfield -> (Int, Int) 183getBitfield :: Int -> Get Bitfield
313completeness bf = (haveCount bf, bitfieldBitCount bf) 184getBitfield = error "getBitfield"
314 185
315-- | Get min index of piece that the peer have. 186putBitfield :: Bitfield -> Put
316findMin :: Bitfield -> Maybe Int 187putBitfield = error "putBitfield"
317findMin (MkBitfield b) = do 188
318 byteIx <- findSet b 189bitfieldByteCount :: Bitfield -> Int
319 bitIx <- findMinWord8 (B.index b byteIx) 190bitfieldByteCount = error "bitfieldByteCount"
320 return $ byteIx * bitSize (undefined :: Word8) + bitIx
321 where
322 -- TODO: bit tricks
323 findMinWord8 :: Word8 -> Maybe Int
324 findMinWord8 byte = L.find (testBit byte) [0..bitSize (undefined :: Word8) - 1]
325 {-# INLINE findMinWord8 #-}
326{-# INLINE findMin #-}
327
328
329findMax :: Bitfield -> Maybe Int
330findMax (MkBitfield b) = do
331 -- TODO avoid reverse
332 byteIx <- (pred (B.length b) -) <$> findSet (B.reverse b)
333 bitIx <- findMaxWord8 (B.index b byteIx)
334 return $ byteIx * bitSize (undefined :: Word8) + bitIx
335 where
336 -- TODO: bit tricks
337 findMaxWord8 :: Word8 -> Maybe Int
338 findMaxWord8 byte = L.find (testBit byte)
339 (reverse [0 :: Int ..
340 bitSize (undefined :: Word8) - 1])
341 191
342{-# INLINE findMax #-} 192{-----------------------------------------------------------------------
193 Debug
194-----------------------------------------------------------------------}
343 195
344frequencies :: [Bitfield] -> [Int] 196mkBitfield :: PieceCount -> [PieceIx] -> Bitfield
345frequencies xs = foldr1 (zipWith (+)) $ map (map fromEnum . toList) xs 197mkBitfield s ixs = Bitfield {
198 bfSize = s
199 , bfSet = S.splitLT s $ S.fromList ixs
200 } \ No newline at end of file
diff --git a/src/Data/Bitfield/Mutable.hs b/src/Data/Bitfield/Mutable.hs
deleted file mode 100644
index a2f736e9..00000000
--- a/src/Data/Bitfield/Mutable.hs
+++ /dev/null
@@ -1,177 +0,0 @@
1-- TODO: update documentation
2-- |
3-- Copyright : (c) Sam T. 2013
4-- License : MIT
5-- Maintainer : pxqr.sta@gmail.com
6-- Stability : experimental
7-- Portability : portable
8--
9--
10-- Set of integers with atomic modification. Internally IntSet
11-- represented as array of tightly packed bits.
12--
13-- Note that:
14--
15-- * Insertion, deletion are atomic, waitfree and failfree.
16--
17-- * You can avoid copying in conversion if you don't care about
18-- referencial transparency or sure that after conversion
19-- bitfields never modified.
20--
21--
22{-# OPTIONS -fno-warn-unused-do-bind #-}
23module Data.Bitfield.Mutable
24 ( Bitfield
25
26 -- * Construction
27 , empty, full
28 , create, releaseIntSet
29
30 -- * Query
31-- , lookup, member, notMember
32-- , size
33 , maxSize
34 , lookupUnsafe
35
36 -- * Modification
37-- , insert, delete
38 , insertUnsafe, deleteUnsafe
39
40 -- * Conversion
41 , fromByteString, toByteString
42 , fromByteStringUnsafe, toByteStringUnsafe
43 ) where
44
45import Control.Applicative hiding (empty)
46import Data.Bits.Atomic
47import Data.ByteString (ByteString)
48import qualified Data.ByteString as B
49import qualified Data.ByteString.Internal as B
50import Foreign
51
52
53-- | Basically 'BitSet' is a wrapper on the 'ForeignPtr'.
54data Bitfield = Bitfield {
55 bfBasePtr :: {-# UNPACK #-} !(ForeignPtr Word8)
56 , bfOffset :: {-# UNPACK #-} !Int
57 , bfByteSize :: {-# UNPACK #-} !Int
58 , bfMaxSize :: {-# UNPACK #-} !Int
59 } deriving Show
60
61
62maxSize :: Bitfield -> Int
63maxSize = bfMaxSize
64
65
66create :: Int -> (Int -> Ptr Word8 -> IO a) -> IO Bitfield
67create n f = do
68 let byteSize = sizeInBytes n
69 fptr <- mallocForeignPtrBytes byteSize
70 withForeignPtr fptr (f byteSize)
71 return (Bitfield fptr 0 byteSize n)
72
73-- | Create a 'IntSet' with a given size in /bits/.
74empty :: Int -> IO Bitfield
75empty n = create n $ \bn ptr ->
76 B.memset ptr 0 (fromIntegral bn)
77
78full :: Int -> IO Bitfield
79full n = create n $ \bn ptr ->
80 B.memset ptr (error "IntSet.full") (fromIntegral bn)
81
82
83-- | Should be used to free scarce resources immediately.
84--
85-- WARNING: After this call 'BitField' should not be used. Also you
86-- can avoid using it at all if resource is not too scarce.
87--
88releaseIntSet :: Bitfield -> IO ()
89releaseIntSet = finalizeForeignPtr . bfBasePtr
90
91-- | Set nth bit in the given BifField to 1.
92--
93-- UNSAFE: no bound checking.
94--
95insertUnsafe :: Int -> Bitfield -> IO ()
96insertUnsafe i s =
97 withByte s i $ \ptr -> do
98 fetchAndOr ptr (bit (bitLoc i))
99 return ()
100{-# INLINE insertUnsafe #-}
101
102
103deleteUnsafe :: Int -> Bitfield -> IO ()
104deleteUnsafe i s =
105 withByte s i $ \ptr -> do
106 fetchAndAnd ptr (complement (bit (bitLoc i)))
107 return ()
108{-# INLINE deleteUnsafe #-}
109
110-- | Get nth bit in the given BitField.
111--
112-- UNSAFE: no bound checking.
113--
114lookupUnsafe :: Int -> Bitfield -> IO Bool
115lookupUnsafe n s = withByte s n $ \ptr -> (`testBit` bitLoc n) <$> peek ptr
116{-# INLINE lookupUnsafe #-}
117
118fromByteString :: Int -> ByteString -> Bitfield
119fromByteString n = fromByteStringUnsafe n . B.copy
120{-# INLINE fromByteString #-}
121
122toByteString :: Bitfield -> ByteString
123toByteString = B.copy . toByteStringUnsafe
124{-# INLINE toByteString #-}
125
126-- | Convert a 'BitField' to the 'ByteString' /without/ copying,
127-- so we can write it to a socket or a file for exsample.
128--
129-- WARNING: Note that using the resulting 'ByteString' might (and
130-- even should) BREAK REFERENCIAL TRANSPARENCY since we can change
131-- bits using 'setBitN' after the conversion. Use this function
132-- wisely and if and only if you understand the consequences,
133-- otherwise the really BAD THINGS WILL HAPPEN or use safe version
134-- instead.
135--
136toByteStringUnsafe :: Bitfield -> ByteString
137toByteStringUnsafe = B.fromForeignPtr <$> bfBasePtr <*> pure 0 <*> bfByteSize
138
139
140-- | Convert a 'ByteString' to 'BitField' /without/ copying, so we can
141-- read it from a file or a socket.
142--
143-- WARNING: Please see 'toByteString' doc, the same apply to this function.
144--
145fromByteStringUnsafe :: Int -> ByteString -> Bitfield
146fromByteStringUnsafe n (B.PS fptr a b) = Bitfield fptr a b n
147
148baseSize :: (Bits a, Integral a) =>
149 a -- ^ Base, should be power of two.
150 -> a -- ^ Size.
151 -> a -- ^ Size in base.
152baseSize base n = (n `div` base) + fromIntegral (fromEnum ((n .&. 0x7) > 0))
153{-# SPECIALIZE baseSize :: Int -> Int -> Int #-}
154{-# SPECIALIZE baseSize :: Word64 -> Word64 -> Word64 #-}
155
156-------------------------------- internal --------------------------------------
157sizeInBytes :: Int -- ^ Length in bits.
158 -> Int -- ^ Length in bytes aligned by size of word.
159sizeInBytes = baseSize 8
160{-# INLINE sizeInBytes #-}
161
162-- TODO: see if shifts and bitwise ands are faster
163-- and make portable version if not
164byteLoc :: Int -> Int
165byteLoc i = i `div` 8 * sizeOf (error "byteLoc" :: Word8)
166{-# INLINE bitLoc #-}
167
168bitLoc :: Int -> Int
169bitLoc i = i `mod` 8 * sizeOf (error "bitLoc" :: Word8)
170{-# INLINE byteLoc #-}
171
172withByte :: Bitfield -> Int -> (Ptr Word8 -> IO a) -> IO a
173withByte s n action = do
174 let offset = bfOffset s + byteLoc n
175 withForeignPtr (bfBasePtr s) $ \ptr ->
176 action (ptr `advancePtr` offset)
177{-# INLINE withByte #-} \ No newline at end of file
diff --git a/src/Network/BitTorrent/PeerWire/Message.hs b/src/Network/BitTorrent/PeerWire/Message.hs
index 5544dca7..f5ad2693 100644
--- a/src/Network/BitTorrent/PeerWire/Message.hs
+++ b/src/Network/BitTorrent/PeerWire/Message.hs
@@ -28,27 +28,27 @@ data Message = KeepAlive
28 28
29 -- | Zero-based index of a piece that has just been 29 -- | Zero-based index of a piece that has just been
30 -- successfully downloaded and verified via the hash. 30 -- successfully downloaded and verified via the hash.
31 | Have PieceIx 31 | Have !PieceIx
32 32
33 -- | The bitfield message may only be sent immediately 33 -- | The bitfield message may only be sent immediately
34 -- after the handshaking sequence is complete, and 34 -- after the handshaking sequence is complete, and
35 -- before any other message are sent. If client have no 35 -- before any other message are sent. If client have no
36 -- pieces then bitfield need not to be sent. 36 -- pieces then bitfield need not to be sent.
37 | Bitfield Bitfield 37 | Bitfield !Bitfield
38 38
39 -- | Request for a particular block. If a client is 39 -- | Request for a particular block. If a client is
40 -- requested a block that another peer do not have the 40 -- requested a block that another peer do not have the
41 -- peer might not answer at all. 41 -- peer might not answer at all.
42 | Request BlockIx 42 | Request !BlockIx
43 43
44 -- | Response for a request for a block. 44 -- | Response for a request for a block.
45 | Piece Block 45 | Piece !Block
46 46
47 -- | Used to cancel block requests. It is typically 47 -- | Used to cancel block requests. It is typically
48 -- used during "End Game". 48 -- used during "End Game".
49 | Cancel BlockIx 49 | Cancel !BlockIx
50 50
51 | Port PortNumber 51 | Port !PortNumber
52 52
53 -- | BEP 6: Then peer have all pieces it might send the 53 -- | BEP 6: Then peer have all pieces it might send the
54 -- 'HaveAll' message instead of 'Bitfield' 54 -- 'HaveAll' message instead of 'Bitfield'
@@ -63,16 +63,16 @@ data Message = KeepAlive
63 -- | BEP 6: This is an advisory message meaning "you 63 -- | BEP 6: This is an advisory message meaning "you
64 -- might like to download this piece." Used to avoid 64 -- might like to download this piece." Used to avoid
65 -- excessive disk seeks and amount of IO. 65 -- excessive disk seeks and amount of IO.
66 | SuggestPiece PieceIx 66 | SuggestPiece !PieceIx
67 67
68 -- | BEP 6: Notifies a requesting peer that its request 68 -- | BEP 6: Notifies a requesting peer that its request
69 -- will not be satisfied. 69 -- will not be satisfied.
70 | RejectRequest BlockIx 70 | RejectRequest !BlockIx
71 71
72 -- | BEP 6: This is an advisory messsage meaning "if 72 -- | BEP 6: This is an advisory messsage meaning "if
73 -- you ask for this piece, I'll give it to you even if 73 -- you ask for this piece, I'll give it to you even if
74 -- you're choked." Used to shorten starting phase. 74 -- you're choked." Used to shorten starting phase.
75 | AllowedFast PieceIx 75 | AllowedFast !PieceIx
76 deriving (Show, Eq) 76 deriving (Show, Eq)
77 77
78 78
diff --git a/src/Network/BitTorrent/PeerWire/Selection.hs b/src/Network/BitTorrent/PeerWire/Selection.hs
index 9d154613..63cca15d 100644
--- a/src/Network/BitTorrent/PeerWire/Selection.hs
+++ b/src/Network/BitTorrent/PeerWire/Selection.hs
@@ -33,6 +33,7 @@ module Network.BitTorrent.PeerWire.Selection
33 ) where 33 ) where
34 34
35import Data.Bitfield 35import Data.Bitfield
36import Data.Ratio
36import Network.BitTorrent.PeerWire.Block 37import Network.BitTorrent.PeerWire.Block
37 38
38 39
@@ -42,10 +43,8 @@ type Selector = Bitfield -- ^ Indices of client /have/ pieces.
42 -> Maybe PieceIx -- ^ Zero-based index of piece to request 43 -> Maybe PieceIx -- ^ Zero-based index of piece to request
43 -- to, if any. 44 -- to, if any.
44 45
45type PieceThreshold = Int
46
47selector :: Selector -- ^ Selector to use at the start. 46selector :: Selector -- ^ Selector to use at the start.
48 -> PieceThreshold 47 -> Ratio PieceCount
49 -> Selector -- ^ Selector to use after the client have the C pieces. 48 -> Selector -- ^ Selector to use after the client have the C pieces.
50 -> Selector -- ^ Selector that changes behaviour based on completeness. 49 -> Selector -- ^ Selector that changes behaviour based on completeness.
51selector start pt ready h a xs = 50selector start pt ready h a xs =
@@ -60,16 +59,14 @@ data StartegyClass
60 | SCEnd 59 | SCEnd
61 deriving (Show, Eq, Ord, Enum, Bounded) 60 deriving (Show, Eq, Ord, Enum, Bounded)
62 61
63endThreshold :: PieceThreshold
64endThreshold = 1
65 62
66strategyClass :: PieceThreshold -> Bitfield -> StartegyClass 63strategyClass :: Ratio PieceCount -> Bitfield -> StartegyClass
67strategyClass pt = classify . completeness 64strategyClass threshold = classify . completeness
68 where 65 where
69 classify (have, total) 66 classify have
70 | have < pt = SCBeginning 67 | have < threshold = SCBeginning
71 | total - have > endThreshold = SCReady 68 | have + 1 % numerator have < 1 = SCReady -- FIXME numerator have is not total count
72 | otherwise = SCEnd 69 | otherwise = SCEnd
73 70
74 71
75-- | Select the first available piece. 72-- | Select the first available piece.
@@ -82,10 +79,9 @@ strictLast h a _ = findMax (difference a h)
82 79
83-- | 80-- |
84rarestFirst :: Selector 81rarestFirst :: Selector
85rarestFirst h a xs = rarest (frequencies (map (intersection want) xs)) 82rarestFirst h a xs = rarest (map (intersection want) xs)
86 where 83 where
87 want = difference h a 84 want = difference h a
88 rarest = Just . head
89 85
90-- | In average random first is faster than rarest first strategy but 86-- | In average random first is faster than rarest first strategy but
91-- only if all pieces are available. 87-- only if all pieces are available.
diff --git a/tests/Encoding.hs b/tests/Encoding.hs
index d43045bc..0b678a25 100644
--- a/tests/Encoding.hs
+++ b/tests/Encoding.hs
@@ -12,6 +12,7 @@ import Test.Framework (Test)
12import Test.Framework.Providers.QuickCheck2 (testProperty) 12import Test.Framework.Providers.QuickCheck2 (testProperty)
13import Test.QuickCheck 13import Test.QuickCheck
14 14
15
15import Network.URI 16import Network.URI
16import Network 17import Network
17 18
@@ -33,7 +34,8 @@ instance Arbitrary BlockIx where
33instance Arbitrary Block where 34instance Arbitrary Block where
34 arbitrary = Block <$> positive <*> positive <*> arbitrary 35 arbitrary = Block <$> positive <*> positive <*> arbitrary
35 36
36deriving instance Arbitrary Bitfield 37instance Arbitrary Bitfield where
38 arbitrary = mkBitfield <$> positive <*> arbitrary
37 39
38instance Arbitrary PortNumber where 40instance Arbitrary PortNumber where
39 arbitrary = fromIntegral <$> (arbitrary :: Gen Word16) 41 arbitrary = fromIntegral <$> (arbitrary :: Gen Word16)
diff --git a/tests/Main.hs b/tests/Main.hs
index 9b0d58e4..bc3f7809 100644
--- a/tests/Main.hs
+++ b/tests/Main.hs
@@ -10,121 +10,5 @@ import Encoding
10import Data.Bitfield as BT 10import Data.Bitfield as BT
11import Network.BitTorrent as BT 11import Network.BitTorrent as BT
12 12
13
14prop_bitfieldDiff0 :: Bitfield -> Bool
15prop_bitfieldDiff0 b = (b `difference` empty (8 * bitfieldByteCount b)) == b
16
17prop_bitfieldDiff1 :: Bitfield -> Bool
18prop_bitfieldDiff1 b = em `difference` b == em
19 where
20 em = empty (8 * bitfieldByteCount b)
21
22prop_bitfieldMaxNothing :: Int -> Bool
23prop_bitfieldMaxNothing n = findMax (empty (n `mod` 1024)) == Nothing
24
25prop_bitfieldMinNothing :: Int -> Bool
26prop_bitfieldMinNothing n = findMax (empty (n `mod` 1024)) == Nothing
27
28prop_bitfieldMaxJust :: Word -> Bool
29prop_bitfieldMaxJust n =
30 let m = findMax (full (8 * s)) in
31 if n == 0 then m == Nothing
32 else m == Just ((s * 8) - 1)
33 where
34 s = fromIntegral n `mod` 1024
35
36prop_bitfieldMinCases :: Bool
37prop_bitfieldMinCases = all mkTestCase
38 [ ("\x0\x3", Just 8)
39 , ("\x0\x127", Just 8)
40 ]
41 where
42 mkTestCase (bs, res) = findMin (MkBitfield bs) == res
43
44prop_bitfieldMaxCases :: Bool
45prop_bitfieldMaxCases = all mkTestCase
46 [ ("\x0\x3", Just 9)
47 , ("\x0\x127", Just 13)
48 ]
49 where
50 mkTestCase (bs, res) = findMax (MkBitfield bs) == res
51
52prop_bitfieldMinJust :: Word -> Bool
53prop_bitfieldMinJust n =
54 let m = findMin (full (fromIntegral n `mod` 1024)) in
55 if n == 0 then m == Nothing
56 else m == Just 0
57
58prop_bitfieldUnionIdentity :: Bitfield -> Bool
59prop_bitfieldUnionIdentity b =
60 ((b `union` empty (8 * bitfieldByteCount b)) == b)
61 && ((empty (8 * bitfieldByteCount b) `union` b) == b)
62
63prop_bitfieldUnionCommutative :: Bitfield -> Bitfield -> Bool
64prop_bitfieldUnionCommutative a b = union a b == union b a
65
66prop_bitfieldUnionAssociative :: Bitfield -> Bitfield -> Bitfield -> Bool
67prop_bitfieldUnionAssociative a b c = union a (union b c) == union (union a b) c
68
69prop_bitfieldUnionIdempotent :: Bitfield -> Bitfield -> Bool
70prop_bitfieldUnionIdempotent a b = union a b == union a (union a b)
71
72prop_bitfieldIntersectionIdentity :: Bitfield -> Bool
73prop_bitfieldIntersectionIdentity b =
74 ((b `intersection` full (8 * bitfieldByteCount b)) == b)
75 && ((full (8 * bitfieldByteCount b) `intersection` b) == b)
76
77prop_bitfieldIntersectionCommutative :: Bitfield -> Bitfield -> Bool
78prop_bitfieldIntersectionCommutative a b = intersection a b == intersection b a
79
80prop_bitfieldIntersectionAssociative :: Bitfield -> Bitfield -> Bitfield -> Bool
81prop_bitfieldIntersectionAssociative a b c =
82 intersection a (intersection b c) == intersection (intersection a b) c
83
84prop_bitfieldIntersectionIndempotent :: Bitfield -> Bitfield -> Bool
85prop_bitfieldIntersectionIndempotent a b = f b == f (f b)
86 where
87 f = intersection a
88
89prop_bitfieldHaveCount :: Bitfield -> Bool
90prop_bitfieldHaveCount b = listHaveCount (toList b) == haveCount b
91 where
92 listHaveCount = foldr f 0
93
94 f :: Bool -> Int -> Int
95 f byte count = fromEnum byte + count
96
97prop_bitfieldCompeteness :: Bitfield -> Bool
98prop_bitfieldCompeteness b = let (have, total) = completeness b in have <= total
99
100main :: IO () 13main :: IO ()
101main = defaultMain $ 14main = defaultMain []
102 [ testProperty "Message encode <-> decode" $ prop_encoding (T :: T Message)
103 , testProperty "PeerID encode <-> decode" $ prop_encoding (T :: T PeerID)
104 , testProperty "Handshake encode <-> decode" $ prop_encoding (T :: T Handshake)
105 ]
106 ++ test_scrape_url ++
107 [
108 testProperty "bitfield `difference` empty bitfield" prop_bitfieldDiff0
109 , testProperty "empty bitfield `difference` bitfield" prop_bitfieldDiff1
110
111 , testProperty "prop_bitfieldMinNothing" prop_bitfieldMinNothing
112 , testProperty "prop_bitfieldMaxNothing" prop_bitfieldMaxNothing
113 , testProperty "prop_bitfieldMaxJust" prop_bitfieldMaxJust
114 , testProperty "prop_bitfieldMinJust" prop_bitfieldMinJust
115 , testProperty "prop_bitfieldMinCases" prop_bitfieldMinCases
116 , testProperty "prop_bitfieldMaxCases" prop_bitfieldMaxCases
117
118 , testProperty "prop_bitfieldUnionIdentity" prop_bitfieldUnionIdentity
119 , testProperty "prop_bitfieldUnionCommutative" prop_bitfieldUnionCommutative
120 , testProperty "prop_bitfieldUnionAssociative" prop_bitfieldUnionAssociative
121 , testProperty "prop_bitfieldUnionIdempotent" prop_bitfieldUnionIdempotent
122
123 , testProperty "prop_bitfieldIntersectionIdentity" prop_bitfieldIntersectionIdentity
124 , testProperty "prop_bitfieldIntersectionCommutative" prop_bitfieldIntersectionCommutative
125 , testProperty "prop_bitfieldIntersectionAssociative" prop_bitfieldIntersectionAssociative
126 , testProperty "prop_bitfieldIntersectionIndempotent" prop_bitfieldIntersectionIndempotent
127
128 , testProperty "prop_bitfieldHaveCount" prop_bitfieldHaveCount
129 , testProperty "prop_bitfieldCompeteness" prop_bitfieldCompeteness
130 ] \ No newline at end of file
generated by cgit v1.2.3 (git 2.39.1) at 2024-09-23 07:32:59 +0000