summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--src/Network/BitTorrent/PeerWire/Bitfield.hs60
1 files changed, 39 insertions, 21 deletions
diff --git a/src/Network/BitTorrent/PeerWire/Bitfield.hs b/src/Network/BitTorrent/PeerWire/Bitfield.hs
index 5cbd4e86..2d2bbd59 100644
--- a/src/Network/BitTorrent/PeerWire/Bitfield.hs
+++ b/src/Network/BitTorrent/PeerWire/Bitfield.hs
@@ -18,7 +18,9 @@ module Network.BitTorrent.PeerWire.Bitfield
18 , fromByteString, toByteString 18 , fromByteString, toByteString
19 19
20 -- * Query 20 -- * Query
21 , findMin, findMax, difference 21 , findMin, findMax
22 , union, intersection, difference
23 , frequencies
22 24
23 -- * Serialization 25 -- * Serialization
24 , getBitfield, putBitfield, bitfieldByteCount 26 , getBitfield, putBitfield, bitfieldByteCount
@@ -29,7 +31,7 @@ import Data.Array.Unboxed
29import Data.Bits 31import Data.Bits
30import Data.ByteString (ByteString) 32import Data.ByteString (ByteString)
31import qualified Data.ByteString as B 33import qualified Data.ByteString as B
32import Data.List as L 34import Data.List as L hiding (union)
33import Data.Maybe 35import Data.Maybe
34import Data.Serialize 36import Data.Serialize
35import Data.Word 37import Data.Word
@@ -45,9 +47,11 @@ newtype Bitfield = MkBitfield {
45 47
46empty :: Int -> Bitfield 48empty :: Int -> Bitfield
47empty n = MkBitfield $ B.replicate (sizeInBase n 8) 0 49empty n = MkBitfield $ B.replicate (sizeInBase n 8) 0
50{-# INLINE empty #-}
48 51
49full :: Int -> Bitfield 52full :: Int -> Bitfield
50full n = MkBitfield $ B.replicate (sizeInBase n 8) (complement 0) 53full n = MkBitfield $ B.replicate (sizeInBase n 8) (complement 0)
54{-# INLINE full #-}
51 55
52fromByteString :: ByteString -> Bitfield 56fromByteString :: ByteString -> Bitfield
53fromByteString = MkBitfield 57fromByteString = MkBitfield
@@ -67,38 +71,52 @@ combine as@(a : _) = return $ foldr andBS empty as
67frequencies :: [Bitfield] -> UArray PieceIx Int 71frequencies :: [Bitfield] -> UArray PieceIx Int
68frequencies = undefined 72frequencies = undefined
69 73
70diffWord8 :: Word8 -> Word8 -> Word8 74zipWithBF :: (Word8 -> Word8 -> Word8) -> Bitfield -> Bitfield -> Bitfield
71diffWord8 a b = a .&. (a `xor` b) 75zipWithBF f a b = MkBitfield $ B.pack $ B.zipWith f (bfBits a) (bfBits b)
72{-# INLINE diffWord8 #-} 76{-# INLINE zipWithBF #-}
73 77
74difference :: Bitfield -> Bitfield -> Bitfield 78union :: Bitfield -> Bitfield -> Bitfield
75difference a b = MkBitfield $ B.pack $ B.zipWith diffWord8 (bfBits a) (bfBits b) 79union = zipWithBF (.|.)
76{-# INLINE difference #-}
77 80
78difference' :: ByteString -> ByteString -> ByteString 81intersection :: Bitfield -> Bitfield -> Bitfield
79difference' a b = undefined 82intersection = zipWithBF (.&.)
83
84difference :: Bitfield -> Bitfield -> Bitfield
85difference = zipWithBF diffWord8
80 where 86 where
81 go i = undefined 87 diffWord8 :: Word8 -> Word8 -> Word8
88 diffWord8 a b = a .&. (a `xor` b)
89 {-# INLINE diffWord8 #-}
90{-# INLINE difference #-}
82 91
83 92
84-- TODO: bit tricks
85findMinWord8 :: Word8 -> Maybe Int
86findMinWord8 b = L.findIndex (testBit b) [0..bitSize (undefined :: Word8) - 1]
87{-# INLINE findMinWord8 #-}
88 93
89-- | Get min index of piece that the peer have. 94-- | Get min index of piece that the peer have.
90findMin :: Bitfield -> Maybe PieceIx 95findMin :: Bitfield -> Maybe PieceIx
91findMin (MkBitfield b) = do 96findMin (MkBitfield b) = do
92 byteIx <- B.findIndex (0 /=) b 97 byteIx <- B.findIndex (0 /=) b
93 bitIx <- findMinWord8 (B.index b byteIx) 98 bitIx <- findMinWord8 (B.index b byteIx)
94 return $ byteIx * bitSize (undefined :: Word8) + bitIx 99 return $ byteIx * bitSize (undefined :: Word8) + bitIx
100 where
101 -- TODO: bit tricks
102 findMinWord8 :: Word8 -> Maybe Int
103 findMinWord8 b = L.find (testBit b) [0..bitSize (undefined :: Word8) - 1]
104 {-# INLINE findMinWord8 #-}
95{-# INLINE findMin #-} 105{-# INLINE findMin #-}
96 106
97findMaxWord8 :: Word8 -> Maybe Int
98findMaxWord8 = error "bitfield: findMaxWord8"
99 107
100findMax :: Bitfield -> Maybe PieceIx 108findMax :: Bitfield -> Maybe PieceIx
101findMax = error "bitfield: findMax" 109findMax (MkBitfield b) = do
110 byteIx <- (pred (B.length b) -) <$> B.findIndex (0 /=) (B.reverse b)
111 bitIx <- findMaxWord8 (B.index b byteIx)
112 return $ byteIx * bitSize (undefined :: Word8) + bitIx
113 where
114 -- TODO: bit tricks
115 findMaxWord8 :: Word8 -> Maybe Int
116 findMaxWord8 b = L.find (testBit b)
117 (reverse [0 :: Int ..
118 bitSize (undefined :: Word8) - 1])
119
102{-# INLINE findMax #-} 120{-# INLINE findMax #-}
103 121
104 122