3 files changed, 467 insertions, 0 deletions

diff --git a/dht/bittorrent/src/System/Torrent/FileMap.hs b/dht/bittorrent/src/System/Torrent/FileMap.hs
new file mode 100644
index 00000000..38c475e8
--- /dev/null
+++ b/dht/bittorrent/src/System/Torrent/FileMap.hs
@@ -0,0 +1,163 @@
+{-# LANGUAGE RecordWildCards  #-}
+{-# LANGUAGE ViewPatterns     #-}
+{-# OPTIONS -fno-warn-orphans #-}
+module System.Torrent.FileMap
+       ( FileMap
+
+         -- * Construction
+       , Mode (..)
+       , def
+       , mmapFiles
+       , unmapFiles
+
+         -- * Query
+       , System.Torrent.FileMap.size
+
+         -- * Modification
+       , readBytes
+       , writeBytes
+       , unsafeReadBytes
+
+         -- * Unsafe conversions
+       , fromLazyByteString
+       , toLazyByteString
+       ) where
+
+import Control.Applicative
+import Control.Monad as L
+import Data.ByteString as BS
+import Data.ByteString.Internal as BS
+import Data.ByteString.Lazy as BL
+import Data.ByteString.Lazy.Internal as BL
+import Data.Default
+import Data.Vector as V -- TODO use unboxed vector
+import Foreign
+import System.IO.MMap
+
+import Data.Torrent
+
+
+data FileEntry = FileEntry
+  { filePosition :: {-# UNPACK #-} !FileOffset
+  , fileBytes    :: {-# UNPACK #-} !BS.ByteString -- XXX: mutable buffer (see 'writeBytes').
+  } deriving (Show, Eq)
+
+type FileMap = Vector FileEntry
+
+instance Default Mode where
+  def = ReadWriteEx
+
+mmapFiles :: Mode -> FileLayout FileSize -> IO FileMap
+mmapFiles mode layout = V.fromList <$> L.mapM mkEntry (accumPositions layout)
+  where
+    mkEntry (path, (pos, expectedSize)) = do
+      let esize = fromIntegral expectedSize -- FIXME does this safe?
+      (fptr, moff, msize) <- mmapFileForeignPtr path mode $ Just (0, esize)
+      if msize /= esize
+        then error "mmapFiles" -- TODO unmap mapped files on exception
+        else return $ FileEntry pos (PS fptr moff msize)
+
+unmapFiles :: FileMap -> IO ()
+unmapFiles = V.mapM_ unmapEntry
+  where
+    unmapEntry (FileEntry _ (PS fptr _ _)) = finalizeForeignPtr fptr
+
+-- Unsafe: FileMap 'writeBytes' will modify supplied bytestrings in place.
+fromLazyByteString :: BL.ByteString -> FileMap
+fromLazyByteString lbs = V.unfoldr f (0, lbs)
+  where
+    f (_,   Empty     ) = Nothing
+    f (pos, Chunk x xs) = Just (FileEntry pos x, ((pos + chunkSize), xs))
+      where chunkSize = fromIntegral $ BS.length x
+
+-- | /O(n)/.
+--
+-- Unsafe: mutable buffers are returned without copy.
+toLazyByteString :: FileMap -> BL.ByteString
+toLazyByteString = V.foldr f Empty
+  where
+    f FileEntry {..} bs = Chunk fileBytes bs
+
+-- | /O(1)/.
+size :: FileMap -> FileOffset
+size m
+  |           V.null m               = 0
+  | FileEntry {..} <- V.unsafeLast m
+  = filePosition + fromIntegral (BS.length fileBytes)
+
+-- | Find the file number for a particular byte offset within a torrent.
+bsearch :: FileOffset -> FileMap -> Maybe Int
+bsearch x m
+  | V.null m  = Nothing
+  | otherwise = branch (V.length m `div` 2)
+  where
+    branch c @ ((m !) -> FileEntry {..})
+      | x <  filePosition            = bsearch x (V.take c m)
+      | x >= filePosition + fileSize = do
+        ix <- bsearch x (V.drop (succ c) m)
+        return $ succ c + ix
+      |          otherwise           = Just c
+      where
+        fileSize = fromIntegral (BS.length fileBytes)
+
+-- | /O(log n)/.
+drop :: FileOffset -> FileMap -> (FileSize, FileMap)
+drop off m
+  | Just ix <- bsearch off m
+  , FileEntry {..} <- m ! ix = (off - filePosition, V.drop ix m)
+  |         otherwise        = (0                 , V.empty)
+
+-- | /O(log n)/.
+take :: FileSize -> FileMap -> (FileMap, FileSize)
+take len m
+  |           len >= s       = (m                 , 0)
+  | Just ix <- bsearch (pred len) m = let m' = V.take (succ ix) m
+                               in (m', System.Torrent.FileMap.size m' - len)
+  |          otherwise       = (V.empty           , 0)
+  where
+    s =  System.Torrent.FileMap.size m
+
+-- | /O(log n + m)/. Do not use this function with 'unmapFiles'.
+--
+-- The returned bytestring points directly into an area memory mapped from a
+-- file.
+unsafeReadBytes :: FileOffset -> FileSize -> FileMap -> BL.ByteString
+unsafeReadBytes off s m
+  | (l  , m') <- System.Torrent.FileMap.drop off     m
+  , (m'', _ ) <- System.Torrent.FileMap.take (off + s) m'
+  = BL.take (fromIntegral s) $ BL.drop (fromIntegral l) $ toLazyByteString m''
+
+-- The returned bytestring is copied and safe to use after the file is
+-- unmapped.
+readBytes :: FileOffset -> FileSize -> FileMap -> IO BL.ByteString
+readBytes off s m = do
+    let bs_copy = BL.copy $ unsafeReadBytes off s m
+    forceLBS bs_copy
+    return bs_copy
+  where
+    forceLBS  Empty      = return ()
+    forceLBS (Chunk _ x) = forceLBS x
+
+-- UNSAFE: Uses the first byte string as a pointer to mutable data and writes
+-- the contents of the second bytestring there.
+bscpy :: BL.ByteString -> BL.ByteString -> IO ()
+bscpy (PS _ _ 0 `Chunk` dest_rest) src                         = bscpy dest_rest src
+bscpy dest                         (PS _ _ 0 `Chunk` src_rest) = bscpy dest      src_rest
+bscpy (PS dest_fptr dest_off dest_size `Chunk` dest_rest)
+      (PS src_fptr  src_off  src_size  `Chunk` src_rest)
+  = do let csize = min dest_size src_size
+       withForeignPtr dest_fptr $ \dest_ptr ->
+         withForeignPtr src_fptr $ \src_ptr ->
+           memcpy (dest_ptr `advancePtr` dest_off)
+                  (src_ptr  `advancePtr` src_off)
+                  (fromIntegral csize) -- TODO memmove?
+       bscpy (PS dest_fptr (dest_off + csize) (dest_size - csize) `Chunk` dest_rest)
+             (PS src_fptr  (src_off  + csize) (src_size  - csize) `Chunk` src_rest)
+bscpy _ _ = return ()
+
+-- UNSAFE: Mutates bytestring contents within the provided FileMap.
+writeBytes :: FileOffset -> BL.ByteString -> FileMap -> IO ()
+writeBytes off lbs m = bscpy dest src
+  where
+    src  = BL.take (fromIntegral (BL.length dest)) lbs
+    dest = unsafeReadBytes off (fromIntegral (BL.length lbs)) m
diff --git a/dht/bittorrent/src/System/Torrent/Storage.hs b/dht/bittorrent/src/System/Torrent/Storage.hs
new file mode 100644
index 00000000..1d77e55d
--- /dev/null
+++ b/dht/bittorrent/src/System/Torrent/Storage.hs
@@ -0,0 +1,221 @@
+-- |
+--   Copyright   :  (c) Sam Truzjan 2013
+--   License     :  BSD3
+--   Maintainer  :  pxqr.sta@gmail.com
+--   Stability   :  experimental
+--   Portability :  portable
+--
+--   This module implements mapping from single continious piece space
+--   to file storage. Storage can be used in two modes:
+--
+--     * As in memory storage - in this case we don't touch filesystem.
+--
+--     * As ordinary mmaped file storage - when we need to store
+--       data in the filesystem.
+--
+{-# LANGUAGE RecordWildCards            #-}
+{-# LANGUAGE DeriveDataTypeable         #-}
+module System.Torrent.Storage
+       ( -- * Storage
+         Storage
+       , StorageFailure (..)
+
+         -- * Construction
+       , Mode (..)
+       , def
+       , open
+       , openInfoDict
+       , close
+       , withStorage
+
+         -- * Query
+       , totalPieces
+       , verifyPiece
+       , genPieceInfo
+       , getBitfield
+
+         -- * Modification
+       , writePiece
+       , readPiece
+       , hintRead
+       , unsafeReadPiece
+
+         -- * Streaming
+       , sourceStorage
+       , sinkStorage
+       ) where
+
+import Control.Applicative
+import Control.Exception
+import Control.Monad as M
+import Control.Monad.Trans
+import Data.ByteString.Lazy as BL
+import Data.Conduit as C
+import Data.Conduit.Binary as C
+import Data.Conduit.List as C
+import Data.Typeable
+
+import Data.Torrent
+import Network.BitTorrent.Exchange.Bitfield as BF
+import System.Torrent.FileMap as FM
+
+
+-- | Some storage operations may throw an exception if misused.
+data StorageFailure
+    -- | Occurs on a write operation if the storage has been opened
+    --   using 'ReadOnly' mode.
+  = StorageIsRO
+
+    -- | Piece index is out of bounds.
+  | InvalidIndex PieceIx
+
+    -- | Piece size do not match with one passed to the 'open'
+    -- function.
+  | InvalidSize  PieceSize
+    deriving (Show, Eq, Typeable)
+
+instance Exception StorageFailure
+
+-- | Pieces store.
+data Storage = Storage
+  { mode     ::                !Mode
+  , pieceLen :: {-# UNPACK #-} !PieceSize
+  , fileMap  :: {-# UNPACK #-} !FileMap
+  }
+
+-- | Map torrent files:
+--
+--   * when torrent first created use 'ReadWriteEx' mode;
+--
+--   * when seeding, validation 'ReadOnly' mode.
+--
+open :: Mode -> PieceSize -> FileLayout FileSize -> IO Storage
+open mode s l
+  |   s <= 0  = throwIO (InvalidSize s)
+  | otherwise = Storage mode s <$> mmapFiles mode l
+
+-- | Like 'open', but use 'InfoDict' file layout.
+openInfoDict :: Mode -> FilePath -> InfoDict -> IO Storage
+openInfoDict mode rootPath InfoDict {..} =
+  open mode (piPieceLength idPieceInfo) (flatLayout rootPath idLayoutInfo)
+
+-- | Unmaps all files forcefully. It is recommended but not required.
+close :: Storage -> IO ()
+close Storage {..} = unmapFiles fileMap
+
+-- | Normally you need to use 'Control.Monad.Trans.Resource.allocate'.
+withStorage :: Mode -> PieceSize -> FileLayout FileSize
+            -> (Storage -> IO ()) -> IO ()
+withStorage m s l = bracket (open m s l) close
+
+-- TODO allocateStorage?
+
+-- | Count of pieces in the storage.
+totalPieces :: Storage -> PieceCount
+totalPieces Storage {..} = FM.size fileMap `sizeInBase` pieceLen
+
+isValidIx :: PieceIx -> Storage -> Bool
+isValidIx i s = 0 <= i && i < totalPieces s
+
+-- | Put piece data at the piece index by overwriting existing
+-- data.
+--
+--   This operation may throw 'StorageFailure'.
+--
+writePiece :: Piece BL.ByteString -> Storage -> IO ()
+writePiece p @ Piece {..} s @ Storage {..}
+  |       mode == ReadOnly    = throwIO  StorageIsRO
+  | isNotValidIx   pieceIndex = throwIO (InvalidIndex pieceIndex)
+  | isNotValidSize pieceIndex (pieceSize p)
+                              = throwIO (InvalidSize  (pieceSize p))
+  |          otherwise        = writeBytes offset pieceData fileMap
+  where
+    isNotValidSize pix psize
+      | succ pix == pcount = psize /= lastPieceLen -- last piece may be shorter
+      |      otherwise     = psize /= pieceLen
+      where
+        lastPieceLen = fromIntegral (FM.size fileMap `rem` fromIntegral pieceLen)
+    {-# INLINE isNotValidSize #-}
+
+    isNotValidIx i = i < 0 || i >= pcount
+    {-# INLINE isNotValidIx #-}
+
+    pcount = totalPieces s
+    offset = fromIntegral pieceIndex * fromIntegral pieceLen
+
+-- | Read specific piece from storage.
+--
+--   This operation may throw 'StorageFailure'.
+--
+readPiece :: PieceIx -> Storage -> IO (Piece BL.ByteString)
+readPiece pix s @ Storage {..}
+  | not (isValidIx pix s) = throwIO (InvalidIndex pix)
+  | otherwise = Piece pix <$> readBytes offset sz fileMap
+  where
+    offset = fromIntegral pix * fromIntegral pieceLen
+    sz     = fromIntegral pieceLen
+
+-- | Hint about the coming 'readPiece'. Ignores invalid indexes, for e.g.:
+--
+--   @forall s. hindRead (-1) s == return ()@
+--
+hintRead :: PieceIx -> Storage -> IO ()
+hintRead _pix Storage {..} = return ()
+
+-- | Zero-copy version of readPiece. Can be used only with 'ReadOnly'
+-- storages.
+unsafeReadPiece :: PieceIx -> Storage -> IO (Piece BL.ByteString)
+unsafeReadPiece pix s @ Storage {..}
+  | not (isValidIx pix s) = throwIO (InvalidIndex pix)
+  | otherwise = return $ Piece pix (unsafeReadBytes offset sz fileMap)
+  where
+    offset = fromIntegral pix * fromIntegral pieceLen
+    sz     = fromIntegral pieceLen
+
+-- | Stream storage pieces from first to the last.
+sourceStorage :: Storage -> Source IO (Piece BL.ByteString)
+sourceStorage s = go 0
+  where
+    go pix
+        | pix < totalPieces s = do
+          piece <- liftIO $ readPiece pix s
+          liftIO $ hintRead (succ pix) s
+          yield piece
+          go (succ pix)
+        |       otherwise     = return ()
+
+-- | Write stream of pieces to the storage. Fail if storage is 'ReadOnly'.
+sinkStorage :: Storage -> Sink (Piece BL.ByteString) IO ()
+sinkStorage s = do
+  awaitForever $ \ piece ->
+    liftIO $ writePiece piece s
+
+-- | This function can be used to generate 'InfoDict' from a set of
+-- opened files.
+genPieceInfo :: Storage -> IO PieceInfo
+genPieceInfo s = do
+  hashes <- sourceStorage s $= C.map hashPiece $$ C.sinkLbs
+  return $ PieceInfo (pieceLen s) (HashList (BL.toStrict hashes))
+
+-- | Verify specific piece using infodict hash list.
+verifyPiece :: Storage -> PieceInfo -> PieceIx -> IO Bool
+verifyPiece s pinfo pix = do
+  piece <- unsafeReadPiece pix s
+  return $! checkPieceLazy pinfo piece
+
+-- | Verify storage.
+--
+--   Throws 'InvalidSize' if piece info size do not match with storage
+--   piece size.
+--
+getBitfield :: Storage -> PieceInfo -> IO Bitfield
+getBitfield s @ Storage {..} pinfo @ PieceInfo {..}
+  | pieceLen /= piPieceLength = throwIO (InvalidSize piPieceLength)
+  | otherwise = M.foldM checkPiece (BF.haveNone total) [0..total - 1]
+  where
+    total = totalPieces s
+
+    checkPiece :: Bitfield -> PieceIx -> IO Bitfield
+    checkPiece bf pix = do
+      valid <- verifyPiece s pinfo pix
+      return $ if valid then BF.insert pix bf else bf
diff --git a/dht/bittorrent/src/System/Torrent/Tree.hs b/dht/bittorrent/src/System/Torrent/Tree.hs
new file mode 100644
index 00000000..41cfb360
--- /dev/null
+++ b/dht/bittorrent/src/System/Torrent/Tree.hs
@@ -0,0 +1,83 @@
+-- |
+--   Copyright   :  (c) Sam Truzjan 2013
+--   License     :  BSD3
+--   Maintainer  :  pxqr.sta@gmail.com
+--   Stability   :  experimental
+--   Portability :  portable
+--
+--   Directory tree can be used to easily manipulate file layout info.
+--
+{-# LANGUAGE FlexibleInstances  #-}
+{-# LANGUAGE TemplateHaskell    #-}
+{-# LANGUAGE DeriveDataTypeable #-}
+module System.Torrent.Tree
+       ( -- * Directory tree
+         DirTree (..)
+
+         -- * Construction
+       , build
+
+         -- * Query
+       , System.Torrent.Tree.lookup
+       , lookupDir
+       , fileNumber
+       , dirNumber
+       ) where
+
+import Data.ByteString as BS
+import Data.ByteString.Char8 as BC
+import Data.Foldable
+import Data.List as L
+import Data.Map  as M
+import Data.Monoid
+
+import Data.Torrent
+
+
+-- | 'DirTree' is more convenient form of 'LayoutInfo'.
+data DirTree a = Dir  { children :: Map ByteString (DirTree a) }
+               | File { node     :: FileInfo a                 }
+                 deriving Show
+
+-- | Build directory tree from a list of files.
+build :: LayoutInfo -> DirTree ()
+build SingleFile {liFile = FileInfo {..}} = Dir
+    { children = M.singleton fiName (File fi) }
+  where
+    fi = FileInfo fiLength fiMD5Sum ()
+build MultiFile {..} = Dir $ M.singleton liDirName files
+  where
+    files = Dir $ M.fromList $ L.map mkFileEntry liFiles
+    mkFileEntry FileInfo {..} = (L.head fiName, ent) -- TODO FIXME
+      where
+        ent = File $ FileInfo fiLength fiMD5Sum ()
+
+--decompress :: DirTree () -> [FileInfo ()]
+--decompress = undefined
+
+-- TODO pretty print
+
+-- | Lookup file by path.
+lookup :: [FilePath] -> DirTree a -> Maybe (DirTree a)
+lookup []        t      = Just t
+lookup (p : ps) (Dir m) | Just subTree <- M.lookup (BC.pack p) m
+                        = System.Torrent.Tree.lookup ps subTree
+lookup _         _      = Nothing
+
+-- | Lookup directory by path.
+lookupDir :: [FilePath] -> DirTree a -> Maybe [(ByteString, DirTree a)]
+lookupDir ps d = do
+  subTree <- System.Torrent.Tree.lookup ps d
+  case subTree of
+    File _  -> Nothing
+    Dir  es -> Just $ M.toList es
+
+-- | Get total count of files in directory and subdirectories.
+fileNumber :: DirTree a -> Sum Int
+fileNumber File {..} = Sum 1
+fileNumber Dir  {..} = foldMap fileNumber children
+
+-- | Get total count of directories in the directory and subdirectories.
+dirNumber :: DirTree a -> Sum Int
+dirNumber File {..} = Sum 0
+dirNumber Dir  {..} = Sum 1 <> foldMap dirNumber children