Piecemap.hs

{-# LANGUAGE CPP, TypeOperators, OverlappingInstances, RankNTypes #-}
{-
 Copyright (c) 2005-2006 Lemmih <lemmih@gmail.com>

 Permission to use, copy, modify, and distribute this software for any
 purpose with or without fee is hereby granted, provided that the above
 copyright notice and this permission notice appear in all copies.

 THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Conjure.InterestTable
-- Copyright   :  (c) Lemmih, 2005
-- License     :  BSD-style
--
-- Maintainer  :  lemmih@gmail.com
-- Stability   :  experimental
-- Portability :  non-portable (requires DiffArrays)
--
--
-- The InterestTable keeps track of what pieces are
-- interesting. Interesting pieces are those pieces which few peers have
-- downloaded. If a piece is rare, it is a prime candidate for sharing
-- with others and most of the Bittorrent protocol relies on this fact
-- for fast piece retrieval.
--
-- This module provides a Table of interests. Every time we learn
-- something about the distribution of pieces, we update this table to
-- track the new information. Thus, we keep a table of things we should
-- be interested in. Of course, it is also possible to query a table to
-- find the next piece we should download.
-----------------------------------------------------------------------------
module Conjure.Piecemap
    ( drawPiecemap  -- :: Piecemap -> ShowS
    , emptyUsecount -- :: Int -> Usecount
    , emptyPiecemap -- :: Int -> Piecemap
    , mkPiecemap    -- :: Int -> ByteString -> Piecemap
    , fromPiecemap  -- :: Piecemap -> ByteString
    , findNewPieces -- :: Piecemap -- ^ Our pieces. @findNewPiece@ may not return (Just i) if this!i == True.
                    -- -> Piecemap -- ^ Remote pieces. @findNewPieces@ may not return (Just i) if this!i == False.
                    -- -> [Int]    -- ^ Lazy list of pieces we can get from the remote peer.
    , orderPieces
    , setPiecemapBit -- TVar Piecemap -> Int -> Bool -> STM ()
    , addPiecemap   -- :: Piecemap -> Usecount -> Usecount
    , addPiece      -- :: Int -> Usecount -> Usecount
    , delPiecemap   -- :: Piecemap -> Usecount -> Usecount
    , scanTorrent   -- :: Torrent -> TorrentHandle -> IO Piecemap
    ) where

import Data.List                ( groupBy, sortBy, group )
import Data.Ord                 ( comparing )
import Data.Word                ( Word8 )
import Data.Bits                ( Bits (..) )
import Data.Array.Base
import Data.Array.Diff
import Control.Monad.ST
import Control.Exception
import Control.Concurrent.STM
import System.Random            ( StdGen )

import Conjure.Types
import Conjure.Utils.Shuffle
import Conjure.Utils.SHA1 (sha1)
import Conjure.Torrent

import qualified Data.ByteString as BS
import Data.ByteString (ByteString)

#ifdef __CABAL_TEST__
import Test.QuickCheck hiding (evaluate)
#endif

import Control.Monad
import GHC.Exts
{-
-- import Data.Array.IO            ( IOUArray )
-- import System.IO.Unsafe         (unsafePerformIO)
-- Stolen from Data.Array.Diff:
-- If the array contains unboxed elements, then the elements of the
-- diff list may also recursively reference the array from inside
-- replaceDiffArray, so we must seq them too.
replaceDiffArray2 :: (MArray a e IO, Ix i)
                => IOToDiffArray a i e
                -> [(Int, e)]
                -> IO (IOToDiffArray a i e)
a `replaceDiffArray2` ies = do
    mapM_ (\(a,b) -> do evaluate a; evaluate b) ies
    a `replaceDiffArray` ies


instance IArray (IOToDiffArray IOUArray) Bool where
    bounds        a      = bounds (unsafeCoerce# a :: IOToDiffArray IOUArray i Int)
    unsafeArray   lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt      a i    = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies  = unsafePerformIO $ a `replaceDiffArray2` ies

instance (Ix ix, Show ix) => Show (DiffUArray ix Bool) where
  showsPrec = showsIArray
-}

drawPiecemap :: Piecemap -> ShowS
drawPiecemap piecemap
    = foldr (\e r -> draw e . r) id $ group $ elems piecemap
    where draw [True]        = showChar '#'
          draw lst@(True:_)  = shows (length lst) . showChar '#'
          draw [False]       = showChar '_'
          draw lst@(False:_) = shows (length lst) . showChar '_'


-- It's important that the sha1 check is strict.
-- Otherwise we can't garbage collect the pieces and
-- we'll end up with the entire file in memory.
-- FIXME: Skip all the pieces in a file when 'readPiece'' returns Nothing.
scanTorrent :: Torrent -> Backend -> IO Piecemap
scanTorrent torrent backend
    = fmap (listArray (0, nPieces-1)) (mapM checker [0 .. nPieces - 1])
    where nPieces = infoNumPieces torrent
          checker pNum
              = do pieceMb <- readPiece' backend pNum
                   case pieceMb of
                     Nothing -> return False
                     Just piece
                       -> return $! pieceCheckSum torrent pNum == sha1 piece

emptyUsecount :: Int -> Usecount
emptyUsecount l
    = listArray (0,l-1) (replicate l 0)

emptyPiecemap :: Int -> Piecemap
emptyPiecemap l
    = listArray (0,l-1) (replicate l False)

mkPiecemap :: Int -> ByteString -> Piecemap
mkPiecemap l bitmap
    = listArray (0,l-1) (take l (bsToBitmap (BS.unpack bitmap)))

bsToBitmap :: [Word8] -> [Bool]
bsToBitmap [] = []
bsToBitmap (x:xs)
    = foldl (\s n l -> s (testBit x n:l)) id [7, 6 .. 0] (bsToBitmap xs)

fromPiecemap :: Piecemap -> ByteString
fromPiecemap arr = BS.pack (bitmapToBS (elems arr))

bitmapToBS :: [Bool] -> [Word8]
bitmapToBS [] = []
bitmapToBS xs
    = foldl (\b a -> (b `shiftL` 1) .|. fromBool a) 0 (take 8 (byte ++ repeat False)):bitmapToBS bytes
    where (byte,bytes) = splitAt 8 xs

fromBool :: (Num t) => Bool -> t
fromBool False = 0
fromBool True = 1

findNewPieces :: Piecemap -- ^ Our pieces. @findNewPiece@ may not return (Just i) if this!i == True.
              -> Piecemap -- ^ Remote pieces. @findNewPieces@ may not return (Just i) if this!i == False.
              -> [Int]    -- ^ Lazy list of pieces we can get from the remote peer.
findNewPieces ourPieces remotePieces
    = assert (bounds ourPieces == bounds remotePieces) $
      loop 0
    where len = rangeSize $ bounds ourPieces
          loop n
              | n < len = let gotPiece = ourPieces `unsafeAt` n
                              isAvailable = remotePieces `unsafeAt` n
                          in if not gotPiece && isAvailable
                             then n:loop (n+1)
                             else loop (n+1)
              | otherwise = []

-- Arrange pieces in rarest first where equally rare pieces are shuffled.
orderPieces :: StdGen -> Usecount -> [Int] -> [Int]
orderPieces gen usecount wantedPieces
    = concatMap (shuffle gen)
      $ map (map snd) $ groupBy (\a b -> fst a == fst b)
      $ sortBy (comparing fst)
            [ (unsafeAt usecount idx,idx) | idx <- wantedPieces ]

runSTDiffU :: (Ix i) => (forall s. ST s (STUArray s i Int)) -> DiffUArray i Int
runSTDiffU st = runST (st >>= unsafeFreeze)

setPiecemapBit :: TVar Piecemap -> Int -> Bool -> STM ()
setPiecemapBit piecemapVar pieceNum status
    = do piecemap <- readTVar piecemapVar
         writeTVar piecemapVar (piecemap // [(pieceNum,status)])

addPiecemap :: Piecemap -> Usecount -> Usecount
addPiecemap = modPiecemap (+)

addPiece :: Int -> Usecount -> Usecount
addPiece idx usecount
    = let n = unsafeAt usecount idx
      in usecount // [(idx, n+1)]

delPiecemap :: Piecemap -> Usecount -> Usecount
delPiecemap = modPiecemap (-)

modPiecemap :: (Int -> Int -> Int) -> Piecemap -> Usecount -> Usecount
modPiecemap ac pieceArr countArr
    = assert (bounds pieceArr == bounds countArr) $
      runSTDiffU mkArray
    where len = rangeSize $ bounds countArr
          mkArray = do arr <- newArray_ (bounds countArr)
                       let loop n
                               | n < len = let newVal = unsafeAt countArr n `ac` if unsafeAt pieceArr n then 1 else 0
                                           in unsafeWrite arr n newVal >> loop (n+1)
                               | otherwise = return ()
                       loop 0
                       return arr


--------------------------------------------------------------
-- Tests.
--------------------------------------------------------------
#ifdef __CABAL_TEST__

instance Arbitrary Word8 where
    arbitrary = fmap fromIntegral (choose (0, 255) :: Gen Int)
instance Arbitrary (DiffArray Int Bool) where
    arbitrary = sized $ \n ->
                liftM (listArray (0,n-1)) (vector n)
    coarbitrary = undefined

instance Arbitrary (DiffUArray Int Int) where
    arbitrary = sized $ \n ->
                liftM (listArray (0,n-1)) (vector n)
    coarbitrary = undefined

{-
findNewPieces :: Piecemap -- ^ Our pieces. @findNewPiece@ may not return (Just i) if this!i == True.
              -> Piecemap -- ^ Remote pieces. @findNewPieces@ may not return (Just i) if this!i == False.
              -> [Int]    -- ^ Lazy list of pieces we can get from the remote peer.
-}

_prop_add_del piecemap usecount
    = assocs usecount == assocs (delPiecemap piecemap (addPiecemap piecemap usecount))

_prop_newPieces our remote
    = let newPieces = findNewPieces our remote
      in assocs our /= assocs remote ==> and [ not (our!p) && remote!p | p <- newPieces ]

_prop_identity piecemap
    = trivial (len==0) $ assocs piecemap == assocs (mkPiecemap len bs)
    where (l,u) = bounds piecemap
          len = u-l+1
          bs = fromPiecemap piecemap

_prop_bs_to_bitmap bs = trivial (null bs) $ bs == bitmapToBS (bsToBitmap bs)
_prop_bitmap_to_bs bitmap = trivial (null bitmap) $
                            bitmap == take (length bitmap) (bsToBitmap (bitmapToBS bitmap))


#endif