From: Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
Date: Thu, 3 Sep 2015 11:02:49 +0000 (+0530)
Subject: Peer: Use Monad Transformers to simplify code
X-Git-Url: https://git.rkrishnan.org/?p=functorrent.git;a=commitdiff_plain;h=755a5f299f8ff9768d98fd780dd8a51420ae3192

Peer: Use Monad Transformers to simplify code

Wrap PState in a State Monad. Reorganization of the code in progress.
---

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index afc2028..65372d9 100644
--- a/src/FuncTorrent/Peer.hs
+++ b/src/FuncTorrent/Peer.hs
@@ -10,7 +10,7 @@ import System.IO (Handle, BufferMode(..), hSetBuffering)
 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, empty)
 import qualified Data.ByteString.Char8 as BC (length)
 import Network (connectTo, PortID(..))
-import Control.Monad (liftM, forM)
+import Control.Monad.State
 import Data.Bits
 import Data.Word (Word8)
 import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust)
@@ -21,13 +21,14 @@ import FuncTorrent.Utils (splitN, splitNum)
 import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset)
 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
 
--- PeerState is a misnomer
-data PeerState = PeerState { handle :: Handle
-                           , peer :: Peer
-                           , meChoking :: Bool
-                           , meInterested :: Bool
-                           , heChoking :: Bool
-                           , heInterested :: Bool}
+data PState = PState { handle :: Handle
+                     , peer :: Peer
+                     , meChoking :: Bool
+                     , meInterested :: Bool
+                     , heChoking :: Bool
+                     , heInterested :: Bool}
+
+type PeerState = State PState
 
 data PieceDlState = Pending
                   | InProgress
@@ -36,7 +37,7 @@ data PieceDlState = Pending
 
 -- todo - map with index to a new data structure (peers who have that piece amd state)
 data PieceData = PieceData { peers :: [Peer]        -- ^ list of peers who have this piece
-                           , state :: PieceDlState  -- ^ state of the piece from download perspective.
+                           , dlstate :: PieceDlState  -- ^ state of the piece from download perspective.
                            , hash  :: ByteString    -- ^ piece hash
                            , len :: Integer }       -- ^ piece length
 
@@ -49,7 +50,7 @@ type PieceMap = Map Integer PieceData
 mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
 mkPieceMap numPieces pieceHash pLengths = fromList kvs
   where kvs = [(i, PieceData { peers = []
-                             , state = Pending
+                             , dlstate = Pending
                              , hash = h
                              , len = pLen })
               | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
@@ -57,7 +58,7 @@ mkPieceMap numPieces pieceHash pLengths = fromList kvs
 
 havePiece :: PieceMap -> Integer -> Bool
 havePiece pm index =
-  state (pm ! index) == Have
+  dlstate (pm ! index) == Have
 
 connectToPeer :: Peer -> IO Handle
 connectToPeer (Peer _ ip port) = do
@@ -82,60 +83,69 @@ bitfieldToList bs = go bs 0
           in
            setBits ++ go bs' (pos + 1)
 
--- recvMsg :: Peer -> Handle -> Msg
-msgLoop :: PeerState -> PieceMap -> FilePath -> IO ()
-msgLoop pState@(PeerState { meInterested = False, heChoking = True }) pieceStatus file =
-  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 file
-msgLoop pState@(PeerState { meInterested = True, heChoking = False }) pieceStatus file =
-  -- 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 file fileOffset pBS
-       msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus) file
-msgLoop pState pieceStatus file = 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 file
-   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' file
-   UnChokeMsg ->
-     msgLoop (pState { heChoking = False }) pieceStatus file
-   _ ->
-     msgLoop pState pieceStatus file
+-- helper functions to manipulate PeerState
+toPeerState :: Handle -> Peer -> Bool -> Bool -> Bool -> Bool -> PState
+toPeerState h p meCh meIn heCh heIn =
+  PState { handle = h
+         , peer = p
+         , heChoking = heCh
+         , heInterested = heIn
+         , meChoking = meCh
+         , meInterested = meIn }
+
+-- -- recvMsg :: Peer -> Handle -> Msg
+-- msgLoop :: PeerState -> PieceMap -> FilePath -> IO ()
+-- msgLoop pState@(PeerState { meInterested = False, heChoking = True }) pieceStatus file = 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 file
+-- msgLoop pState@(PeerState { meInterested = True, heChoking = False }) pieceStatus file =
+--   -- 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 file fileOffset pBS
+--        msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus) file
+-- msgLoop pState pieceStatus file = 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 file
+--    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' file
+--    UnChokeMsg ->
+--      msgLoop (pState { heChoking = False }) pieceStatus file
+--    _ ->
+--      msgLoop pState pieceStatus file
 
 -- 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) -> state v == Pending) pieceList
+      allPending = filter (\(_, v) -> dlstate v == Pending) pieceList
   in
    case allPending of
     [] -> Nothing
@@ -151,12 +161,7 @@ handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
 handlePeerMsgs p m peerId = do
   h <- connectToPeer p
   doHandshake h p (infoHash m) peerId
-  let state = PeerState { handle = h
-                        , peer = p
-                        , heInterested = False
-                        , heChoking = True
-                        , meInterested = False
-                        , meChoking = True }
+  let pstate = toPeerState h p False True False True
       pieceHash = pieces (info m)
       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
       pLen = pieceLength (info m)
@@ -164,8 +169,21 @@ handlePeerMsgs p m peerId = do
       fileName = name (info m)
       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
   createDummyFile fileName (fromIntegral fileLen)
-  msgLoop state pieceStatus fileName
-  
+  (r, _) <- runStateT (msgLoop pieceStatus fileName) pstate
+  return ()
+
+msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
+msgLoop pieceStatus file =
+  StateT(\pState -> do
+             let h = handle pState
+             msg <- getMsg h
+             liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
+             case msg of
+               KeepAliveMsg -> do
+                 sendMsg h KeepAliveMsg
+                 liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
+                 runStateT (msgLoop pieceStatus file) pState)
+
 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 downloadPiece h index pieceLength = do
   let chunks = splitNum pieceLength 16384