Peer: refactor pickPiece function

[functorrent.git] / src / FuncTorrent / Peer.hs

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index 5e23fa112574c0ca07f6ddfbde0f36212a932720..0b9ece75bc9cc88db2fc5783255ea5cc103fb436 100644 (file)
--- a/src/FuncTorrent/Peer.hs
+++ b/src/FuncTorrent/Peer.hs
@@ -4,7 +4,7 @@ module FuncTorrent.Peer
      handlePeerMsgs
     ) where
 
-import Prelude hiding (lookup, concat, replicate, splitAt, take)
+import Prelude hiding (lookup, concat, replicate, splitAt, take, filter)
 
 import System.IO (Handle, BufferMode(..), hSetBuffering)
 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, empty)
@@ -13,8 +13,9 @@ import Network (connectTo, PortID(..))
 import Control.Monad.State
 import Data.Bits
 import Data.Word (Word8)
-import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust)
+import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust, filter)
 import qualified Crypto.Hash.SHA1 as SHA1 (hash)
+import Safe (headMay)
 
 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
 import FuncTorrent.Utils (splitN, splitNum)
@@ -102,13 +103,8 @@ toPeerState h p meCh meIn heCh heIn =
 -- simple algorithm to pick piece.
 -- pick the first piece from 0 that is not downloaded yet.
 pickPiece :: PieceMap -> Maybe Integer
-pickPiece m =
-  let pieceList = toList m
-      allPending = filter (\(_, v) -> dlstate v == Pending) pieceList
-  in
-   case allPending of
-    [] -> Nothing
-    ((i, _):_) -> Just i
+pickPiece =
+  (fst `liftM`) . headMay . toList . filter (\v -> dlstate v == Pending)
 
 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 updatePieceAvailability pieceStatus p pieceList =