2 files changed, 37 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