Implement getBitfield and genPieceInfo functions

1 files changed, 13 insertions, 0 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 {