Implement getBitfield and genPieceInfo functions

3 files changed, 44 insertions, 5 deletions

diff --git a/src/Data/Torrent/Bitfield.hs b/src/Data/Torrent/Bitfield.hs
index 02a4c14f..8cdae69f 100644
--- a/src/Data/Torrent/Bitfield.hs
+++ b/src/Data/Torrent/Bitfield.hs
@@ -62,6 +62,7 @@ module Data.Torrent.Bitfield
        , rarest
 
          -- * Combine
+       , insert
        , union
        , intersection
        , difference
@@ -196,6 +197,10 @@ findMax = S.findMax . bfSet
 isSubsetOf :: Bitfield -> Bitfield -> Bool
 isSubsetOf a b = bfSet a `S.isSubsetOf` bfSet b
 
+{-----------------------------------------------------------------------
+-- Availability
+-----------------------------------------------------------------------}
+
 -- | Frequencies are needed in piece selection startegies which use
 -- availability quantity to find out the optimal next piece index to
 -- download.
@@ -240,6 +245,14 @@ rarest xs
     Combine
 -----------------------------------------------------------------------}
 
+insert :: PieceIx -> Bitfield -> Bitfield
+insert pix bf @ Bitfield {..}
+  | 0 <= pix && pix < bfSize = Bitfield
+    { bfSet  = S.insert pix bfSet
+    , bfSize = bfSize
+    }
+  | otherwise = bf
+
 -- | Find indices at least one peer have.
 union :: Bitfield -> Bitfield -> Bitfield
 union a b = {-# SCC union #-} Bitfield {
diff --git a/src/System/Torrent/Storage.hs b/src/System/Torrent/Storage.hs
index 8aa1aa99..a8b0bdc6 100644
--- a/src/System/Torrent/Storage.hs
+++ b/src/System/Torrent/Storage.hs
@@ -44,9 +44,12 @@ module System.Torrent.Storage
 
 import Control.Applicative
 import Control.Exception
+import Control.Monad as M
 import Control.Monad.Trans
 import Data.ByteString.Lazy as BL
-import Data.Conduit
+import Data.Conduit as C
+import Data.Conduit.Binary as C
+import Data.Conduit.List as C
 import Data.Typeable
 
 import Data.Torrent.Bitfield as BF
@@ -156,10 +159,26 @@ sinkStorage s = do
   awaitForever $ \ piece ->
     liftIO $ writePiece piece s
 
--- | TODO examples of use
+-- | This function can be used to generate 'InfoDict' from a set of
+-- opened files.
 genPieceInfo :: Storage -> IO PieceInfo
-genPieceInfo = undefined
+genPieceInfo s = do
+  hashes <- sourceStorage s $= C.map hashPiece $$ C.sinkLbs
+  return $ PieceInfo (pieceLen s) (HashList (BL.toStrict hashes))
 
--- | TODO examples of use
+-- | Verify storage.
+--
+--   Throws 'InvalidSize' if piece info size do not match with storage
+--   piece size.
+--
 getBitfield :: Storage -> PieceInfo -> IO Bitfield
-getBitfield = undefined
\ No newline at end of file
+getBitfield s @ Storage {..} pinfo @ PieceInfo {..}
+  | pieceLen /= piPieceLength = throwIO (InvalidSize piPieceLength)
+  | otherwise = M.foldM verifyPiece (BF.haveNone total) [0..total - 1]
+  where
+    total = totalPieces s
+
+    verifyPiece :: Bitfield -> PieceIx -> IO Bitfield
+    verifyPiece bf pix = do
+      valid <- checkPieceLazy pinfo <$> readPiece pix s
+      return $ if valid then BF.insert pix bf else bf
diff --git a/tests/System/Torrent/StorageSpec.hs b/tests/System/Torrent/StorageSpec.hs
index 2e94ccf1..40eaa9c7 100644
--- a/tests/System/Torrent/StorageSpec.hs
+++ b/tests/System/Torrent/StorageSpec.hs
@@ -8,6 +8,7 @@ import System.Directory
 import System.IO.Unsafe
 import Test.Hspec
 
+import Data.Torrent.Bitfield as BF
 import Data.Torrent.Layout
 import Data.Torrent.Piece
 import System.Torrent.Storage
@@ -86,3 +87,9 @@ spec = before createLayout $ do
         sourceStorage s $= C.map bzeroPiece $$ sinkStorage s
         b <- sourceStorage s $$ C.fold (\ b p -> b && isZeroPiece p) True
         b `shouldBe` True
+
+  describe "genPieceInfo" $ do
+    it "" $ do
+      withStorage ReadWrite psize layout $ \ s -> do
+        bf <- genPieceInfo s >>= getBitfield s
+        bf `shouldSatisfy` BF.full
\ No newline at end of file