4 files changed, 80 insertions, 5 deletions

diff --git a/bench/Main.hs b/bench/Main.hs
index 6662b953..6162eeb1 100644
--- a/bench/Main.hs
+++ b/bench/Main.hs
@@ -49,6 +49,9 @@ bitfieldInter n = BT.empty n `intersection` BT.empty n
 bitfieldUnion :: Int -> Bitfield
 bitfieldUnion n = BT.empty n `union` BT.empty n
 
+bitfieldHaveCount :: Int -> Int
+bitfieldHaveCount n = haveCount (BT.full n)
+
 selectionStrictFirst :: Int -> Maybe Int
 selectionStrictFirst n = strictFirst (BT.empty n) (BT.empty n) []
 
@@ -84,6 +87,7 @@ main = do
     , bench "bitfield/difference"   $ nf bitfieldDiff  (10 * m)
     , bench "bitfield/intersection" $ nf bitfieldInter (10 * m)
     , bench "bitfield/union"        $ nf bitfieldUnion (10 * m)
+    , bench "bitfield/haveCount"    $ nf bitfieldHaveCount (10 * m)
 
     , bench "selection/strictFirst" $ nf selectionStrictFirst  (10 * m)
     , bench "selection/strictLast"  $ nf selectionStrictLast   (10 * m)
diff --git a/src/Network/BitTorrent/PeerWire/Bitfield.hs b/src/Network/BitTorrent/PeerWire/Bitfield.hs
index 9d88e784..03273899 100644
--- a/src/Network/BitTorrent/PeerWire/Bitfield.hs
+++ b/src/Network/BitTorrent/PeerWire/Bitfield.hs
@@ -22,12 +22,14 @@ module Network.BitTorrent.PeerWire.Bitfield
        , fromByteString, toByteString
 
          -- * Query
+       , haveCount, completeness
        , findMin, findMax
        , union, intersection, difference, combine
        , frequencies
 
          -- * Serialization
-       , getBitfield, putBitfield, bitfieldByteCount
+       , getBitfield, putBitfield
+       , bitfieldByteCount, bitfieldBitCount
 
 
        , aligned, alignLow, alignedZip
@@ -46,9 +48,11 @@ import Data.Word
 
 import Foreign
 
-import Network.BitTorrent.PeerWire.Block
+--import Network.BitTorrent.PeerWire.Block
 import Data.Torrent
 
+-- one good idea is to aggregate frequently used stats in reducer
+-- it should give a big boost
 newtype Bitfield = MkBitfield {
     bfBits :: ByteString
 --  , bfSize :: Int
@@ -90,6 +94,11 @@ bitfieldByteCount :: Bitfield -> Int
 bitfieldByteCount = B.length . bfBits
 {-# INLINE bitfieldByteCount #-}
 
+-- WARN
+-- TODO
+bitfieldBitCount :: Bitfield -> Int
+bitfieldBitCount bf = bitSize (undefined :: Word8) * bitfieldByteCount bf
+{-# INLINE bitfieldBitCount #-}
 
 align :: Storable a => Ptr a -> (Ptr a, Int)
 align p = tie (alignPtr p) undefined
@@ -235,7 +244,44 @@ findSet b =
                  else go (succ i)
              | otherwise = Nothing
 
+foldBS :: (Word8 -> Int -> Int) -> (Word -> Int -> Int) -> Int -> ByteString -> Int
+foldBS f g acc b =
+    let (fptr, off, len) = B.toForeignPtr b in
+    B.inlinePerformIO $ withForeignPtr fptr $ \_ptr -> do
+      let ptr = _ptr `advancePtr` off
+
+      let (low, mid, hgh) = aligned (ptr, len)
+      let resL = goFold   low acc
+      let resM = goFoldW (mid :: Mem Word) resL
+      let resH = goFold   hgh resM
+
+      -- computation of res should not escape withForeignPtr
+      case resH of
+        0 -> return ()
+        _ -> return ()
 
+      return resH
+
+  where
+    goFold :: Mem Word8 -> Int -> Int
+    goFold (ptr, n) = go 0
+      where
+        go :: Int -> Int -> Int
+        go i !a
+          |   i < n   =
+            let v = B.inlinePerformIO (peekElemOff ptr i)
+            in go (succ i) (f v a)
+          | otherwise = a
+
+    goFoldW :: Mem Word -> Int -> Int
+    goFoldW (ptr, n) = go 0
+      where
+        go :: Int -> Int -> Int
+        go i !a
+          |   i < n   =
+            let v = B.inlinePerformIO (peekElemOff ptr i)
+            in go (succ i) (g v a)
+          | otherwise = a
 
 union :: Bitfield -> Bitfield -> Bitfield
 union = zipWithBF (.|.)
@@ -257,8 +303,16 @@ combine :: [Bitfield] -> Maybe Bitfield
 combine [] = Nothing
 combine as = return $ foldr1 intersection as
 
+haveCount :: Bitfield -> Int
+haveCount (MkBitfield b) = foldBS f f 0 b
+  where
+    f byte count = popCount byte + count
+
+completeness :: Bitfield -> (Int, Int)
+completeness bf = (haveCount bf, bitfieldBitCount bf)
+
 -- | Get min index of piece that the peer have.
-findMin :: Bitfield -> Maybe PieceIx
+findMin :: Bitfield -> Maybe Int
 findMin (MkBitfield b) = do
     byteIx <- findSet b
     bitIx  <- findMinWord8 (B.index b byteIx)
@@ -271,7 +325,7 @@ findMin (MkBitfield b) = do
 {-# INLINE findMin #-}
 
 
-findMax :: Bitfield -> Maybe PieceIx
+findMax :: Bitfield -> Maybe Int
 findMax (MkBitfield b) = do
     -- TODO avoid reverse
     byteIx <- (pred (B.length b) -) <$> findSet (B.reverse b)
diff --git a/src/Network/BitTorrent/PeerWire/Selection.hs b/src/Network/BitTorrent/PeerWire/Selection.hs
index 92285501..2e17317e 100644
--- a/src/Network/BitTorrent/PeerWire/Selection.hs
+++ b/src/Network/BitTorrent/PeerWire/Selection.hs
@@ -24,7 +24,9 @@
 module Network.BitTorrent.PeerWire.Selection
        ( Selector
        , strictFirst, strictLast
-       , rarestFirst, randomFirst, endGame, autoSelector
+       , rarestFirst, randomFirst, endGame
+
+       , autoSelector
        ) where
 
 import Network.BitTorrent.PeerWire.Block
diff --git a/tests/Main.hs b/tests/Main.hs
index 4045f793..5a9a50ee 100644
--- a/tests/Main.hs
+++ b/tests/Main.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE OverloadedStrings #-}
 module Main (main) where
 
+import Data.Bits
 import Data.Word
 import Test.Framework (defaultMain)
 import Test.Framework.Providers.QuickCheck2 (testProperty)
@@ -84,6 +85,17 @@ prop_bitfieldIntersectionIndempotent a b = f b == f (f b)
   where
     f = intersection a
 
+prop_bitfieldHaveCount :: Bitfield -> Bool
+prop_bitfieldHaveCount b = listHaveCount (toList b) == haveCount b
+  where
+    listHaveCount = foldr f 0
+
+    f :: Bool -> Int -> Int
+    f byte count = fromEnum byte + count
+
+prop_bitfieldCompeteness :: Bitfield -> Bool
+prop_bitfieldCompeteness b = let (have, total) = completeness b in have <= total
+
 main :: IO ()
 main = defaultMain $
        [ testProperty "Message encode <-> decode"   $ prop_encoding (T :: T Message)
@@ -111,4 +123,7 @@ main = defaultMain $
        , testProperty "prop_bitfieldIntersectionCommutative" prop_bitfieldIntersectionCommutative
        , testProperty "prop_bitfieldIntersectionAssociative" prop_bitfieldIntersectionAssociative
        , testProperty "prop_bitfieldIntersectionIndempotent" prop_bitfieldIntersectionIndempotent
+
+       , testProperty "prop_bitfieldHaveCount"               prop_bitfieldHaveCount
+       , testProperty "prop_bitfieldCompeteness"             prop_bitfieldCompeteness
        ]
\ No newline at end of file