~ Use IntSet instead of ByteString for bitfields.

There are several reasons for this:

* IntSet is stored in ordinary heap, while ByteStrings in pinned memory;
* Our IntSet's should be much faster 90% time. (in typical BT client)
  Hovewer in worst case IntSet is slower, but difference should is not
  so big. (We should measure this although)
* It's pure, tested, error-free and much more convenient.

Moreover we have kill a lot of ugly code!

1 files changed, 149 insertions, 294 deletions

diff --git a/src/Data/Bitfield.hs b/src/Data/Bitfield.hs
index f7240c8a..546c68e9 100644
--- a/src/Data/Bitfield.hs
+++ b/src/Data/Bitfield.hs
@@ -5,341 +5,196 @@
 --   Stability   :  experimental
 --   Portability :  portable
 --
---
---   This module provides Bitfield datatype used to represent sets of
---   piece indexes any peer have. All associated operations should be
---   defined here as well.
+--   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 RankNTypes #-}
+{-# LANGUAGE RecordWildCards #-}
 module Data.Bitfield
--- TODO: move to Data.Bitfield
-       ( Bitfield(..)
+       ( Bitfield, PieceCount
 
          -- * Construction
-       , empty, full
+       , empty
+       , insert
+       , haveAll, haveNone, have
 
          -- * Query
-       , bitfieldByteCount, bitfieldBitCount
-
-       , haveCount, completeness
+       , haveCount, totalCount, completeness
        , findMin, findMax
-       , union, intersection, difference, combine
-       , frequencies
+       , frequencies, rarest
+
+         -- * Combine
+       , union
+       , intersection
+       , difference
 
          -- * Serialization
        , getBitfield, putBitfield
+       , bitfieldByteCount
 
-         -- * Conversion
-       , toList
-       , fromByteString, toByteString
-
-         -- * Debug
-       , aligned, alignLow, alignedZip
+       , -- * Debug
+         mkBitfield
        ) where
 
-import Control.Applicative hiding (empty)
-import Data.Bits
-import           Data.ByteString (ByteString)
-import qualified Data.ByteString as B
-import qualified Data.ByteString.Internal as B
-import Data.List as L hiding (union)
-import Data.Serialize
-import Data.Word
+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
 
-import Foreign
 
---import Network.BitTorrent.PeerWire.Block
-import Data.Torrent
+type PieceCount = Int
 
--- TODO: one good idea is to aggregate frequently used stats in reducer
--- it should give a big boost
-newtype Bitfield = MkBitfield {
-    bfBits :: ByteString
---  , bfSize :: Int
-  } deriving (Show, Eq, Ord)
+-- 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)
 
-empty :: Int -> Bitfield
-empty n = MkBitfield $ B.replicate (sizeInBase n 8) 0
-{-# INLINE empty #-}
+-- Invariants: all elements of bfSet lie in [0..bfSize - 1];
 
-full :: Int -> Bitfield
-full n = MkBitfield $ B.replicate (sizeInBase n 8)  (complement 0)
-{-# INLINE full #-}
+instance Monoid Bitfield where
+  {-# SPECIALIZE instance Monoid Bitfield #-}
+  mempty  = empty 0
+  mappend = union
+  mconcat = unions
 
-toList :: Bitfield -> [Bool]
-toList (MkBitfield bs) = concatMap unpkg (B.unpack bs)
-  where
-    unpkg :: Word8 -> [Bool]
-    unpkg byte = L.map (testBit byte) [0..bitSize (undefined :: Word8) - 1]
-{-# INLINE toList #-}
+-- TODO documentation
+{-----------------------------------------------------------------------
+    Construction
+-----------------------------------------------------------------------}
 
-fromByteString :: ByteString -> Bitfield
-fromByteString = MkBitfield
-{-# INLINE fromByteString #-}
+empty :: PieceCount -> Bitfield
+empty s = Bitfield s S.empty
 
-toByteString :: Bitfield -> ByteString
-toByteString = bfBits
-{-# INLINE toByteString #-}
+insert :: PieceIx -> Bitfield -> Bitfield
+insert ix Bitfield {..}
+  | 0 <= ix && ix < bfSize = Bitfield bfSize (S.insert ix bfSet)
+  |      otherwise         = Bitfield bfSize bfSet
 
-getBitfield :: Int -> Get Bitfield
-getBitfield n = MkBitfield <$> getBytes n
-{-# INLINE getBitfield #-}
+haveNone :: PieceCount -> Bitfield
+haveNone = empty
 
-putBitfield :: Bitfield -> Put
-putBitfield = putByteString . bfBits
-{-# INLINE putBitfield #-}
+haveAll :: PieceCount -> Bitfield
+haveAll s = Bitfield s (S.interval 0 (s - 1))
 
-bitfieldByteCount :: Bitfield -> Int
-bitfieldByteCount = B.length . bfBits
-{-# INLINE bitfieldByteCount #-}
+have :: PieceIx -> Bitfield -> Bitfield
+have = insert
 
--- WARN
--- TODO
-bitfieldBitCount :: Bitfield -> Int
-bitfieldBitCount bf = bitSize (undefined :: Word8) * bitfieldByteCount bf
-{-# INLINE bitfieldBitCount #-}
+{-----------------------------------------------------------------------
+    Query
+-----------------------------------------------------------------------}
 
-align :: Storable a => Ptr a -> (Ptr a, Int)
-align p = tie (alignPtr p) undefined
-  where
-    tie :: Storable a => (Int -> Ptr a) -> a -> (Ptr a, Int)
-    tie f a = (f (alignment a), (alignment a))
-
-alignLow :: Ptr Word8 -> Ptr Word
-alignLow ptr =
-  let alg  = alignment (undefined :: Word)
-      aptr = alignPtr (castPtr ptr) alg :: Ptr Word
-  in
-    if ptr == castPtr aptr
-    then aptr
-    else castPtr ((castPtr aptr :: Ptr Word8) `advancePtr` negate alg)
-
-isAlignedBy :: Storable a => Ptr a -> Int -> Bool
-isAlignedBy ptr alg = alignPtr ptr alg == ptr
-
-type Mem a    = (Ptr a, Int)
-
-aligned :: Storable a => Mem Word8 -> (Mem Word8, Mem a, Mem Word8)
-aligned (ptr, len) =
-  let lowPtr  = ptr
-      lowLen  = midPtr `minusPtr` ptr
-      midOff  = lowLen
-      (midPtr, alg) = align (castPtr ptr)
-      midLen  = alg * div (len - midOff) alg
-      midLenA = midLen `div` alg
-      hghOff  = midOff + midLen
-      hghPtr  = ptr `advancePtr` hghOff
-      hghLen  = len - hghOff
-   in
-     ((lowPtr, lowLen), (midPtr, midLenA), (hghPtr, hghLen))
-  where
-{-# INLINE aligned #-}
-
-type Mem3 a = (Ptr a, Ptr a, Ptr a, Int)
-
-emptyMem3 :: Mem3 a
-emptyMem3 = (nullPtr, nullPtr, nullPtr, 0)
-
--- assume resulting memory is aligned
-alignedZip :: Ptr Word8 -> Ptr Word8 -> Ptr Word8 -> Int
-              -> (Mem3 Word, Mem3 Word8)
-alignedZip aptr bptr cptr size =
-  let alg = alignment (undefined :: Word) in
-  if (aptr `isAlignedBy` alg) && (bptr `isAlignedBy` alg)
-  then
-    let asize = alignLow (nullPtr `plusPtr` size) `minusPtr` nullPtr
-    in
-       ( (castPtr aptr, castPtr bptr, castPtr cptr, asize `div` alg)
-       , ( aptr `advancePtr` asize
-         , bptr `advancePtr` asize
-         , cptr `advancePtr` asize
-         , (size - asize)
-         )
-       )
-  else (emptyMem3, (aptr, bptr, cptr, size))
-
--- force specialization
-zipWithBS :: (Word -> Word -> Word)
-          -> (Word8 -> Word8 -> Word8)
-          -> ByteString -> ByteString -> ByteString
-zipWithBS f g a b =
-    let (afptr, aoff, asize) = B.toForeignPtr a
-        (bfptr, boff, bsize) = B.toForeignPtr b
-        size                = min asize bsize in
-    B.unsafeCreate size $ \rptr -> do
-      withForeignPtr afptr $ \_aptr -> do
-        withForeignPtr bfptr $ \_bptr -> do
-          let aptr = _aptr `advancePtr` aoff
-          let bptr = _bptr `advancePtr` boff
-
-          let (mid, hgh) = alignedZip aptr bptr rptr size
-          zipWords mid
-          zipBytes hgh
-  where
-    zipBytes :: (Ptr Word8, Ptr Word8, Ptr Word8, Int) -> IO ()
-    zipBytes (aptr, bptr, rptr, n) = go 0
-      where
-        go :: Int -> IO ()
-        go i |   i < n   = do -- TODO unfold
-               av <- peekElemOff aptr i
-               bv <- peekElemOff bptr i
-               pokeElemOff rptr i (g av bv)
-               go (succ i)
-             | otherwise = return ()
-
-    zipWords :: (Ptr Word, Ptr Word, Ptr Word, Int) -> IO ()
-    zipWords (aptr, bptr, rptr, n) = go 0
-      where
-        go :: Int -> IO ()
-        go i |   i < n   = do -- TODO unfold
-               av <- peekElemOff aptr i
-               bv <- peekElemOff bptr i
-               pokeElemOff rptr i (f av bv)
-               go (succ i)
-             | otherwise = return ()
-
-
-
-zipWithBF :: (forall a. Bits a => a -> a -> a) -> Bitfield -> Bitfield -> Bitfield
-zipWithBF f a b = MkBitfield $ zipWithBS f f (bfBits a) (bfBits b)
-{-# INLINE zipWithBF #-}
-
-findSet :: ByteString -> Maybe Int
-findSet 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 lowOff = fst low `minusPtr` ptr
-      let midOff = fst mid `minusPtr` ptr
-      let hghOff = fst hgh `minusPtr` ptr
-
-      let resL = (lowOff +) <$> goFind  low
-      let resM = (midOff +) <$> goFind (mid :: Mem Word) -- tune size here
-                                       --  TODO: with Word8
-                                       -- bytestring findIndex works 2
-                                       -- times faster.
-      let resH = (hghOff +) <$> goFind  hgh
-
-      let res  = resL <|> resM <|> resH
-
-      -- computation of res should not escape withForeignPtr
-      case res of
-        Nothing -> return ()
-        Just _  -> return ()
-
-      return res
+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
-    goFind :: (Storable a, Eq a, Num a) => Mem a -> Maybe Int
-    goFind (ptr, n) = go 0
-      where
-        go :: Int -> Maybe Int
-        go i |   i < n   =
-               let v = B.inlinePerformIO (peekElemOff ptr i) in
-               if v /= 0
-                 then Just i
-                 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
+    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
-    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
+    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 = zipWithBF (.|.)
-{-# INLINE union #-}
+union a b = Bitfield {
+    bfSize = bfSize a `max` bfSize b
+  , bfSet  = bfSet a `S.union` bfSet b
+  }
 
 intersection :: Bitfield -> Bitfield -> Bitfield
-intersection = zipWithBF (.&.)
-{-# INLINE intersection #-}
+intersection a b = Bitfield {
+    bfSize = bfSize a `min` bfSize b
+  , bfSet  = bfSet a `S.intersection` bfSet b
+  }
 
 difference :: Bitfield -> Bitfield -> Bitfield
-difference = zipWithBF diffWord8
-  where
-    diffWord8 :: Bits a => a -> a -> a
-    diffWord8 a b = a .&. (a `xor` b)
-    {-# INLINE diffWord8 #-}
-{-# INLINE difference #-}
+difference a b = Bitfield {
+    bfSize = bfSize a -- FIXME is it more reasonable?
+  , bfSet  = bfSet a `S.difference` bfSet b
+  }
 
-combine :: [Bitfield] -> Maybe Bitfield
-combine [] = Nothing
-combine as = return $ foldr1 intersection as
+unions :: [Bitfield] -> Bitfield
+unions = foldl' union (empty 0)
 
-haveCount :: Bitfield -> Int
-haveCount (MkBitfield b) = foldBS f f 0 b
-  where
-    f byte count = popCount byte + count
+{-----------------------------------------------------------------------
+    Serialization
+-----------------------------------------------------------------------}
 
-completeness :: Bitfield -> (Int, Int)
-completeness bf = (haveCount bf, bitfieldBitCount bf)
+getBitfield :: Int -> Get Bitfield
+getBitfield = error "getBitfield"
 
--- | Get min index of piece that the peer have.
-findMin :: Bitfield -> Maybe Int
-findMin (MkBitfield b) = do
-    byteIx <- findSet b
-    bitIx  <- findMinWord8 (B.index b byteIx)
-    return $ byteIx * bitSize (undefined :: Word8) + bitIx
-  where
-    -- TODO: bit tricks
-    findMinWord8 :: Word8 -> Maybe Int
-    findMinWord8 byte = L.find (testBit byte) [0..bitSize (undefined :: Word8) - 1]
-    {-# INLINE findMinWord8 #-}
-{-# INLINE findMin #-}
-
-
-findMax :: Bitfield -> Maybe Int
-findMax (MkBitfield b) = do
-    -- TODO avoid reverse
-    byteIx <- (pred (B.length b) -) <$> findSet (B.reverse b)
-    bitIx  <- findMaxWord8 (B.index b byteIx)
-    return $ byteIx * bitSize (undefined :: Word8) + bitIx
-  where
-    -- TODO: bit tricks
-    findMaxWord8 :: Word8 -> Maybe Int
-    findMaxWord8 byte = L.find (testBit byte)
-                             (reverse [0 :: Int ..
-                                            bitSize (undefined :: Word8) - 1])
+putBitfield :: Bitfield -> Put
+putBitfield = error "putBitfield"
+
+bitfieldByteCount :: Bitfield -> Int
+bitfieldByteCount = error "bitfieldByteCount"
 
-{-# INLINE findMax #-}
+{-----------------------------------------------------------------------
+    Debug
+-----------------------------------------------------------------------}
 
-frequencies :: [Bitfield] -> [Int]
-frequencies xs = foldr1 (zipWith (+)) $ map (map fromEnum . toList) xs
+mkBitfield :: PieceCount -> [PieceIx] -> Bitfield
+mkBitfield s ixs = Bitfield {
+    bfSize = s
+  , bfSet  = S.splitLT s $ S.fromList ixs
+  }
\ No newline at end of file