Merge pull request #5 from cuviper/iter-opt

Optimize the iterators and release 0.1.1
indexmap-rs · Jan 30, 2025 · 6b9915e · 6b9915e
2 parents 2d9e930 + 1100437
commit 6b9915e
Show file tree

Hide file tree

Showing 4 changed files with 88 additions and 27 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "ringmap"
 edition = "2021"
-version = "0.1.0"
+version = "0.1.1"
 documentation = "https://docs.rs/ringmap/"
 repository = "https://github.com/indexmap-rs/ringmap"
 license = "Apache-2.0 OR MIT"

diff --git a/RELEASES.md b/RELEASES.md
@@ -1,5 +1,9 @@
 # Releases
 
+## 0.1.1
+
+- Optimized the branch behavior of the iterators.
+
 ## 0.1.0
 
 - Initial release, based on `indexmap v2.7.1`.
diff --git a/src/macros.rs b/src/macros.rs
@@ -98,6 +98,13 @@ macro_rules! iterator_methods {
             self.next_back()
         }
 
+        fn fold<Acc, F>(self, acc: Acc, f: F) -> Acc
+        where
+            F: FnMut(Acc, Self::Item) -> Acc,
+        {
+            self.iter.map($map_elt).fold(acc, f)
+        }
+
         fn collect<C>(self) -> C
         where
             C: FromIterator<Self::Item>,
@@ -120,6 +127,13 @@ macro_rules! double_ended_iterator_methods {
         fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
             self.iter.nth_back(n).map($map_elt)
         }
+
+        fn rfold<Acc, F>(self, acc: Acc, f: F) -> Acc
+        where
+            F: FnMut(Acc, Self::Item) -> Acc,
+        {
+            self.iter.map($map_elt).rfold(acc, f)
+        }
     };
 }
 

diff --git a/src/map/iter.rs b/src/map/iter.rs
@@ -2,11 +2,10 @@ use super::core::RingMapCore;
 use super::{Bucket, Entries, RingMap};
 
 use alloc::collections::vec_deque::{self, VecDeque};
-use core::fmt;
 use core::hash::{BuildHasher, Hash};
 use core::iter::FusedIterator;
 use core::ops::{Index, RangeBounds};
-use core::slice;
+use core::{fmt, mem, slice};
 
 impl<'a, K, V, S> IntoIterator for &'a RingMap<K, V, S> {
     type Item = (&'a K, &'a V);
@@ -64,7 +63,12 @@ impl<'a, K, V> Iterator for Buckets<'a, K, V> {
     fn next(&mut self) -> Option<Self::Item> {
         match self.head.next() {
             next @ Some(_) => next,
-            None => self.tail.next(),
+            None => {
+                // Swap so the rest is found on the first branch next time.
+                // (Like `VecDeque` does in its own iterators.)
+                mem::swap(&mut self.head, &mut self.tail);
+                self.head.next()
+            }
         }
     }
 
@@ -78,20 +82,26 @@ impl<'a, K, V> Iterator for Buckets<'a, K, V> {
     }
 
     fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
-        if n < self.head.len() {
-            return self.head.nth(n);
-        }
-        if self.head.len() > 0 {
+        if n >= self.head.len() {
             n -= self.head.len();
             self.head = [].iter();
+            mem::swap(&mut self.head, &mut self.tail);
         }
-        self.tail.nth(n)
+        self.head.nth(n)
     }
 
     fn last(mut self) -> Option<Self::Item> {
         self.next_back()
     }
 
+    fn fold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
+    where
+        F: FnMut(Acc, Self::Item) -> Acc,
+    {
+        acc = self.head.fold(acc, &mut f);
+        self.tail.fold(acc, &mut f)
+    }
+
     fn collect<C>(self) -> C
     where
         C: FromIterator<Self::Item>,
@@ -104,19 +114,30 @@ impl<K, V> DoubleEndedIterator for Buckets<'_, K, V> {
     fn next_back(&mut self) -> Option<Self::Item> {
         match self.tail.next_back() {
             next @ Some(_) => next,
-            None => self.head.next_back(),
+            None => {
+                // Swap so the rest is found on the first branch next time.
+                // (Like `VecDeque` does in its own iterators.)
+                mem::swap(&mut self.head, &mut self.tail);
+                self.tail.next_back()
+            }
         }
     }
 
     fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> {
-        if n < self.tail.len() {
-            return self.tail.nth_back(n);
-        }
-        if self.tail.len() > 0 {
+        if n >= self.tail.len() {
             n -= self.tail.len();
             self.tail = [].iter();
+            mem::swap(&mut self.head, &mut self.tail);
         }
-        self.head.nth_back(n)
+        self.tail.nth_back(n)
+    }
+
+    fn rfold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
+    where
+        F: FnMut(Acc, Self::Item) -> Acc,
+    {
+        acc = self.tail.rfold(acc, &mut f);
+        self.head.rfold(acc, &mut f)
     }
 }
 
@@ -180,7 +201,12 @@ impl<'a, K, V> Iterator for BucketsMut<'a, K, V> {
     fn next(&mut self) -> Option<Self::Item> {
         match self.head.next() {
             next @ Some(_) => next,
-            None => self.tail.next(),
+            None => {
+                // Swap so the rest is found on the first branch next time.
+                // (Like `VecDeque` does in its own iterators.)
+                mem::swap(&mut self.head, &mut self.tail);
+                self.head.next()
+            }
         }
     }
 
@@ -194,20 +220,26 @@ impl<'a, K, V> Iterator for BucketsMut<'a, K, V> {
     }
 
     fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
-        if n < self.head.len() {
-            return self.head.nth(n);
-        }
-        if self.head.len() > 0 {
+        if n >= self.head.len() {
             n -= self.head.len();
             self.head = [].iter_mut();
+            mem::swap(&mut self.head, &mut self.tail);
         }
-        self.tail.nth(n)
+        self.head.nth(n)
     }
 
     fn last(mut self) -> Option<Self::Item> {
         self.next_back()
     }
 
+    fn fold<Acc, F>(self, acc: Acc, mut f: F) -> Acc
+    where
+        F: FnMut(Acc, Self::Item) -> Acc,
+    {
+        let acc = self.head.fold(acc, &mut f);
+        self.tail.fold(acc, &mut f)
+    }
+
     fn collect<C>(self) -> C
     where
         C: FromIterator<Self::Item>,
@@ -220,19 +252,30 @@ impl<K, V> DoubleEndedIterator for BucketsMut<'_, K, V> {
     fn next_back(&mut self) -> Option<Self::Item> {
         match self.tail.next_back() {
             next @ Some(_) => next,
-            None => self.head.next_back(),
+            None => {
+                // Swap so the rest is found on the first branch next time.
+                // (Like `VecDeque` does in its own iterators.)
+                mem::swap(&mut self.head, &mut self.tail);
+                self.tail.next_back()
+            }
         }
     }
 
     fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> {
-        if n < self.tail.len() {
-            return self.tail.nth_back(n);
-        }
-        if self.tail.len() > 0 {
+        if n >= self.tail.len() {
             n -= self.tail.len();
             self.tail = [].iter_mut();
+            mem::swap(&mut self.head, &mut self.tail);
         }
-        self.head.nth_back(n)
+        self.tail.nth_back(n)
+    }
+
+    fn rfold<Acc, F>(self, acc: Acc, mut f: F) -> Acc
+    where
+        F: FnMut(Acc, Self::Item) -> Acc,
+    {
+        let acc = self.tail.rfold(acc, &mut f);
+        self.head.rfold(acc, &mut f)
     }
 }