src/FuncTorrent/PieceManager.hs

   1 module FuncTorrent.PieceManager
   2        (PieceDlState(..),
   3         PieceData(..),
   4         PieceMap,
   5         pieceNumToOffset,
   6         updatePieceAvailability,
   7         pickPiece,
   8         bytesDownloaded,
   9         initPieceMap,
  10        ) where
  11
  12 import           Prelude hiding (filter)
  13
  14 import qualified Data.ByteString.Char8 as BC (length)
  15 import           Control.Monad (liftM)
  16 import           Data.ByteString (ByteString)
  17 import           Data.Map (Map, (!), fromList, toList, mapWithKey, filter)
  18 import           Safe (headMay)
  19
  20 import           FuncTorrent.PeerMsgs (Peer)
  21 import           FuncTorrent.Utils (splitN, splitNum)
  22
  23 data PieceDlState = Pending
  24                   | Downloading
  25                   | Have
  26                   deriving (Show, Eq)
  27
  28 -- todo - map with index to a new data structure (peers who have that piece and state)
  29 data PieceData = PieceData { peers :: [Peer]        -- ^ list of peers who have this piece
  30                            , dlstate :: PieceDlState  -- ^ state of the piece from download perspective.
  31                            , hash  :: ByteString    -- ^ piece hash
  32                            , len :: Integer }       -- ^ piece length
  33
  34 -- which piece is with which peers
  35 type PieceMap = Map Integer PieceData
  36
  37 pieceNumToOffset :: PieceMap -> Integer -> Integer
  38 pieceNumToOffset _        0 = 0
  39 pieceNumToOffset pieceMap k = k * len (pieceMap ! (k - 1))
  40
  41 -- simple algorithm to pick piece.
  42 -- pick the first piece from 0 that is not downloaded yet.
  43 pickPiece :: PieceMap -> Maybe Integer
  44 pickPiece =
  45   (fst `liftM`) . headMay . toList . filter (\v -> dlstate v == Pending)
  46
  47 bytesDownloaded :: PieceMap -> Integer
  48 bytesDownloaded =
  49   sum . map (len . snd) . toList . filter (\v -> dlstate v == Have)
  50
  51 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
  52 updatePieceAvailability pieceStatus p pieceList =
  53   mapWithKey (\k pd -> if k `elem` pieceList
  54                        then (pd { peers = p : peers pd })
  55                        else pd) pieceStatus
  56
  57 -- Make the initial Piece map, with the assumption that no peer has the
  58 -- piece and that every piece is pending download.
  59 initPieceMap :: ByteString  -> Integer -> Integer -> PieceMap
  60 initPieceMap pieceHash fileLen pieceLen = fromList kvs
  61   where
  62     numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
  63     kvs = [(i, PieceData { peers = []
  64                          , dlstate = Pending
  65                          , hash = h
  66                          , len = pLen })
  67           | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
  68     hashes = splitN 20 pieceHash
  69     pLengths = splitNum fileLen pieceLen