Keep track inflight set in assembler

1 files changed, 39 insertions, 56 deletions

diff --git a/src/Network/BitTorrent/Exchange/Assembler.hs b/src/Network/BitTorrent/Exchange/Assembler.hs
index a490d2dc..b1a80c9b 100644
--- a/src/Network/BitTorrent/Exchange/Assembler.hs
+++ b/src/Network/BitTorrent/Exchange/Assembler.hs
@@ -1,82 +1,65 @@
 module Network.BitTorrent.Exchange.Assembler
        ( Assembler
-       , insert
-
-         -- * Query
-       , pendingPieces
-       , completeBlocks
-       , incompleteBlocks
        ) where
 
+import Control.Applicative
 import Data.IntMap.Strict as IM
+import Data.List as L
 import Data.Maybe
+
 import Data.Torrent.Piece
-import Data.Torrent.Block
+import Network.BitTorrent.Exchange.Block
+
 
 type PieceMap = IntMap
 
--- TODO move to Data.Torrent.Piece ?
--- assembler is also a block selector?
 data Assembler a = Assembler
-  { piecePending :: PieceMap [Block a]
+  { inflightSet  :: PieceMap [BlockIx]
+  , piecePending :: PieceMap [Block a]
   , pieceInfo    :: PieceInfo
   }
 
-
 data Result a
-  = Assembled (Piece a)
-  | Failed     PieceIx
+  = Assembled   (Piece a)
+  | Corrupted    PieceIx
+  | NotRequested PieceIx
+  | Overlapped   BlockIx
 
--- | You should check if a returned by peer block is actually have
--- been requested and in-flight. This is needed to avoid "I send
--- random corrupted block" attacks.
-insert :: Block a -> Assembler a -> (Assembler a, Maybe (Result a))
-insert Block {..} Assembler {..} = undefined
---  updateWithKey bixPiece
+empty :: PieceInfo -> Assembler a
+empty = Assembler IM.empty IM.empty
 
-pendingPieces :: Assembler a -> [PieceIx]
-pendingPieces Assembler {..} = keys piecePending
+inflightPieces :: Assembler a -> [PieceIx]
+inflightPieces Assembler {..} = keys piecePending
 
-completeBlocks :: Assembler a -> PieceIx -> [Block a]
-completeBlocks Assembler {..} pix = fromMaybe [] $ IM.lookup pix piecePending
+completeBlocks :: PieceIx -> Assembler a -> [Block a]
+completeBlocks pix Assembler {..}  = fromMaybe [] $ IM.lookup pix piecePending
 
-incompleteBlocks :: Assembler a -> PieceIx -> [BlockIx]
+incompleteBlocks :: PieceIx -> Assembler a -> [BlockIx]
 incompleteBlocks = undefined
 
--- TODO merge BlockSelector with Assembler?
-data BlockSelector a = BlockSelector
-  { assembler   :: Assembler a -- we do not select already transfered blocks
-  , inflightSet :: Set BlockIx -- we do not select blocks in flight
-  }
-
-insert :: BlockSelector -> (BlockSelector a, Maybe (Result a))
-insert = undefined
-
+nextBlock :: Assembler a -> Maybe (Assembler a, BlockIx)
+nextBlock Assembler {..} = undefined
 
-data StorageAdapter = StorageAdapter
-  { bitfield :: Bitfield
-  , requestQ :: Queue PieceIx
-  }
--- we do select 'incompleteBlocks' that is not in flight
+allowPiece :: PieceIx -> Assembler a -> Assembler a
+allowPiece = undefined
 
---assembler :: Assembler -> Conduit (Block a) (Result a)
---assembler = undefined
+insert' :: Block a -> [Block a] -> [Block a]
+insert' a (x : xs) = undefined
 
--- by priority
---   foreign request queue (max queue size)
---   assembler block information (max queue size)
---   selection strategies (bitfields)
+insertBlock :: Block a -> [Block a] -> Either [Block a] (Piece a)
+insertBlock  = undefined
 
--- when in flight queue is full we do not use selector
---   in flight queue information (max queue size)
-
--- piece select is used when
-data PieceSelector = Selector
-  { forceQueue        :: TVar (Queue PieceIx)
-  , forcePendingQueue :: TVar (Queue PieceIx)
-  , assembler :: TVar Assembler
-  , strategy  :: Bool
-  }
+inserta :: Block a -> PieceMap [Block a] -> (PieceMap [Block a], Maybe (Piece a))
+inserta = undefined
 
-select :: Selector -> (Selector, PieceIx)
-select = undefined
\ No newline at end of file
+-- | You should check if a returned by peer block is actually have
+-- been requested and in-flight. This is needed to avoid "I send
+-- random corrupted block" attacks.
+insert :: Block a -> Assembler a -> (Assembler a, Maybe (Result a))
+insert blk @ Block {..} a @ Assembler {..}
+  = let (pending, mpiece) = inserta blk piecePending
+    in (Assembler inflightSet pending pieceInfo, f <$> mpiece)
+  where
+    f p = undefined
+--      | checkPieceLazy pieceInfo p = Assembled p
+--      |          otherwise         = Corrupted ixPiece