From 8346085fb9e8df7e5161cea6b18997bde22ca273 Mon Sep 17 00:00:00 2001
From: Andreas Fackler <AndreasFackler@gmx.de>
Date: Mon, 30 Jul 2018 15:02:20 +0200
Subject: [PATCH] Fix removal of terminated CS instances in HB.

---
 src/honey_badger.rs | 31 ++++++++++++++++---------------
 1 file changed, 16 insertions(+), 15 deletions(-)

diff --git a/src/honey_badger.rs b/src/honey_badger.rs
index a48eae1..1b19f55 100644
--- a/src/honey_badger.rs
+++ b/src/honey_badger.rs
@@ -252,7 +252,7 @@ where
             cs.handle_message(sender_id, message)?
         };
         let step = self.process_output(cs_step, epoch)?;
-        self.remove_terminated(epoch);
+        self.remove_terminated();
         Ok(step)
     }
 
@@ -558,20 +558,21 @@ where
     }
 
     /// Removes all `CommonSubset` instances from _past_ epochs that have terminated.
-    fn remove_terminated(&mut self, from_epoch: u64) {
-        for epoch in from_epoch..self.epoch {
-            if self
-                .common_subsets
-                .get(&epoch)
-                .map_or(false, CommonSubset::terminated)
-            {
-                debug!(
-                    "{:?} Epoch {} has terminated.",
-                    self.netinfo.our_uid(),
-                    epoch
-                );
-                self.common_subsets.remove(&epoch);
-            }
+    fn remove_terminated(&mut self) {
+        let terminated_epochs: Vec<u64> = self
+            .common_subsets
+            .iter()
+            .take_while(|&(epoch, _)| *epoch < self.epoch)
+            .filter(|&(_, cs)| cs.terminated())
+            .map(|(epoch, _)| *epoch)
+            .collect();
+        for epoch in terminated_epochs {
+            debug!(
+                "{:?} Epoch {} has terminated.",
+                self.netinfo.our_uid(),
+                epoch
+            );
+            self.common_subsets.remove(&epoch);
         }
     }
 }