6 files changed, 55 insertions, 5 deletions

diff --git a/bench/serialization.hs b/bench/serialization.hs
index 8ed90958..c6d14328 100644
--- a/bench/serialization.hs
+++ b/bench/serialization.hs
@@ -15,6 +15,9 @@ instance NFData BlockIx where
 instance NFData Block where
   rnf (Block a b c) = a `deepseq` b `deepseq` rnf c
 
+instance NFData Bitfield where
+  rnf = rnf . bfBits
+
 instance NFData Message where
   rnf (Have i)     = rnf i
   rnf (Bitfield b) = rnf b
diff --git a/network-bittorrent.cabal b/network-bittorrent.cabal
index a26068f6..00ed2162 100644
--- a/network-bittorrent.cabal
+++ b/network-bittorrent.cabal
@@ -32,8 +32,9 @@ library
                      , Network.BitTorrent.Tracker.Scrape
 
                      , Network.BitTorrent.PeerWire
-                     , Network.BitTorrent.PeerWire.ClientInfo
+                     , Network.BitTorrent.PeerWire.Bitfield
                      , Network.BitTorrent.PeerWire.Block
+                     , Network.BitTorrent.PeerWire.ClientInfo
                      , Network.BitTorrent.PeerWire.Message
                      , Network.BitTorrent.PeerWire.Handshake
 
diff --git a/src/Network/BitTorrent/PeerWire.hs b/src/Network/BitTorrent/PeerWire.hs
index fde98695..9f1e02ca 100644
--- a/src/Network/BitTorrent/PeerWire.hs
+++ b/src/Network/BitTorrent/PeerWire.hs
@@ -8,6 +8,7 @@
 {-# LANGUAGE DoAndIfThenElse #-}
 module Network.BitTorrent.PeerWire (module PW) where
 
+import Network.BitTorrent.PeerWire.Bitfield as PW
 import Network.BitTorrent.PeerWire.Block as PW
 import Network.BitTorrent.PeerWire.Message as PW
 import Network.BitTorrent.PeerWire.Handshake as PW
diff --git a/src/Network/BitTorrent/PeerWire/Bitfield.hs b/src/Network/BitTorrent/PeerWire/Bitfield.hs
new file mode 100644
index 00000000..c9357a25
--- /dev/null
+++ b/src/Network/BitTorrent/PeerWire/Bitfield.hs
@@ -0,0 +1,39 @@
+-- |
+--   Copyright   :  (c) Sam T. 2013
+--   License     :  MIT
+--   Maintainer  :  pxqr.sta@gmail.com
+--   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.
+--
+module Network.BitTorrent.PeerWire.Bitfield
+       ( Bitfield(..)
+       , getBitfield, putBitfield, bitfieldByteCount
+       ) where
+
+import Control.Applicative
+import           Data.ByteString (ByteString)
+import qualified Data.ByteString as B
+
+import Data.Serialize
+
+
+newtype Bitfield = MkBitfield {
+    bfBits :: ByteString
+  } deriving (Show, Eq, Ord)
+
+bitfieldByteCount :: Bitfield -> Int
+bitfieldByteCount = B.length . bfBits
+{-# INLINE bitfieldByteCount #-}
+
+getBitfield :: Int -> Get Bitfield
+getBitfield n = MkBitfield <$> getBytes n
+{-# INLINE getBitfield #-}
+
+putBitfield :: Bitfield -> Put
+putBitfield = putByteString . bfBits
+{-# INLINE putBitfield #-}
\ No newline at end of file
diff --git a/src/Network/BitTorrent/PeerWire/Message.hs b/src/Network/BitTorrent/PeerWire/Message.hs
index 3f6647dd..9cbc2e38 100644
--- a/src/Network/BitTorrent/PeerWire/Message.hs
+++ b/src/Network/BitTorrent/PeerWire/Message.hs
@@ -1,5 +1,6 @@
 module Network.BitTorrent.PeerWire.Message
        ( Message(..)
+       , Bitfield
        , ppMessage
        ) where
 
@@ -10,6 +11,7 @@ import qualified Data.ByteString as B
 import Data.Serialize
 
 import Network.BitTorrent.PeerWire.Block
+import Network.BitTorrent.PeerWire.Bitfield
 
 
 -- | Messages used in communication between peers.
@@ -32,7 +34,7 @@ data Message = KeepAlive
                -- after the handshaking sequence is complete, and
                -- before any other message are sent. If client have no
                -- pieces then bitfield need not to be sent.
-             | Bitfield ByteString
+             | Bitfield Bitfield
 
                -- | Request for a particular block. If a client is
                -- requested a block that another peer do not have the
@@ -87,7 +89,7 @@ instance Serialize Message where
           0x02 -> return Interested
           0x03 -> return NotInterested
           0x04 -> Have     <$> getInt
-          0x05 -> Bitfield <$> getBytes (pred len)
+          0x05 -> Bitfield <$> getBitfield (pred len)
           0x06 -> Request  <$> get
           0x07 -> Piece    <$> getBlock (len - 9)
           0x08 -> Cancel   <$> get
@@ -111,8 +113,8 @@ instance Serialize Message where
   put Interested    = putInt 1  >> putWord8 0x02
   put NotInterested = putInt 1  >> putWord8 0x03
   put (Have i)      = putInt 5  >> putWord8 0x04 >> putInt i
-  put (Bitfield b)  = putInt l  >> putWord8 0x05 >> putByteString b
-    where l = succ (B.length b)
+  put (Bitfield b)  = putInt l  >> putWord8 0x05 >> putBitfield b
+    where l = succ (bitfieldByteCount b)
           {-# INLINE l #-}
   put (Request blk) = putInt 13 >> putWord8 0x06 >> put blk
   put (Piece   blk) = putInt l  >> putWord8 0x07 >> putBlock
diff --git a/tests/encoding.hs b/tests/encoding.hs
index a8490551..405d2eb1 100644
--- a/tests/encoding.hs
+++ b/tests/encoding.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# OPTIONS -fno-warn-orphans #-}
 module Main (main) where
 
@@ -29,6 +31,8 @@ instance Arbitrary BlockIx where
 instance Arbitrary Block where
   arbitrary = Block <$> positive <*> positive <*> arbitrary
 
+deriving instance Arbitrary Bitfield
+
 instance Arbitrary Message where
   arbitrary = oneof
     [ pure KeepAlive