From dc210341d595f9abc934396e194f01d5102cd459 Mon Sep 17 00:00:00 2001
From: David Frank <david.frank@dfinity.org>
Date: Sun, 24 Nov 2024 23:50:52 +0000
Subject: [PATCH] Support the case when an element is popped while an iterator
 is alive

---
 src/base_vec.rs  |  5 ++++-
 src/vec/tests.rs | 16 +++++++++++++++-
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/src/base_vec.rs b/src/base_vec.rs
index f4e68b65..707cb303 100644
--- a/src/base_vec.rs
+++ b/src/base_vec.rs
@@ -379,7 +379,10 @@ where
     }
 
     fn count(self) -> usize {
-        self.range.count()
+        min(self.vec.len(), self.range.end)
+            .saturating_sub(self.range.start)
+            .try_into()
+            .expect("Cannot express count as usize")
     }
 
     fn nth(&mut self, n: usize) -> Option<T> {
diff --git a/src/vec/tests.rs b/src/vec/tests.rs
index 088404bc..7895c94d 100644
--- a/src/vec/tests.rs
+++ b/src/vec/tests.rs
@@ -227,11 +227,25 @@ fn test_iter() {
     assert_eq!(sv.iter().skip(3).count(), 0);
     assert_eq!(sv.iter().skip(4).count(), 0);
     assert_eq!(sv.iter().skip(usize::MAX).count(), 0);
+}
 
+#[test]
+fn test_iter_count() {
+    let sv = StableVec::<u64, M>::new(M::default()).unwrap();
+    sv.push(&1).unwrap();
+    sv.push(&2).unwrap();
+    sv.push(&3).unwrap();
+    sv.push(&4).unwrap();
+    {
+        let mut iter = sv.iter();
+        iter.next_back();
+        assert_eq!(iter.count(), 3);
+    }
     {
-        assert_eq!(sv.len(), 3);
         let mut iter = sv.iter();
         iter.next_back();
+        sv.pop(); // this pops the element that we iterated through on the previous line
+        sv.pop();
         assert_eq!(iter.count(), 2);
     }
 }