From: Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
Date: Fri, 4 Sep 2015 06:15:39 +0000 (+0530)
Subject: Peer: refactor pickPiece function
X-Git-Url: https://git.rkrishnan.org/%5B/%5D%20/file/URI:LIT:krugkidfnzsc4/@@named=/index.php?a=commitdiff_plain;h=205bce74b5786543949b162203a6e8e61e72f5f9;p=functorrent.git

Peer: refactor pickPiece function

Use the headMay from the Safe package to get a Maybe list head and
lift fst to Maybe to get the key.
---

diff --git a/functorrent.cabal b/functorrent.cabal
index 2db1119..2ec81bb 100644
--- a/functorrent.cabal
+++ b/functorrent.cabal
@@ -42,6 +42,7 @@ library
                        network-uri,
                        parsec,
                        QuickCheck,
+                       safe,
                        tasty,
                        tasty-hunit
 
@@ -63,7 +64,8 @@ executable functorrent
                        network,
                        network-uri,
                        parsec,
-                       QuickCheck
+                       QuickCheck,
+                       safe
 
 test-suite functorrent-test
   type:              exitcode-stdio-1.0
diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index 5e23fa1..0b9ece7 100644
--- 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 =