1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
|
-- |
-- Copyright : (c) Sam T. 2013
-- License : MIT
-- Maintainer : pxqr.sta@gmail.com
-- Stability : experimental
-- Portability : portable
--
-- This modules provides all necessary machinery to work with
-- bitfields. Bitfields are used to keep track indices of complete
-- pieces either peer have or client have.
--
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RecordWildCards #-}
module Data.Bitfield
( Bitfield, PieceCount
-- * Construction
, empty
, insert
, haveAll, haveNone, have
-- * Query
, haveCount, totalCount, completeness
, findMin, findMax
, frequencies, rarest
-- * Combine
, union
, intersection
, difference
-- * Serialization
, getBitfield, putBitfield
, bitfieldByteCount
, -- * Debug
mkBitfield
) where
import Control.Monad
import Control.Monad.ST
import Data.Vector.Unboxed (Vector)
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Unboxed.Mutable as VM
import Data.IntervalSet (IntSet)
import qualified Data.IntervalSet as S
import Data.List (foldl')
import Data.Monoid
import Data.Ratio
import Data.Serialize
import Network.BitTorrent.PeerWire.Block
type PieceCount = Int
-- TODO cache some operations
-- | Bitfields are represented just as integer sets but with
-- restriction: the each set should be within given interval (or
-- subset of). Size is used to specify interval, so bitfield of size
-- 10 might contain only indices in interval [0..9].
--
data Bitfield = Bitfield {
bfSize :: !PieceCount
, bfSet :: !IntSet
} deriving (Show, Read, Eq)
-- Invariants: all elements of bfSet lie in [0..bfSize - 1];
instance Monoid Bitfield where
{-# SPECIALIZE instance Monoid Bitfield #-}
mempty = empty 0
mappend = union
mconcat = unions
-- TODO documentation
{-----------------------------------------------------------------------
Construction
-----------------------------------------------------------------------}
empty :: PieceCount -> Bitfield
empty s = Bitfield s S.empty
insert :: PieceIx -> Bitfield -> Bitfield
insert ix Bitfield {..}
| 0 <= ix && ix < bfSize = Bitfield bfSize (S.insert ix bfSet)
| otherwise = Bitfield bfSize bfSet
haveNone :: PieceCount -> Bitfield
haveNone = empty
haveAll :: PieceCount -> Bitfield
haveAll s = Bitfield s (S.interval 0 (s - 1))
have :: PieceIx -> Bitfield -> Bitfield
have = insert
{-----------------------------------------------------------------------
Query
-----------------------------------------------------------------------}
haveCount :: Bitfield -> PieceCount
haveCount = S.size . bfSet
totalCount :: Bitfield -> PieceCount
totalCount = bfSize
-- |
--
-- > forall bf. 0 <= completeness bf <= 1
--
completeness :: Bitfield -> Ratio PieceCount
completeness b = haveCount b % totalCount b
findMin :: Bitfield -> Maybe PieceIx
findMin Bitfield {..}
| S.null bfSet = Nothing
| otherwise = Just (S.findMin bfSet)
findMax :: Bitfield -> Maybe PieceIx
findMax Bitfield {..}
| S.null bfSet = Nothing
| otherwise = Just (S.findMax bfSet)
type Frequency = Int
frequencies :: [Bitfield] -> Vector Frequency
frequencies [] = V.fromList []
frequencies xs = runST $ do
v <- VM.new size
VM.set v 0
forM_ xs $ \ Bitfield {..} -> do
forM_ (S.toList bfSet) $ \ x -> do
fr <- VM.read v x
VM.write v x (succ fr)
V.unsafeFreeze v
where
size = maximum (map bfSize xs)
rarest :: [Bitfield] -> Maybe PieceIx
rarest xs
| V.null freqMap = Nothing
| otherwise = Just $ fst $ V.ifoldr minIx (0, freqMap V.! 0) freqMap
where
freqMap = frequencies xs
minIx :: PieceIx -> Frequency -> (PieceIx, Frequency) -> (PieceIx, Frequency)
minIx ix fr acc@(_, fra)
| fr < fra && fr > 0 = (ix, fr)
| otherwise = acc
{-----------------------------------------------------------------------
Combine
-----------------------------------------------------------------------}
union :: Bitfield -> Bitfield -> Bitfield
union a b = Bitfield {
bfSize = bfSize a `max` bfSize b
, bfSet = bfSet a `S.union` bfSet b
}
intersection :: Bitfield -> Bitfield -> Bitfield
intersection a b = Bitfield {
bfSize = bfSize a `min` bfSize b
, bfSet = bfSet a `S.intersection` bfSet b
}
difference :: Bitfield -> Bitfield -> Bitfield
difference a b = Bitfield {
bfSize = bfSize a -- FIXME is it more reasonable?
, bfSet = bfSet a `S.difference` bfSet b
}
unions :: [Bitfield] -> Bitfield
unions = foldl' union (empty 0)
{-----------------------------------------------------------------------
Serialization
-----------------------------------------------------------------------}
getBitfield :: Int -> Get Bitfield
getBitfield = error "getBitfield"
putBitfield :: Bitfield -> Put
putBitfield = error "putBitfield"
bitfieldByteCount :: Bitfield -> Int
bitfieldByteCount = error "bitfieldByteCount"
{-----------------------------------------------------------------------
Debug
-----------------------------------------------------------------------}
mkBitfield :: PieceCount -> [PieceIx] -> Bitfield
mkBitfield s ixs = Bitfield {
bfSize = s
, bfSet = S.splitLT s $ S.fromList ixs
}
|
