2 files changed, 23 insertions, 4 deletions

diff --git a/src/Data/Torrent/Bitfield.hs b/src/Data/Torrent/Bitfield.hs
index 8cdae69f..b65f058b 100644
--- a/src/Data/Torrent/Bitfield.hs
+++ b/src/Data/Torrent/Bitfield.hs
@@ -57,6 +57,7 @@ module Data.Torrent.Bitfield
        , isSubsetOf
 
          -- ** Availability
+       , complement
        , Frequency
        , frequencies
        , rarest
@@ -194,8 +195,20 @@ findMax :: Bitfield -> PieceIx
 findMax = S.findMax . bfSet
 {-# INLINE findMax #-}
 
+-- | Check if all pieces from first bitfield present if the second bitfield
 isSubsetOf :: Bitfield -> Bitfield -> Bool
 isSubsetOf a b = bfSet a `S.isSubsetOf` bfSet b
+{-# INLINE isSubsetOf #-}
+
+-- | Resulting bitfield includes only missing pieces.
+complement :: Bitfield -> Bitfield
+complement Bitfield {..} = Bitfield
+  { bfSet  = uni `S.difference` bfSet
+  , bfSize = bfSize
+  }
+  where
+    Bitfield _ uni = haveAll bfSize
+{-# INLINE complement #-}
 
 {-----------------------------------------------------------------------
 -- Availability
diff --git a/src/System/Torrent/Storage.hs b/src/System/Torrent/Storage.hs
index a8b0bdc6..b5092b2e 100644
--- a/src/System/Torrent/Storage.hs
+++ b/src/System/Torrent/Storage.hs
@@ -28,6 +28,8 @@ module System.Torrent.Storage
        , withStorage
 
          -- * Query
+       , totalPieces
+       , verifyPiece
        , genPieceInfo
        , getBitfield
 
@@ -166,6 +168,10 @@ 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 = checkPieceLazy pinfo <$> readPiece pix s
+
 -- | Verify storage.
 --
 --   Throws 'InvalidSize' if piece info size do not match with storage
@@ -174,11 +180,11 @@ genPieceInfo s = do
 getBitfield :: Storage -> PieceInfo -> IO Bitfield
 getBitfield s @ Storage {..} pinfo @ PieceInfo {..}
   | pieceLen /= piPieceLength = throwIO (InvalidSize piPieceLength)
-  | otherwise = M.foldM verifyPiece (BF.haveNone total) [0..total - 1]
+  | otherwise = M.foldM checkPiece (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
+    checkPiece :: Bitfield -> PieceIx -> IO Bitfield
+    checkPiece bf pix = do
+      valid <- verifyPiece s pinfo pix
       return $ if valid then BF.insert pix bf else bf