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