{- Copyright (c) 2005 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. ----------------------------------------------------------------------------- -- FIXME: Rename this to Conjure.FindAppropriateNameForThisModule module Conjure.InterestTable ( emptyUsecount -- :: Int -> Usecount , emptyPiecemap -- :: Int -> Piecemap , mkPiecemap -- :: Int -> FastString -> Piecemap , fromPiecemap -- :: Piecemap -> FastString , 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. , setPiecemapBit -- TVar Piecemap -> Int -> Bool -> STM () , addPiecemap -- :: Piecemap -> Usecount -> Usecount , delPiecemap -- :: Piecemap -> Usecount -> Usecount , scanTorrent -- :: Torrent -> TorrentHandle -> IO Piecemap -- tests , prop_identity , prop_newPieces ) where import Data.Word import Data.Bits import Data.Array.IO import Data.Array.Base import Data.Array.Diff import Control.Monad.ST import Control.Exception import Control.Concurrent.STM import System.IO.Unsafe import Conjure.Types import SHA1 (sha1) import Conjure.Torrent import Conjure.FileSystem.Interface import qualified Data.FastPackedString as FS import Data.FastPackedString (FastString) import Test.QuickCheck hiding (evaluate) import Control.Monad -- 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 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 -- 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 -> TorrentHandle -> IO Piecemap scanTorrent torrent handle = do bits <- mapM checker [0 .. nPieces - 1] return (listArray (0,nPieces-1) bits) where nPieces = infoNumPieces torrent checker pNum = do pieceMb <- readPiece' handle 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 -> FastString -> Piecemap mkPiecemap l bitmap = listArray (0,l-1) (take l (fsToBitmap (FS.unpackWords bitmap) [])) fsToBitmap :: [Word8] -> [Bool] -> [Bool] fsToBitmap [] = id fsToBitmap (x:xs) = foldl (\s n l -> s (testBit x n:l)) id [7,6..0] . fsToBitmap xs fromPiecemap :: Piecemap -> FastString fromPiecemap arr = FS.packWords (bitmapToFS (elems arr++padding) []) where (s,e) = bounds arr padding = replicate (8 - ((e-s+1) `mod` 8)) False bitmapToFS :: [Bool] -> [Word8] -> [Word8] bitmapToFS [] = id bitmapToFS xs = \l -> foldl (\b a -> (b `shiftL` 1) .|. fromBool a) 0 byte : bitmapToFS bytes l where (byte,bytes) = splitAt 8 xs 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 (l,u) = bounds ourPieces loop n | n > (u-l) = [] | otherwise = let gotPiece = ourPieces `unsafeAt` n isAvailable = remotePieces `unsafeAt` n in if not gotPiece && isAvailable then n:loop (n+1) else loop (n+1) 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 (+) 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 (l,u) = bounds countArr mkArray = do arr <- newArray_ (bounds countArr) let loop n | n > (u-l) = return () | otherwise = let newVal = unsafeAt countArr n `ac` if unsafeAt pieceArr n then 1 else 0 in unsafeWrite arr n newVal >> loop (n+1) loop 0 return arr -------------------------------------------------------------- -- Tests. -------------------------------------------------------------- instance Arbitrary (DiffUArray 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_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 fs) where (l,u) = bounds piecemap len = u-l+1 fs = fromPiecemap piecemap