summaryrefslogtreecommitdiff
path: root/src/Network
diff options
context:
space:
mode:
Diffstat (limited to 'src/Network')
-rw-r--r--src/Network/BitTorrent/PeerWire.hs1
-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
4 files changed, 3 insertions, 347 deletions
diff --git a/src/Network/BitTorrent/PeerWire.hs b/src/Network/BitTorrent/PeerWire.hs
index 6d8ff156..d0583bff 100644
--- a/src/Network/BitTorrent/PeerWire.hs
+++ b/src/Network/BitTorrent/PeerWire.hs
@@ -9,7 +9,6 @@
9module Network.BitTorrent.PeerWire (module PW) where 9module Network.BitTorrent.PeerWire (module PW) where
10 10
11import Network.BitTorrent.PeerWire.Block as PW 11import Network.BitTorrent.PeerWire.Block as PW
12import Network.BitTorrent.PeerWire.Bitfield as PW
13import Network.BitTorrent.PeerWire.Selection as PW 12import Network.BitTorrent.PeerWire.Selection as PW
14import Network.BitTorrent.PeerWire.Message as PW 13import Network.BitTorrent.PeerWire.Message as PW
15import Network.BitTorrent.PeerWire.Handshake as PW 14import Network.BitTorrent.PeerWire.Handshake as PW
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-09-22 23:33:37 +0000