src/FuncTorrent/PieceManager.hs

   1 {-
   2  - Copyright (C) 2015-2016 Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
   3  -
   4  - This file is part of FuncTorrent.
   5  -
   6  - FuncTorrent is free software; you can redistribute it and/or modify
   7  - it under the terms of the GNU General Public License as published by
   8  - the Free Software Foundation; either version 3 of the License, or
   9  - (at your option) any later version.
  10  -
  11  - FuncTorrent is distributed in the hope that it will be useful,
  12  - but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14  - GNU General Public License for more details.
  15  -
  16  - You should have received a copy of the GNU General Public License
  17  - along with FuncTorrent; if not,  see <http://www.gnu.org/licenses/>
  18  -}
  19
  20 module FuncTorrent.PieceManager
  21        (PieceDlState(..),
  22         PieceData(..),
  23         PieceMap,
  24         pieceNumToOffset,
  25         updatePieceAvailability,
  26         pickPiece,
  27         bytesDownloaded,
  28         initPieceMap,
  29        ) where
  30
  31 import           Prelude hiding (filter)
  32
  33 import qualified Data.ByteString.Char8 as BC (length)
  34 import           Control.Monad (liftM)
  35 import           Data.ByteString (ByteString)
  36 import           Data.Map (Map, (!), fromList, toList, mapWithKey, filter)
  37 import           Safe (headMay)
  38
  39 import           FuncTorrent.PeerMsgs (Peer)
  40 import           FuncTorrent.Utils (splitN, splitNum)
  41
  42 data PieceDlState = Pending
  43                   | Downloading
  44                   | Have
  45                   deriving (Show, Eq)
  46
  47 -- todo - map with index to a new data structure (peers who have that piece and state)
  48 data PieceData = PieceData { peers :: [Peer]        -- ^ list of peers who have this piece
  49                            , dlstate :: PieceDlState  -- ^ state of the piece from download perspective.
  50                            , hash  :: ByteString    -- ^ piece hash
  51                            , len :: Integer }       -- ^ piece length
  52
  53 -- which piece is with which peers
  54 type PieceMap = Map Integer PieceData
  55
  56 pieceNumToOffset :: PieceMap -> Integer -> Integer
  57 pieceNumToOffset _        0 = 0
  58 pieceNumToOffset pieceMap k = k * len (pieceMap ! (k - 1))
  59
  60 -- simple algorithm to pick piece.
  61 -- pick the first piece from 0 that is not downloaded yet.
  62 pickPiece :: PieceMap -> Maybe Integer
  63 pickPiece =
  64   (fst `liftM`) . headMay . toList . filter (\v -> dlstate v == Pending)
  65
  66 bytesDownloaded :: PieceMap -> Integer
  67 bytesDownloaded =
  68   sum . map (len . snd) . toList . filter (\v -> dlstate v == Have)
  69
  70 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
  71 updatePieceAvailability pieceStatus p pieceList =
  72   mapWithKey (\k pd -> if k `elem` pieceList
  73                        then (pd { peers = p : peers pd })
  74                        else pd) pieceStatus
  75
  76 -- Make the initial Piece map, with the assumption that no peer has the
  77 -- piece and that every piece is pending download.
  78 initPieceMap :: ByteString  -> Integer -> Integer -> PieceMap
  79 initPieceMap pieceHash fileLen pieceLen = fromList kvs
  80   where
  81     numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
  82     kvs = [(i, PieceData { peers = []
  83                          , dlstate = Pending
  84                          , hash = h
  85                          , len = pLen })
  86           | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
  87     hashes = splitN 20 pieceHash
  88     pLengths = splitNum fileLen pieceLen