{-# LANGUAGE CPP, TypeOperators, OverlappingInstances, RankNTypes #-} {- Copyright (c) 2005-2006 Lemmih 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