Merge branch 'minify-histo' into next

jlouis · May 24, 2010 · 0f6c19f · 0f6c19f
2 parents 8525df4 + 956aadf
commit 0f6c19f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/src/Data/PendingSet.hs b/src/Data/PendingSet.hs
@@ -2,6 +2,7 @@ module Data.PendingSet
     ( PendingSet
     , Data.PendingSet.empty
     , Data.PendingSet.size
+    , remove
     , have
     , unhave
     , haves
@@ -33,10 +34,14 @@ have pn = PendingSet . alter f pn . unPS
 -- | A Peer does not have a given piece anymore (TODO: Not used in practice)
 unhave :: PieceNum -> PendingSet -> PendingSet
 unhave pn = PendingSet . alter f pn . unPS
-  where f Nothing  =  error "Data.PendingSet.unhave"
+  where f Nothing  =  Nothing
         f (Just 1) = Nothing
         f (Just x) = Just (x-1)
 
+-- | Remove a piece from the histogram. Used when it completes
+remove :: PieceNum -> PendingSet -> PendingSet
+remove pn = PendingSet . delete pn . unPS
+
 -- | Add all pieces in a bitfield
 haves :: [PieceNum] -> PendingSet -> PendingSet
 haves pns = flip (foldl f) pns

diff --git a/src/Process/PieceMgr.hs b/src/Process/PieceMgr.hs
@@ -225,7 +225,9 @@ peerHave idxs tmv = do
     inp <- gets inProgress
     interesting <- filterM (mem ps inp) idxs
     liftIO . atomically $ putTMVar tmv interesting
-    modify (\db -> db { histogram = PendS.haves idxs (histogram db)})
+    if null interesting
+        then return ()
+        else modify (\db -> db { histogram = PendS.haves interesting (histogram db)})
   where mem ps inp p = do
             q <- PS.member p ps
             if q
@@ -234,7 +236,8 @@ peerHave idxs tmv = do
 
 
 peerUnhave :: [PieceNum] -> Process CF ST ()
-peerUnhave idxs = modify (\db -> db { histogram = PendS.unhaves idxs (histogram db)})
+peerUnhave idxs =
+    modify (\db -> db { histogram = PendS.unhaves idxs (histogram db)})
 
 endgameBroadcast :: PieceNum -> Block -> Process CF ST ()
 endgameBroadcast pn blk = {-# SCC "endgameBroadCast" #-} do
@@ -275,7 +278,8 @@ createPieceDb mmap pmap = do
 completePiece :: PieceNum -> PieceMgrProcess ()
 completePiece pn = do
     PS.insert pn =<< gets donePiece
-    modify (\db -> db { inProgress = M.delete pn (inProgress db) })
+    modify (\db -> db { inProgress = M.delete pn (inProgress db),
+                        histogram  = PendS.remove pn (histogram db )})
 
 -- | Handle torrent completion
 checkFullCompletion :: PieceMgrProcess ()