summaryrefslogtreecommitdiff
path: root/src/Network/BitTorrent/PeerWire
diff options
context:
space:
mode:
authorSam T <pxqr.sta@gmail.com>2013-05-05 05:01:28 +0400
committerSam T <pxqr.sta@gmail.com>2013-05-05 05:01:28 +0400
commit7f63b5554fcf31d9b71c79fa9d0ee66d2ef95c94 (patch)
tree008058dc9bfe4b6f75cf1add4f493235a7ee8ed1 /src/Network/BitTorrent/PeerWire
parentf4122eec550671a646310106224ee6523ea8e369 (diff)
~ Move bitfields to Data.
Diffstat (limited to 'src/Network/BitTorrent/PeerWire')
-rw-r--r--src/Network/BitTorrent/PeerWire/Bitfield.hs344
-rw-r--r--src/Network/BitTorrent/PeerWire/Message.hs2
-rw-r--r--src/Network/BitTorrent/PeerWire/Selection.hs3
3 files changed, 3 insertions, 346 deletions
diff --git a/src/Network/BitTorrent/PeerWire/Bitfield.hs b/src/Network/BitTorrent/PeerWire/Bitfield.hs
deleted file mode 100644
index 03273899..00000000
--- a/src/Network/BitTorrent/PeerWire/Bitfield.hs
+++ /dev/null
@@ -1,344 +0,0 @@
1-- |
2-- Copyright : (c) Sam T. 2013
3-- License : MIT
4-- Maintainer : pxqr.sta@gmail.com
5-- Stability : experimental
6-- Portability : portable
7--
8--
9-- This module provides Bitfield datatype used to represent sets of
10-- piece indexes any peer have. All associated operations should be
11-- defined here as well.
12--
13{-# LANGUAGE BangPatterns #-}
14{-# LANGUAGE RankNTypes #-}
15module Network.BitTorrent.PeerWire.Bitfield
16-- TODO: move to Data.Bitfield
17 ( Bitfield(..)
18
19 -- * Construction
20 , empty, full
21 , toList
22 , fromByteString, toByteString
23
24 -- * Query
25 , haveCount, completeness
26 , findMin, findMax
27 , union, intersection, difference, combine
28 , frequencies
29
30 -- * Serialization
31 , getBitfield, putBitfield
32 , bitfieldByteCount, bitfieldBitCount
33
34
35 , aligned, alignLow, alignedZip
36 ) where
37
38import Control.Applicative hiding (empty)
39import Data.Array.Unboxed
40import Data.Bits
41import Data.ByteString (ByteString)
42import qualified Data.ByteString as B
43import qualified Data.ByteString.Internal as B
44import Data.List as L hiding (union)
45import Data.Maybe
46import Data.Serialize
47import Data.Word
48
49import Foreign
50
51--import Network.BitTorrent.PeerWire.Block
52import Data.Torrent
53
54-- one good idea is to aggregate frequently used stats in reducer
55-- it should give a big boost
56newtype Bitfield = MkBitfield {
57 bfBits :: ByteString
58-- , bfSize :: Int
59 } deriving (Show, Eq, Ord)
60
61
62empty :: Int -> Bitfield
63empty n = MkBitfield $ B.replicate (sizeInBase n 8) 0
64{-# INLINE empty #-}
65
66full :: Int -> Bitfield
67full n = MkBitfield $ B.replicate (sizeInBase n 8) (complement 0)
68{-# INLINE full #-}
69
70toList :: Bitfield -> [Bool]
71toList (MkBitfield bs) = concatMap unpkg (B.unpack bs)
72 where
73 unpkg :: Word8 -> [Bool]
74 unpkg byte = L.map (testBit byte) [0..bitSize (undefined :: Word8) - 1]
75{-# INLINE toList #-}
76
77fromByteString :: ByteString -> Bitfield
78fromByteString = MkBitfield
79{-# INLINE fromByteString #-}
80
81toByteString :: Bitfield -> ByteString
82toByteString = bfBits
83{-# INLINE toByteString #-}
84
85getBitfield :: Int -> Get Bitfield
86getBitfield n = MkBitfield <$> getBytes n
87{-# INLINE getBitfield #-}
88
89putBitfield :: Bitfield -> Put
90putBitfield = putByteString . bfBits
91{-# INLINE putBitfield #-}
92
93bitfieldByteCount :: Bitfield -> Int
94bitfieldByteCount = B.length . bfBits
95{-# INLINE bitfieldByteCount #-}
96
97-- WARN
98-- TODO
99bitfieldBitCount :: Bitfield -> Int
100bitfieldBitCount bf = bitSize (undefined :: Word8) * bitfieldByteCount bf
101{-# INLINE bitfieldBitCount #-}
102
103align :: Storable a => Ptr a -> (Ptr a, Int)
104align p = tie (alignPtr p) undefined
105 where
106 tie :: Storable a => (Int -> Ptr a) -> a -> (Ptr a, Int)
107 tie f a = (f (alignment a), (alignment a))
108
109alignLow :: Ptr Word8 -> Ptr Word
110alignLow ptr =
111 let alg = alignment (undefined :: Word)
112 aptr = alignPtr (castPtr ptr) alg :: Ptr Word
113 in
114 if ptr == castPtr aptr
115 then aptr
116 else castPtr ((castPtr aptr :: Ptr Word8) `advancePtr` negate alg)
117
118isAlignedBy :: Storable a => Ptr a -> Int -> Bool
119isAlignedBy ptr alg = alignPtr ptr alg == ptr
120
121type Mem a = (Ptr a, Int)
122
123aligned :: Storable a => Mem Word8 -> (Mem Word8, Mem a, Mem Word8)
124aligned (ptr, len) =
125 let lowPtr = ptr
126 lowLen = midPtr `minusPtr` ptr
127 midOff = lowLen
128 (midPtr, alg) = align (castPtr ptr)
129 midLen = alg * div (len - midOff) alg
130 midLenA = midLen `div` alg
131 hghOff = midOff + midLen
132 hghPtr = ptr `advancePtr` hghOff
133 hghLen = len - hghOff
134 in
135 ((lowPtr, lowLen), (midPtr, midLenA), (hghPtr, hghLen))
136 where
137{-# INLINE aligned #-}
138
139type Mem3 a = (Ptr a, Ptr a, Ptr a, Int)
140
141emptyMem3 :: Mem3 a
142emptyMem3 = (nullPtr, nullPtr, nullPtr, 0)
143
144-- assume resulting memory is aligned
145alignedZip :: Ptr Word8 -> Ptr Word8 -> Ptr Word8 -> Int
146 -> (Mem3 Word, Mem3 Word8)
147alignedZip aptr bptr cptr size =
148 let alg = alignment (undefined :: Word) in
149 if (aptr `isAlignedBy` alg) && (bptr `isAlignedBy` alg)
150 then
151 let asize = alignLow (nullPtr `plusPtr` size) `minusPtr` nullPtr
152 in
153 ( (castPtr aptr, castPtr bptr, castPtr cptr, asize `div` alg)
154 , ( aptr `advancePtr` asize
155 , bptr `advancePtr` asize
156 , cptr `advancePtr` asize
157 , (size - asize)
158 )
159 )
160 else (emptyMem3, (aptr, bptr, cptr, size))
161
162-- force specialization
163zipWithBS :: (Word -> Word -> Word)
164 -> (Word8 -> Word8 -> Word8)
165 -> ByteString -> ByteString -> ByteString
166zipWithBS f g a b =
167 let (afptr, aoff, asize) = B.toForeignPtr a
168 (bfptr, boff, bsize) = B.toForeignPtr b
169 size = min asize bsize in
170 B.unsafeCreate size $ \rptr -> do
171 withForeignPtr afptr $ \_aptr -> do
172 withForeignPtr bfptr $ \_bptr -> do
173 let aptr = _aptr `advancePtr` aoff
174 let bptr = _bptr `advancePtr` boff
175
176 let (mid, hgh) = alignedZip aptr bptr rptr size
177 zipWords mid
178 zipBytes hgh
179 where
180 zipBytes :: (Ptr Word8, Ptr Word8, Ptr Word8, Int) -> IO ()
181 zipBytes (aptr, bptr, rptr, n) = go 0
182 where
183 go :: Int -> IO ()
184 go i | i < n = do -- TODO unfold
185 av <- peekElemOff aptr i
186 bv <- peekElemOff bptr i
187 pokeElemOff rptr i (g av bv)
188 go (succ i)
189 | otherwise = return ()
190
191 zipWords :: (Ptr Word, Ptr Word, Ptr Word, Int) -> IO ()
192 zipWords (aptr, bptr, rptr, n) = go 0
193 where
194 go :: Int -> IO ()
195 go i | i < n = do -- TODO unfold
196 av <- peekElemOff aptr i
197 bv <- peekElemOff bptr i
198 pokeElemOff rptr i (f av bv)
199 go (succ i)
200 | otherwise = return ()
201
202
203
204zipWithBF :: (forall a. Bits a => a -> a -> a) -> Bitfield -> Bitfield -> Bitfield
205zipWithBF f a b = MkBitfield $ zipWithBS f f (bfBits a) (bfBits b)
206{-# INLINE zipWithBF #-}
207
208findSet :: ByteString -> Maybe Int
209findSet b =
210 let (fptr, off, len) = B.toForeignPtr b in
211 B.inlinePerformIO $ withForeignPtr fptr $ \_ptr -> do
212 let ptr = _ptr `advancePtr` off
213
214 let (low, mid, hgh) = aligned (ptr, len)
215 let lowOff = fst low `minusPtr` ptr
216 let midOff = fst mid `minusPtr` ptr
217 let hghOff = fst hgh `minusPtr` ptr
218
219 let resL = (lowOff +) <$> goFind low
220 let resM = (midOff +) <$> goFind (mid :: Mem Word) -- tune size here
221 -- TODO: with Word8
222 -- bytestring findIndex works 2
223 -- times faster.
224 let resH = (hghOff +) <$> goFind hgh
225
226 let res = resL <|> resM <|> resH
227
228 -- computation of res should not escape withForeignPtr
229 case res of
230 Nothing -> return ()
231 Just _ -> return ()
232
233 return res
234
235 where
236 goFind :: (Storable a, Eq a, Num a) => Mem a -> Maybe Int
237 goFind (ptr, n) = go 0
238 where
239 go :: Int -> Maybe Int
240 go i | i < n =
241 let v = B.inlinePerformIO (peekElemOff ptr i) in
242 if v /= 0
243 then Just i
244 else go (succ i)
245 | otherwise = Nothing
246
247foldBS :: (Word8 -> Int -> Int) -> (Word -> Int -> Int) -> Int -> ByteString -> Int
248foldBS f g acc b =
249 let (fptr, off, len) = B.toForeignPtr b in
250 B.inlinePerformIO $ withForeignPtr fptr $ \_ptr -> do
251 let ptr = _ptr `advancePtr` off
252
253 let (low, mid, hgh) = aligned (ptr, len)
254 let resL = goFold low acc
255 let resM = goFoldW (mid :: Mem Word) resL
256 let resH = goFold hgh resM
257
258 -- computation of res should not escape withForeignPtr
259 case resH of
260 0 -> return ()
261 _ -> return ()
262
263 return resH
264
265 where
266 goFold :: Mem Word8 -> Int -> Int
267 goFold (ptr, n) = go 0
268 where
269 go :: Int -> Int -> Int
270 go i !a
271 | i < n =
272 let v = B.inlinePerformIO (peekElemOff ptr i)
273 in go (succ i) (f v a)
274 | otherwise = a
275
276 goFoldW :: Mem Word -> Int -> Int
277 goFoldW (ptr, n) = go 0
278 where
279 go :: Int -> Int -> Int
280 go i !a
281 | i < n =
282 let v = B.inlinePerformIO (peekElemOff ptr i)
283 in go (succ i) (g v a)
284 | otherwise = a
285
286union :: Bitfield -> Bitfield -> Bitfield
287union = zipWithBF (.|.)
288{-# INLINE union #-}
289
290intersection :: Bitfield -> Bitfield -> Bitfield
291intersection = zipWithBF (.&.)
292{-# INLINE intersection #-}
293
294difference :: Bitfield -> Bitfield -> Bitfield
295difference = zipWithBF diffWord8
296 where
297 diffWord8 :: Bits a => a -> a -> a
298 diffWord8 a b = a .&. (a `xor` b)
299 {-# INLINE diffWord8 #-}
300{-# INLINE difference #-}
301
302combine :: [Bitfield] -> Maybe Bitfield
303combine [] = Nothing
304combine as = return $ foldr1 intersection as
305
306haveCount :: Bitfield -> Int
307haveCount (MkBitfield b) = foldBS f f 0 b
308 where
309 f byte count = popCount byte + count
310
311completeness :: Bitfield -> (Int, Int)
312completeness bf = (haveCount bf, bitfieldBitCount bf)
313
314-- | Get min index of piece that the peer have.
315findMin :: Bitfield -> Maybe Int
316findMin (MkBitfield b) = do
317 byteIx <- findSet b
318 bitIx <- findMinWord8 (B.index b byteIx)
319 return $ byteIx * bitSize (undefined :: Word8) + bitIx
320 where
321 -- TODO: bit tricks
322 findMinWord8 :: Word8 -> Maybe Int
323 findMinWord8 byte = L.find (testBit byte) [0..bitSize (undefined :: Word8) - 1]
324 {-# INLINE findMinWord8 #-}
325{-# INLINE findMin #-}
326
327
328findMax :: Bitfield -> Maybe Int
329findMax (MkBitfield b) = do
330 -- TODO avoid reverse
331 byteIx <- (pred (B.length b) -) <$> findSet (B.reverse b)
332 bitIx <- findMaxWord8 (B.index b byteIx)
333 return $ byteIx * bitSize (undefined :: Word8) + bitIx
334 where
335 -- TODO: bit tricks
336 findMaxWord8 :: Word8 -> Maybe Int
337 findMaxWord8 byte = L.find (testBit byte)
338 (reverse [0 :: Int ..
339 bitSize (undefined :: Word8) - 1])
340
341{-# INLINE findMax #-}
342
343frequencies :: [Bitfield] -> [Int]
344frequencies xs = foldr1 (zipWith (+)) $ map (map fromEnum . toList) xs
diff --git a/src/Network/BitTorrent/PeerWire/Message.hs b/src/Network/BitTorrent/PeerWire/Message.hs
index cc771966..39102eed 100644
--- a/src/Network/BitTorrent/PeerWire/Message.hs
+++ b/src/Network/BitTorrent/PeerWire/Message.hs
@@ -10,7 +10,7 @@ import qualified Data.ByteString as B
10import Data.Serialize 10import Data.Serialize
11 11
12import Network.BitTorrent.PeerWire.Block 12import Network.BitTorrent.PeerWire.Block
13import Network.BitTorrent.PeerWire.Bitfield 13import Data.Bitfield
14 14
15import Data.Array 15import Data.Array
16 16
diff --git a/src/Network/BitTorrent/PeerWire/Selection.hs b/src/Network/BitTorrent/PeerWire/Selection.hs
index 2e17317e..83ab8311 100644
--- a/src/Network/BitTorrent/PeerWire/Selection.hs
+++ b/src/Network/BitTorrent/PeerWire/Selection.hs
@@ -29,9 +29,10 @@ module Network.BitTorrent.PeerWire.Selection
29 , autoSelector 29 , autoSelector
30 ) where 30 ) where
31 31
32import Data.Bitfield
32import Network.BitTorrent.PeerWire.Block 33import Network.BitTorrent.PeerWire.Block
33import Network.BitTorrent.PeerWire.Message 34import Network.BitTorrent.PeerWire.Message
34import Network.BitTorrent.PeerWire.Bitfield 35
35 36
36 37
37type Selector = Bitfield -- ^ Indices of client "have" pieces. 38type Selector = Bitfield -- ^ Indices of client "have" pieces.
generated by cgit v1.2.3 (git 2.39.1) at 2024-11-11 19:54:50 +0000