refactor msgloop a bit

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index 88e5a7105e73c8e74719519659cedbcb30e280c1..cdc0ea612f991622879bfea1193e340832baeef4 100644 (file)
--- a/src/FuncTorrent/Peer.hs
+++ b/src/FuncTorrent/Peer.hs
@@ -190,50 +190,51 @@ bitfieldToList bs = go bs 0
 
 -- recvMsg :: Peer -> Handle -> Msg
 msgLoop :: PeerState -> PieceMap -> IO ()
-msgLoop pState pieceStatus | not (meInterested pState) && heChoking pState = do
-                               -- if me NOT Interested and she is Choking, tell her that
-                               -- I am interested.
-                               let h = handle pState
-                               sendMsg h InterestedMsg
-                               putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
-                               msgLoop (pState { meInterested = True }) pieceStatus
-                          | meInterested pState && not (heChoking pState) =
-                              -- if me Interested and she not Choking, send her a request
-                              -- for a piece.
-                              case pickPiece pieceStatus of
-                               Nothing -> putStrLn "Nothing to download"
-                               Just workPiece -> do
-                                 let pLen = len (pieceStatus ! workPiece)
-                                 putStrLn $ "piece length = " ++ show pLen
-                                 pBS <- downloadPiece (handle pState) workPiece pLen
-                                 if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
-                                   then
-                                   putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
-                                   else do
-                                   let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
-                                   putStrLn $ "Write into file at offset: " ++ show fileOffset
-                                   writeFileAtOffset "/tmp/download.file" fileOffset pBS
-                                   msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
-                          | otherwise = do
-                              msg <- getMsg (handle pState)
-                              putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
-                              case msg of
-                               KeepAliveMsg -> do
-                                 sendMsg (handle pState) KeepAliveMsg
-                                 putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
-                                 msgLoop pState pieceStatus
-                               BitFieldMsg bss -> do
-                                 let pieceList = bitfieldToList (unpack bss)
-                                     pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
-                                 putStrLn $ show (length pieceList) ++ " Pieces"
-                                 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
-                                 -- map with pieceIndex as the key and modify the value to add the peer.
-                                 -- download each of the piece in order
-                                 msgLoop pState pieceStatus'
-                               UnChokeMsg ->
-                                 msgLoop (pState { heChoking = False }) pieceStatus
-                               _ ->
-                                 msgLoop pState pieceStatus
+msgLoop pState@(PeerState { meInterested = False , heChoking = True }) pieceStatus =
+  do
+    -- if me NOT Interested and she is Choking, tell her that
+    -- I am interested.
+    let h = handle pState
+    sendMsg h InterestedMsg
+    putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
+    msgLoop (pState { meInterested = True }) pieceStatus
+msgLoop pState@(PeerState { meInterested = True, heChoking = False }) pieceStatus =
+  -- if me Interested and she not Choking, send her a request
+  -- for a piece.
+  case pickPiece pieceStatus of
+   Nothing -> putStrLn "Nothing to download"
+   Just workPiece -> do
+     let pLen = len (pieceStatus ! workPiece)
+     putStrLn $ "piece length = " ++ show pLen
+     pBS <- downloadPiece (handle pState) workPiece pLen
+     if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
+       then
+       putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
+       else do
+       let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
+       putStrLn $ "Write into file at offset: " ++ show fileOffset
+       writeFileAtOffset "/tmp/download.file" fileOffset pBS
+       msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
+msgLoop pState pieceStatus = do
+  msg <- getMsg (handle pState)
+  putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
+  case msg of
+   KeepAliveMsg -> do
+     sendMsg (handle pState) KeepAliveMsg
+     putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
+     msgLoop pState pieceStatus
+   BitFieldMsg bss -> do
+     let pieceList = bitfieldToList (unpack bss)
+         pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
+     putStrLn $ show (length pieceList) ++ " Pieces"
+     -- for each pieceIndex in pieceList, make an entry in the pieceStatus
+     -- map with pieceIndex as the key and modify the value to add the peer.
+     -- download each of the piece in order
+     msgLoop pState pieceStatus'
+   UnChokeMsg ->
+     msgLoop (pState { heChoking = False }) pieceStatus
+   _ ->
+     msgLoop pState pieceStatus
 
 -- simple algorithm to pick piece.
 -- pick the first piece from 0 that is not downloaded yet.