Implemented more efficient version of nth() for EnumIter.

Peternator7 · Dec 13, 2019 · dbc7a25 · dbc7a25
1 parent 4c08bcf
commit dbc7a25
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 16 deletions.
diff --git a/strum_macros/src/enum_iter.rs b/strum_macros/src/enum_iter.rs
@@ -88,21 +88,24 @@ pub fn enum_iter_inner(ast: &syn::DeriveInput) -> TokenStream {
         impl #impl_generics Iterator for #iter_name #ty_generics #where_clause {
             type Item = #name #ty_generics;
 
-            fn next(&mut self) -> Option<#name #ty_generics> {
-                if self.idx + self.back_idx >= #variant_count {
-                    return None
-                }
-
-                let output = self.get(self.idx);
-
-                self.idx += 1;
-                output
+            fn next(&mut self) -> Option<Self::Item> {
+                self.nth(0)
             }
 
             fn size_hint(&self) -> (usize, Option<usize>) {
                 let t = #variant_count - self.idx - self.back_idx;
                 (t, Some(t))
             }
+
+            fn nth(&mut self, n: usize) -> Option<Self::Item> {
+                self.idx += n + 1;
+
+                if self.idx + self.back_idx > #variant_count {
+                    None
+                } else {
+                    self.get(self.idx - 1)
+                }   
+            }
         }
 
         impl #impl_generics ExactSizeIterator for #iter_name #ty_generics #where_clause {
@@ -112,15 +115,18 @@ pub fn enum_iter_inner(ast: &syn::DeriveInput) -> TokenStream {
         }
 
         impl #impl_generics DoubleEndedIterator for #iter_name #ty_generics #where_clause {
-            fn next_back(&mut self) -> Option<#name #ty_generics> {
-                if self.idx + self.back_idx >= #variant_count {
-                    return None
-                }
+            fn next_back(&mut self) -> Option<Self::Item> {
+                self.nth_back(0)
+            }
 
-                let output = self.get(#variant_count - self.back_idx - 1);
+            fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
+                self.back_idx += n + 1;
 
-                self.back_idx += 1;
-                output
+                if self.idx + self.back_idx > #variant_count {
+                    None
+                } else {
+                    self.get(#variant_count - self.back_idx)
+                }  
             }
         }
 

diff --git a/strum_tests/tests/enum_iter.rs b/strum_tests/tests/enum_iter.rs
@@ -65,6 +65,21 @@ fn len_test() {
     assert_eq!(0, i.len());
 }
 
+#[test]
+fn double_ended_len_test() {
+    let mut i = Complicated::<(), ()>::iter();
+    assert_eq!(3, i.len());
+    i.next_back();
+
+    assert_eq!(2, i.len());
+    i.next();
+
+    assert_eq!(1, i.len());
+    i.next_back();
+
+    assert_eq!(0, i.len());
+}
+
 #[test]
 fn clone_test() {
     let mut i = Week::iter();
@@ -129,4 +144,16 @@ fn take_from_both_sides_test() {
     assert_eq!(Some(Week::Thursday), iter.next_back());
     assert_eq!(None, iter.next());
     assert_eq!(None, iter.next_back());
+}
+
+#[test]
+fn take_nth_test() {
+    let mut iter = Week::iter();
+
+    assert_eq!(Some(Week::Tuesday), iter.nth(2));
+    assert_eq!(Some(Week::Saturday), iter.nth_back(0));
+    assert_eq!(Some(Week::Thursday), iter.nth_back(1));
+    assert_eq!(None, iter.nth(1));
+    assert_eq!(None, iter.next());
+    assert_eq!(None, iter.next_back());
 }