Fix iterators

Gianmarco Garrisi · Gianmarco Garrisi · commit c09aac20f1cf · 2025-10-15T22:10:09.000+02:00
diff --git a/src/core_iterators.rs b/src/core_iterators.rs
@@ -51,6 +51,9 @@ impl<'a, I: 'a, P: 'a> Iterator for Drain<'a, I, P> {
     fn next(&mut self) -> Option<(I, P)> {
         self.iter.next()
     }
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
+    }
 }
 
 impl<I, P> DoubleEndedIterator for Drain<'_, I, P> {
@@ -80,6 +83,9 @@ impl<'a, I: 'a, P: 'a> Iterator for Iter<'a, I, P> {
     fn next(&mut self) -> Option<(&'a I, &'a P)> {
         self.iter.next()
     }
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
+    }
 }
 
 impl<I, P> DoubleEndedIterator for Iter<'_, I, P> {
@@ -109,6 +115,9 @@ impl<I, P> Iterator for IntoIter<I, P> {
     fn next(&mut self) -> Option<(I, P)> {
         self.iter.next()
     }
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
+    }
 }
 
 impl<I, P> DoubleEndedIterator for IntoIter<I, P> {
diff --git a/src/double_priority_queue/iterators.rs b/src/double_priority_queue/iterators.rs
@@ -141,6 +141,7 @@ where
 {
     pq: &'a mut DoublePriorityQueue<I, P, H>,
     pos: usize,
+    pos_back: usize,
 }
 
 #[cfg(not(feature = "std"))]
@@ -150,14 +151,20 @@ where
 {
     pq: &'a mut DoublePriorityQueue<I, P, H>,
     pos: usize,
+    pos_back: usize,
 }
 
 impl<'a, I: 'a, P: 'a, H: 'a> IterMut<'a, I, P, H>
 where
     P: Ord,
 {
     pub(crate) fn new(pq: &'a mut DoublePriorityQueue<I, P, H>) -> Self {
-        IterMut { pq, pos: 0 }
+        let pos_back = pq.len();
+        IterMut {
+            pq,
+            pos: 0,
+            pos_back,
+        }
     }
 }
 
@@ -169,6 +176,9 @@ where
     type Item = (&'a mut I, &'a mut P);
     fn next(&mut self) -> Option<Self::Item> {
         use indexmap::map::MutableKeys;
+        if self.pos == self.pos_back {
+            return None;
+        }
         let r: Option<(&'a mut I, &'a mut P)> = self
             .pq
             .store
@@ -179,6 +189,11 @@ where
         self.pos += 1;
         r
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.pos_back - self.pos;
+        (len, Some(len))
+    }
 }
 
 impl<'a, I: 'a, P: 'a, H: 'a> DoubleEndedIterator for IterMut<'a, I, P, H>
@@ -188,15 +203,16 @@ where
 {
     fn next_back(&mut self) -> Option<Self::Item> {
         use indexmap::map::MutableKeys;
-        let r: Option<(&'a mut I, &'a mut P)> = self
-            .pq
+        if self.pos == self.pos_back {
+            return None;
+        }
+        self.pos_back -= 1;
+        self.pq
             .store
             .map
-            .get_index_mut2(self.pos)
+            .get_index_mut2(self.pos_back)
             .map(|(i, p)| (i as *mut I, p as *mut P))
-            .map(|(i, p)| unsafe { (i.as_mut().unwrap(), p.as_mut().unwrap()) });
-        self.pos -= 1;
-        r
+            .map(|(i, p)| unsafe { (i.as_mut().unwrap(), p.as_mut().unwrap()) })
     }
 }
 
@@ -205,9 +221,6 @@ where
     P: Ord,
     H: BuildHasher,
 {
-    fn len(&self) -> usize {
-        self.pq.len()
-    }
 }
 
 impl<I, P, H> FusedIterator for IterMut<'_, I, P, H>
@@ -258,6 +271,10 @@ where
     fn next(&mut self) -> Option<(I, P)> {
         self.pq.pop_min()
     }
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.pq.len();
+        (len, Some(len))
+    }
 }
 
 impl<I, P, H> DoubleEndedIterator for IntoSortedIter<I, P, H>
@@ -269,13 +286,6 @@ where
     }
 }
 
-impl<I, P, H> ExactSizeIterator for IntoSortedIter<I, P, H>
-where
-    P: Ord,
-{
-    fn len(&self) -> usize {
-        self.pq.len()
-    }
-}
+impl<I, P, H> ExactSizeIterator for IntoSortedIter<I, P, H> where P: Ord {}
 
 impl<I, P, H> FusedIterator for IntoSortedIter<I, P, H> where P: Ord {}
diff --git a/src/priority_queue/iterators.rs b/src/priority_queue/iterators.rs
@@ -112,6 +112,10 @@ where
             }
         }
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (0, Some(self.pq.len()))
+    }
 }
 
 impl<'a, I: 'a, P: 'a, F, H: 'a> Drop for ExtractIf<'a, I, P, F, H>
@@ -141,6 +145,7 @@ where
 {
     pq: &'a mut PriorityQueue<I, P, H>,
     pos: usize,
+    pos_back: usize,
 }
 
 #[cfg(not(feature = "std"))]
@@ -150,14 +155,20 @@ where
 {
     pq: &'a mut PriorityQueue<I, P, H>,
     pos: usize,
+    pos_back: usize,
 }
 
 impl<'a, I: 'a, P: 'a, H: 'a> IterMut<'a, I, P, H>
 where
     P: Ord,
 {
     pub(crate) fn new(pq: &'a mut PriorityQueue<I, P, H>) -> Self {
-        IterMut { pq, pos: 0 }
+        let pos_back = pq.len();
+        IterMut {
+            pq,
+            pos: 0,
+            pos_back,
+        }
     }
 }
 
@@ -170,6 +181,10 @@ where
     fn next(&mut self) -> Option<Self::Item> {
         use indexmap::map::MutableKeys;
 
+        if self.pos == self.pos_back {
+            return None;
+        }
+
         let r: Option<(&'a mut I, &'a mut P)> = self
             .pq
             .store
@@ -180,6 +195,47 @@ where
         self.pos += 1;
         r
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.pos_back - self.pos;
+        (len, Some(len))
+    }
+}
+
+impl<'a, I: 'a, P: 'a, H: 'a> DoubleEndedIterator for IterMut<'a, I, P, H>
+where
+    P: Ord,
+    H: BuildHasher,
+{
+    fn next_back(&mut self) -> Option<Self::Item> {
+        use indexmap::map::MutableKeys;
+
+        if self.pos == self.pos_back {
+            return None;
+        }
+
+        self.pos_back -= 1;
+        self.pq
+            .store
+            .map
+            .get_index_mut2(self.pos_back)
+            .map(|(i, p)| (i as *mut I, p as *mut P))
+            .map(|(i, p)| unsafe { (i.as_mut().unwrap(), p.as_mut().unwrap()) })
+    }
+}
+
+impl<'a, I: 'a, P: 'a, H: 'a> ExactSizeIterator for IterMut<'a, I, P, H>
+where
+    P: Ord,
+    H: BuildHasher,
+{
+}
+
+impl<'a, I: 'a, P: 'a, H: 'a> FusedIterator for IterMut<'a, I, P, H>
+where
+    P: Ord,
+    H: BuildHasher,
+{
 }
 
 impl<'a, I: 'a, P: 'a, H: 'a> Drop for IterMut<'a, I, P, H>
@@ -213,4 +269,13 @@ where
     fn next(&mut self) -> Option<(I, P)> {
         self.pq.pop()
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.pq.len();
+        (len, Some(len))
+    }
 }
+
+impl<I, P, H> ExactSizeIterator for IntoSortedIter<I, P, H> where P: Ord {}
+
+impl<I, P, H> FusedIterator for IntoSortedIter<I, P, H> where P: Ord {}
diff --git a/tests/double_priority_queue.rs b/tests/double_priority_queue.rs
@@ -782,7 +782,9 @@ mod doublepq_tests {
         pq.push("b", 2);
         pq.push("f", 7);
 
-        assert_eq!(pq.iter().count(), 3);
+        let iter = pq.iter();
+        assert_eq!(iter.size_hint(), (3, Some(3)));
+        assert_eq!(iter.count(), 3);
     }
 
     #[test]
@@ -794,14 +796,20 @@ mod doublepq_tests {
         pq.push("g", 4);
         pq.push("h", 3);
 
-        for (i, p) in &mut pq {
+        let iter = pq.iter_mut();
+        assert_eq!(iter.size_hint(), (5, Some(5)));
+        assert_eq!(iter.len(), 5);
+
+        for (i, p) in iter {
             if *i < "f" {
                 *p += 18;
             }
         }
 
         assert_eq!(pq.pop_max(), Some(("b", 20)));
 
+        assert_eq!(pq.iter_mut().rev().count(), 4);
+
         /*
         // As expected, this does not compile
         let iter_mut = pq.iter_mut();
@@ -810,6 +818,19 @@ mod doublepq_tests {
         iter_mut.for_each(|(_, p)| {*p += 2}); */
     }
 
+    #[test]
+    fn iter_mut_next_back() {
+        let mut pq = DoublePriorityQueue::new();
+        pq.push("a", 1);
+        pq.push("b", 2);
+        pq.push("f", 7);
+        pq.push("g", 4);
+        pq.push("h", 3);
+
+        let mut iter = pq.iter_mut();
+        assert_eq!(iter.next_back(), Some((&mut "h", &mut 3)));
+    }
+
     #[test]
     fn retain() {
         #[derive(Hash, PartialEq, Eq, Debug)]
diff --git a/tests/priority_queue.rs b/tests/priority_queue.rs

Original file line number	Diff line number	Diff line change
`@@ -51,6 +51,9 @@ impl<'a, I: 'a, P: 'a> Iterator for Drain<'a, I, P> {`
`51`	`51`	`fn next(&mut self) -> Option<(I, P)> {`
`52`	`52`	`self.iter.next()`
`53`	`53`	`}`
	`54`	`+ fn size_hint(&self) -> (usize, Option<usize>) {`
	`55`	`+ self.iter.size_hint()`
	`56`	`+ }`
`54`	`57`	`}`
`55`	`58`
`56`	`59`	`impl<I, P> DoubleEndedIterator for Drain<'_, I, P> {`
`@@ -80,6 +83,9 @@ impl<'a, I: 'a, P: 'a> Iterator for Iter<'a, I, P> {`
`80`	`83`	`fn next(&mut self) -> Option<(&'a I, &'a P)> {`
`81`	`84`	`self.iter.next()`
`82`	`85`	`}`
	`86`	`+ fn size_hint(&self) -> (usize, Option<usize>) {`
	`87`	`+ self.iter.size_hint()`
	`88`	`+ }`
`83`	`89`	`}`
`84`	`90`
`85`	`91`	`impl<I, P> DoubleEndedIterator for Iter<'_, I, P> {`
`@@ -109,6 +115,9 @@ impl<I, P> Iterator for IntoIter<I, P> {`
`109`	`115`	`fn next(&mut self) -> Option<(I, P)> {`
`110`	`116`	`self.iter.next()`
`111`	`117`	`}`
	`118`	`+ fn size_hint(&self) -> (usize, Option<usize>) {`
	`119`	`+ self.iter.size_hint()`
	`120`	`+ }`
`112`	`121`	`}`
`113`	`122`
`114`	`123`	`impl<I, P> DoubleEndedIterator for IntoIter<I, P> {`
Original file line number	Diff line number	Diff line change
`@@ -141,6 +141,7 @@ where`
`141`	`141`	`{`
`142`	`142`	`pq: &'a mut DoublePriorityQueue<I, P, H>,`
`143`	`143`	`pos: usize,`
	`144`	`+ pos_back: usize,`
`144`	`145`	`}`
`145`	`146`
`146`	`147`	`#[cfg(not(feature = "std"))]`
`@@ -150,14 +151,20 @@ where`
`150`	`151`	`{`
`151`	`152`	`pq: &'a mut DoublePriorityQueue<I, P, H>,`
`152`	`153`	`pos: usize,`
	`154`	`+ pos_back: usize,`
`153`	`155`	`}`
`154`	`156`
`155`	`157`	`impl<'a, I: 'a, P: 'a, H: 'a> IterMut<'a, I, P, H>`
`156`	`158`	`where`
`157`	`159`	`P: Ord,`
`158`	`160`	`{`
`159`	`161`	`pub(crate) fn new(pq: &'a mut DoublePriorityQueue<I, P, H>) -> Self {`
`160`		`- IterMut { pq, pos: 0 }`
	`162`	`+ let pos_back = pq.len();`
	`163`	`+ IterMut {`
	`164`	`+ pq,`
	`165`	`+ pos: 0,`
	`166`	`+ pos_back,`
	`167`	`+ }`
`161`	`168`	`}`
`162`	`169`	`}`
`163`	`170`
`@@ -169,6 +176,9 @@ where`
`169`	`176`	`type Item = (&'a mut I, &'a mut P);`
`170`	`177`	`fn next(&mut self) -> Option<Self::Item> {`
`171`	`178`	`use indexmap::map::MutableKeys;`
	`179`	`+ if self.pos == self.pos_back {`
	`180`	`+ return None;`
	`181`	`+ }`
`172`	`182`	`let r: Option<(&'a mut I, &'a mut P)> = self`
`173`	`183`	`.pq`
`174`	`184`	`.store`
`@@ -179,6 +189,11 @@ where`
`179`	`189`	`self.pos += 1;`
`180`	`190`	`r`
`181`	`191`	`}`
	`192`	`+`
	`193`	`+ fn size_hint(&self) -> (usize, Option<usize>) {`
	`194`	`+ let len = self.pos_back - self.pos;`
	`195`	`+ (len, Some(len))`
	`196`	`+ }`
`182`	`197`	`}`
`183`	`198`
`184`	`199`	`impl<'a, I: 'a, P: 'a, H: 'a> DoubleEndedIterator for IterMut<'a, I, P, H>`
`@@ -188,15 +203,16 @@ where`
`188`	`203`	`{`
`189`	`204`	`fn next_back(&mut self) -> Option<Self::Item> {`
`190`	`205`	`use indexmap::map::MutableKeys;`
`191`		`- let r: Option<(&'a mut I, &'a mut P)> = self`
`192`		`- .pq`
	`206`	`+ if self.pos == self.pos_back {`
	`207`	`+ return None;`
	`208`	`+ }`
	`209`	`+ self.pos_back -= 1;`
	`210`	`+ self.pq`
`193`	`211`	`.store`
`194`	`212`	`.map`
`195`		`- .get_index_mut2(self.pos)`
	`213`	`+ .get_index_mut2(self.pos_back)`
`196`	`214`	`.map(\|(i, p)\| (i as mut I, p as mut P))`
`197`		`- .map(\|(i, p)\| unsafe { (i.as_mut().unwrap(), p.as_mut().unwrap()) });`
`198`		`- self.pos -= 1;`
`199`		`- r`
	`215`	`+ .map(\|(i, p)\| unsafe { (i.as_mut().unwrap(), p.as_mut().unwrap()) })`
`200`	`216`	`}`
`201`	`217`	`}`
`202`	`218`
`@@ -205,9 +221,6 @@ where`
`205`	`221`	`P: Ord,`
`206`	`222`	`H: BuildHasher,`
`207`	`223`	`{`
`208`		`- fn len(&self) -> usize {`
`209`		`- self.pq.len()`
`210`		`- }`
`211`	`224`	`}`
`212`	`225`
`213`	`226`	`impl<I, P, H> FusedIterator for IterMut<'_, I, P, H>`
`@@ -258,6 +271,10 @@ where`
`258`	`271`	`fn next(&mut self) -> Option<(I, P)> {`
`259`	`272`	`self.pq.pop_min()`
`260`	`273`	`}`
	`274`	`+ fn size_hint(&self) -> (usize, Option<usize>) {`
	`275`	`+ let len = self.pq.len();`
	`276`	`+ (len, Some(len))`
	`277`	`+ }`
`261`	`278`	`}`
`262`	`279`
`263`	`280`	`impl<I, P, H> DoubleEndedIterator for IntoSortedIter<I, P, H>`
`@@ -269,13 +286,6 @@ where`
`269`	`286`	`}`
`270`	`287`	`}`
`271`	`288`
`272`		`-impl<I, P, H> ExactSizeIterator for IntoSortedIter<I, P, H>`
`273`		`-where`
`274`		`- P: Ord,`
`275`		`-{`
`276`		`- fn len(&self) -> usize {`
`277`		`- self.pq.len()`
`278`		`- }`
`279`		`-}`
	`289`	`+impl<I, P, H> ExactSizeIterator for IntoSortedIter<I, P, H> where P: Ord {}`
`280`	`290`
`281`	`291`	`impl<I, P, H> FusedIterator for IntoSortedIter<I, P, H> where P: Ord {}`