Correcting the rotate_nested_vector

- renamed for better expression of what it does
- updated the doc string of the function to not reference zip
- removed the limitation of hints returned, allowing the caller to set
  the number of items retrieved

Author: upjohnc

lightning-invoice/src/utils.rs

@@ -203,7 +203,9 @@ where
 		invoice = invoice.amount_milli_satoshis(amt);
 	}
 
-	for route_hint in select_phantom_hints(amt_msat, phantom_route_hints, logger) {
+	const MAX_HINTS: usize = 3;
+
+	for route_hint in select_phantom_hints(amt_msat, phantom_route_hints, logger).take(MAX_HINTS) {
 		invoice = invoice.private_route(route_hint);
 	}
 
@@ -268,46 +270,35 @@ where
 	// the hints across our real nodes we add one hint from each in turn until no node has any hints
 	// left (if one node has more hints than any other, these will accumulate at the end of the
 	// vector).
-	const MAX_HINTS: usize = 3;
-
-    zip_nested_vectors(phantom_hints, MAX_HINTS)
+    rotate_nested_vectors(phantom_hints)
 
 }
 
-// Zip multiple nested vectors
-// Similar to the zip method on a vector, it combines the vectors by index - zero index comes first
-// then second, etc.
-// The difference is that this function handles many or just one vector, whereas the builtin method
-// works on just two vectors.
-fn zip_nested_vectors<T: Clone>(vecs: Vec<Vec<T>>, result_size: usize) -> impl Iterator<Item = T> {
+// Draw items iteratively from multiple nested vectors.  The items are retrieved by index and
+// rotates through the vectors - first the zero index then the first index then second index, etc.
+fn rotate_nested_vectors<T: Clone>(vecs: Vec<Vec<T>>) -> impl Iterator<Item = T> {
     let max_vector_length: usize = vecs.iter().map(|x| x.len()).max().unwrap();
     let mut hint_index = 0;
     let mut vector_index = 0;
     let number_inner_vectors: usize = vecs.len();
-    let mut total_hints_returned = 0;
 
     core::iter::from_fn(move || loop {
-        let return_value = loop {
-            if total_hints_returned == result_size || hint_index == max_vector_length {
-                break None;
-            };
-            let hint_value =
-                if vecs[vector_index].len() != 0 && vecs[vector_index].len() > hint_index {
-                    Some(vecs[vector_index][hint_index].clone())
-                } else {
-                    None // no value retrieved - continue looping
-                };
-            vector_index += 1;
-            if hint_index < max_vector_length && vector_index == number_inner_vectors {
-                vector_index = 0;
-                hint_index += 1;
-            };
-            if !hint_value.is_none() {
-                total_hints_returned += 1;
-                break hint_value;
-            };
+        if hint_index == max_vector_length {
+            return None;
+        };
+        let hint_value = if vecs[vector_index].len() != 0 && vecs[vector_index].len() > hint_index {
+            Some(vecs[vector_index][hint_index].clone())
+        } else {
+            None // no value retrieved - continue looping
+        };
+        vector_index += 1;
+        if hint_index < max_vector_length && vector_index == number_inner_vectors {
+            vector_index = 0;
+            hint_index += 1;
+        };
+        if !hint_value.is_none() {
+            return hint_value;
         };
-        return return_value;
     })
 }
 
@@ -796,7 +787,7 @@ mod test {
 	use lightning::routing::router::{PaymentParameters, RouteParameters};
 	use lightning::util::test_utils;
 	use lightning::util::config::UserConfig;
-	use crate::utils::{create_invoice_from_channelmanager_and_duration_since_epoch, zip_nested_vectors};
+	use crate::utils::{create_invoice_from_channelmanager_and_duration_since_epoch, rotate_nested_vectors};
 	use std::collections::HashSet;
 
 	#[test]
@@ -1909,60 +1900,58 @@ mod test {
 
     #[test]
     fn test_zip_nested_vectors() {
-		const TEST_RESULT_SIZE: usize = 3;
-
 		// two nested vectors
         let a = vec![vec!["a0", "b0", "c0"], vec!["a1", "b1"]];
-        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+        let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
-        let expected = vec!["a0", "a1", "b0"];
+        let expected = vec!["a0", "a1", "b0", "b1", "c0"];
         assert_eq!(expected, result);
 
         // test single nested vector
         let a = vec![vec!["a0", "b0", "c0"]];
-        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+        let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
         let expected = vec!["a0", "b0", "c0"];
         assert_eq!(expected, result);
 
         // test second vector with only one element
         let a = vec![vec!["a0", "b0", "c0"], vec!["a1"]];
-        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+        let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
-        let expected = vec!["a0", "a1", "b0"];
+        let expected = vec!["a0", "a1", "b0", "c0"];
         assert_eq!(expected, result);
 
         // test three nestend vectors
         let a = vec![vec!["a0"], vec!["a1", "b1", "c1"], vec!["a2"]];
-        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+        let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
-        let expected = vec!["a0", "a1", "a2"];
+        let expected = vec!["a0", "a1", "a2", "b1", "c1"];
         assert_eq!(expected, result);
 
         // test single nested vector with a single value
         let a = vec![vec!["a0"]];
-        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+        let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
         let expected = vec!["a0"];
         assert_eq!(expected, result);
 
         // test single empty nested vector
         let a:Vec<Vec<&str>> = vec![vec![]];
-        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+        let result = rotate_nested_vectors(a).collect::<Vec<_>>();
         let expected:Vec<&str> = vec![];
 
         assert_eq!(expected, result);
 
         // test first nested vector is empty
         let a = vec![vec![], vec!["a1", "b1", "c1"]];
-        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+        let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
         let expected = vec!["a1", "b1", "c1"];
         assert_eq!(expected, result);
 
         // test two empty vectors
         let a:Vec<Vec<&str>> = vec![vec![], vec![]];
-        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+        let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
         let expected:Vec<&str> = vec![];
         assert_eq!(expected, result);
@@ -1974,21 +1963,21 @@ mod test {
             vec!["a1", "b1", "c1"],
             vec!["a2", "b2", "c2"],
         ];
-        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+        let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
-        let expected = vec!["a0", "a1", "a2"];
+        let expected = vec!["a0", "a1", "a2", "b0", "b1", "b2", "c0", "c1", "c2"];
         assert_eq!(expected, result);
 
         // test a filled vector between two empty vectors
         let a = vec![vec![], vec!["a1", "b1", "c1"], vec![]];
-        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+        let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
         let expected = vec!["a1", "b1", "c1"];
         assert_eq!(expected, result);
 
         // test an empty vector at the end of the vectors
         let a = vec![vec!["a0", "b0", "c0"], vec![]];
-        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+        let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
         let expected = vec!["a0", "b0", "c0"];
         assert_eq!(expected, result);
@@ -2002,15 +1991,15 @@ mod test {
             vec![],
         ];
 
-        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+        let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
-        let expected = vec!["a1", "a3", "b1"];
+        let expected = vec!["a1", "a3", "b1", "b3", "c1"];
         assert_eq!(expected, result);
 
         // test one element in the first nested vectore and two elements in the second nested
         // vector
         let a = vec![vec!["a0"], vec!["a1", "b1"]];
-        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+        let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
         let expected = vec!["a0", "a1", "b1"];
         assert_eq!(expected, result);