PendingSet.hs

module Data.PendingSet
    ( PendingSet
    , Data.PendingSet.empty
    , Data.PendingSet.size
    , have
    , unhave
    , haves
    , unhaves
    , pick
    )
where

import Data.PSQueue hiding (foldl)

import Torrent

-- | Representation of Pending Sets.
newtype PendingSet = PendingSet { unPS :: PSQ PieceNum Int }

-- | The empty pending set.
empty :: PendingSet
empty = PendingSet Data.PSQueue.empty

size :: PendingSet -> Int
size = Data.PSQueue.size . unPS

-- | A peer has a given piece. Reflect this in the PendingSet.
have :: PieceNum -> PendingSet -> PendingSet
have pn = PendingSet . alter f pn . unPS
  where f Nothing = Just 1
        f (Just x) = Just (x + 1)

-- | A Peer does not have a given piece anymore (TODO: Not used in practice)
unhave :: PieceNum -> PendingSet -> PendingSet
unhave pn = PendingSet . alter f pn . unPS
  where f Nothing  =  error "Data.PendingSet.unhave"
        f (Just 1) = Nothing
        f (Just x) = Just (x-1)

-- | Add all pieces in a bitfield
haves :: [PieceNum] -> PendingSet -> PendingSet
haves pns = flip (foldl f) pns
  where f e = flip have e

-- | Remove all pieces in a bitfield
unhaves :: [PieceNum] -> PendingSet -> PendingSet
unhaves pns = flip (foldl f) pns
  where f e = flip unhave e

-- | Crawl through the set of pending pieces in decreasing order of rarity.
-- Each piece is discriminated by a selector function until the first hit is
-- found. Then all Pieces of the same priority accepted by the selector is
-- chosen for return.
pick :: (PieceNum -> IO Bool) -> PendingSet -> IO (Maybe [PieceNum])
pick selector ps = findPri (minView . unPS $ ps)
  where findPri Nothing = return Nothing
        findPri (Just (pn :-> p, rest)) = do
            r <- selector pn
            if r
                then pickAtPri numToPick p [pn] (minView rest)
                else findPri $ minView rest
        pickAtPri 0 _p acc _ = return $ Just acc
        pickAtPri _ _p acc Nothing = return $ Just acc
        pickAtPri k p acc (Just (pn :-> p', rest))
            | p == p' = do
                r <- selector pn
                if r
                    then pickAtPri (k-1) p (pn : acc) $ minView rest
                    else pickAtPri k p acc $ minView rest
            | otherwise = return $ Just acc

-- | Number of pieces to pick with the picker. Setting an upper limit here because if a lot
--   of peers have all pieces, these numbers grow insanely big, leading to allocation we
--   don't really need.
numToPick :: Int
numToPick = 7