From 16626766ee4d82ff7b84e3a9fc6c1341765bed2a Mon Sep 17 00:00:00 2001
From: David Frank <david.frank@dfinity.org>
Date: Sun, 24 Nov 2024 22:48:13 +0000
Subject: [PATCH 1/2] Fix Vec Iter::count() after next_back

---
 src/base_vec.rs  | 6 +-----
 src/vec/tests.rs | 7 +++++++
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/base_vec.rs b/src/base_vec.rs
index cc5b102f..f4e68b65 100644
--- a/src/base_vec.rs
+++ b/src/base_vec.rs
@@ -379,11 +379,7 @@ where
     }
 
     fn count(self) -> usize {
-        let n = self.vec.len().saturating_sub(self.range.start);
-        if n > usize::MAX as u64 {
-            panic!("The number of items in the vec {n} does not fit into usize");
-        }
-        n as usize
+        self.range.count()
     }
 
     fn nth(&mut self, n: usize) -> Option<T> {
diff --git a/src/vec/tests.rs b/src/vec/tests.rs
index d91e9ed0..088404bc 100644
--- a/src/vec/tests.rs
+++ b/src/vec/tests.rs
@@ -227,6 +227,13 @@ 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);
+
+    {
+        assert_eq!(sv.len(), 3);
+        let mut iter = sv.iter();
+        iter.next_back();
+        assert_eq!(iter.count(), 2);
+    }
 }
 
 // A struct with a bugg implementation of storable where the max_size can

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 2/2] 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);
     }
 }