3 files changed, 33 insertions, 11 deletions

diff --git a/bench/Main.hs b/bench/Main.hs
index e4f756a3..70eb4822 100644
--- a/bench/Main.hs
+++ b/bench/Main.hs
@@ -5,6 +5,7 @@ import Control.Applicative
 import Control.DeepSeq
 import Criterion.Main
 import Data.ByteString (ByteString)
+import qualified Data.ByteString as B
 import Data.Serialize
 import Network.BitTorrent as BT
 
@@ -51,22 +52,38 @@ bitfieldUnion n = BT.empty n `union` BT.empty n
 selectionStrictFirst :: Int -> Maybe Int
 selectionStrictFirst n = strictFirst (BT.empty n) (BT.empty n) []
 
+selectionStrictLast :: Int -> Maybe Int
+selectionStrictLast n = strictLast (BT.empty n) (BT.empty n) []
+
 main :: IO ()
 main = do
-  let datas = replicate 10000 (Request (BlockIx 0 0 0))
-  let m = 1024 * 1024
+  let blockixs  = replicate 5000 (Request (BlockIx 0 0 0))
+  let bitfields = replicate 5000 (Bitfield (MkBitfield (B.replicate 1000 0)))
+  let chokes    = replicate 5000 Choke
+  let havenones = replicate 5000 HaveNone
 
-  defaultMain
-    [ datas `deepseq` bench "message/encode"   $ nf encodeMessages datas
-    , let binary = encodeMessages datas in
-      binary `deepseq` bench "message/decode"  $ nf decodeMessages binary
+  let m = 1024 * 1024
 
-      -- ~ 256KiB * 10M = 2.5TiB
-    , bench "bitfield/min"          $ nf bitfieldMin   (10 * m)
+  defaultMain $
+    concatMap (uncurry mkMsgBench)
+    [ ("blockIx",  blockixs)
+    , ("bitfield", bitfields)
+    , ("choke",    chokes)
+    , ("havenone", havenones)
+    ]
+     ++ -- 256KiB * 10M = 2.5TiB
+    [ bench "bitfield/min"          $ nf bitfieldMin   (10 * m)
     , bench "bitfield/max"          $ nf bitfieldMax   (10 * m)
     , bench "bitfield/difference"   $ nf bitfieldDiff  (10 * m)
     , bench "bitfield/intersection" $ nf bitfieldInter (10 * m)
     , bench "bitfield/union"        $ nf bitfieldUnion (10 * m)
 
     , bench "selection/strictFirst" $ nf selectionStrictFirst  (10 * m)
-    ]
\ No newline at end of file
+    , bench "selection/strictFirst" $ nf selectionStrictLast   (10 * m)
+    ]
+ where
+    mkMsgBench name msgs =
+      [ msgs `deepseq` bench ("message/" ++ name ++ "/encode") $ nf encodeMessages msgs
+      , let binary = encodeMessages msgs in
+        binary `deepseq` bench ("message/" ++ name ++ "/decode") $ nf decodeMessages binary
+      ]
diff --git a/src/Network/BitTorrent/PeerWire/Message.hs b/src/Network/BitTorrent/PeerWire/Message.hs
index 188f0572..cc771966 100644
--- a/src/Network/BitTorrent/PeerWire/Message.hs
+++ b/src/Network/BitTorrent/PeerWire/Message.hs
@@ -12,6 +12,7 @@ import Data.Serialize
 import Network.BitTorrent.PeerWire.Block
 import Network.BitTorrent.PeerWire.Bitfield
 
+import Data.Array
 
 -- | Messages used in communication between peers.
 --
@@ -78,7 +79,7 @@ data Message = KeepAlive
 instance Serialize Message where
   get = do
     len <- getInt
-    _   <- lookAhead $ ensure len
+--    _   <- lookAhead $ ensure len
     if len == 0 then return KeepAlive
       else do
         mid <- getWord8
diff --git a/src/Network/BitTorrent/PeerWire/Selection.hs b/src/Network/BitTorrent/PeerWire/Selection.hs
index 04049812..2e412e06 100644
--- a/src/Network/BitTorrent/PeerWire/Selection.hs
+++ b/src/Network/BitTorrent/PeerWire/Selection.hs
@@ -23,7 +23,8 @@
 --
 module Network.BitTorrent.PeerWire.Selection
        ( Selector
-       , strictFirst, rarestFirst, randomFirst, endGame, autoSelector
+       , strictFirst, strictLast
+       , rarestFirst, randomFirst, endGame, autoSelector
        ) where
 
 import Network.BitTorrent.PeerWire.Block
@@ -41,6 +42,9 @@ type Selector =  Bitfield      -- ^ Indices of client "have" pieces.
 strictFirst :: Selector
 strictFirst h a _ = findMin (difference a h)
 
+-- | Select the last available piece.
+strictLast :: Selector
+strictLast h a _ = findMax (difference a h)
 
 -- |
 rarestFirst :: Selector