3 files changed, 15 insertions, 8 deletions

diff --git a/src/Network/BitTorrent/Exchange.hs b/src/Network/BitTorrent/Exchange.hs
index 505360a4..66112f14 100644
--- a/src/Network/BitTorrent/Exchange.hs
+++ b/src/Network/BitTorrent/Exchange.hs
@@ -53,6 +53,7 @@ module Network.BitTorrent.Exchange
        , getHaveCount
        , getWantCount
        , getPieceCount
+       , peerOffer
 
          -- * Events
        , Event(..)
@@ -295,6 +296,10 @@ data Event
   | Fragment  Block
     deriving Show
 
+-- INVARIANT:
+--
+--   * Available Bitfield is never empty
+--
 
 -- | You could think of 'awaitEvent' as wait until something interesting occur.
 --
@@ -316,9 +321,7 @@ data Event
 --     forall (Fragment block). isPiece block == True
 --
 awaitEvent :: P2P Event
-awaitEvent = do
-
-    awaitMessage >>= go
+awaitEvent = awaitMessage >>= go
   where
     go KeepAlive = awaitEvent
     go Choke     = do
@@ -341,8 +344,7 @@ awaitEvent = do
       awaitEvent
 
     go (Have idx)      = do
-      new <- singletonBF idx
-      bitfield %= BF.union new
+      bitfield %= have idx
       _ <- revise
 
       offer <- peerOffer
diff --git a/src/Network/BitTorrent/Exchange/Protocol.hs b/src/Network/BitTorrent/Exchange/Protocol.hs
index 8d42e3a8..4cf4685d 100644
--- a/src/Network/BitTorrent/Exchange/Protocol.hs
+++ b/src/Network/BitTorrent/Exchange/Protocol.hs
@@ -237,7 +237,7 @@ data Block = Block {
   , blkOffset :: {-# UNPACK #-} !Int
 
     -- | Payload.
-  , blkData   :: !ByteString
+  , blkData   :: !ByteString -- TODO make lazy bytestring
   } deriving (Show, Eq)
 
 -- | Format block in human readable form. Payload is ommitted.
diff --git a/src/Network/BitTorrent/Internal.hs b/src/Network/BitTorrent/Internal.hs
index 38388b9a..bf47b87b 100644
--- a/src/Network/BitTorrent/Internal.hs
+++ b/src/Network/BitTorrent/Internal.hs
@@ -48,7 +48,7 @@ module Network.BitTorrent.Internal
        , leaveSwarm
        , waitVacancy
 
-       , available
+       , pieceLength
 
          -- * Peer
        , PeerSession( PeerSession, connectedPeerAddr
@@ -57,6 +57,9 @@ module Network.BitTorrent.Internal
                     )
        , SessionState
        , withPeerSession
+
+         -- ** Broadcasting
+       , available
        , getPending
 
          -- ** Exceptions
@@ -388,6 +391,7 @@ waitVacancy se =
 
 pieceLength :: SwarmSession -> Int
 pieceLength = ciPieceLength . tInfo . torrentMeta
+{-# INLINE pieceLength #-}
 
 {-----------------------------------------------------------------------
     Peer session
@@ -532,7 +536,8 @@ findPieceCount = pieceCount . tInfo . torrentMeta . swarmSession
 -- 3. Signal to the all other peer about this.
 
 available :: Bitfield -> SwarmSession -> IO ()
-available bf se @ SwarmSession {..} = mark >> atomically broadcast
+available bf se @ SwarmSession {..} = do
+    mark >> atomically broadcast
   where
     mark = do
       let bytes = pieceLength se * BF.haveCount bf