Setting return type of PhantomHints to Iterator

- return value of the function select_phantom_hints is an iterator
  rather than vector
- adding utility function that zips nested vectors

Author: upjohnc

lightning-invoice/src/utils.rs

@@ -18,6 +18,7 @@ use lightning::util::logger::Logger;
 use secp256k1::PublicKey;
 use core::ops::Deref;
 use core::time::Duration;
+use core::iter::Iterator;
 
 /// Utility to create an invoice that can be paid to one of multiple nodes, or a "phantom invoice."
 /// See [`PhantomKeysManager`] for more information on phantom node payments.
@@ -229,7 +230,7 @@ where
 ///
 /// [`PhantomKeysManager`]: lightning::sign::PhantomKeysManager
 fn select_phantom_hints<L: Deref>(amt_msat: Option<u64>, phantom_route_hints: Vec<PhantomRouteHints>,
-	logger: L) -> Vec<RouteHint>
+	logger: L) -> impl Iterator<Item = RouteHint>
 where
 	L::Target: Logger,
 {
@@ -267,29 +268,47 @@ 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).
-	let mut invoice_hints: Vec<RouteHint> = Vec::new();
-	let mut hint_idx = 0;
+	const MAX_HINTS: usize = 3;
 
-	loop {
-		let mut remaining_hints = false;
+    zip_nested_vectors(phantom_hints, MAX_HINTS)
 
-		for hints in phantom_hints.iter() {
-			if invoice_hints.len() == 3 {
-				return invoice_hints
-			}
-
-			if hint_idx < hints.len() {
-				invoice_hints.push(hints[hint_idx].clone());
-				remaining_hints = true
-			}
-		}
-
-		if !remaining_hints {
-			return invoice_hints
-		}
+}
 
-		hint_idx +=1;
-	}
+// 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> {
+    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 my_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 !my_value.is_none() {
+                total_hints_returned += 1;
+                break my_value;
+            };
+        };
+        return return_value;
+    })
 }
 
 #[cfg(feature = "std")]
@@ -777,7 +796,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;
+	use crate::utils::{create_invoice_from_channelmanager_and_duration_since_epoch, zip_nested_vectors};
 	use std::collections::HashSet;
 
 	#[test]
@@ -1886,4 +1905,106 @@ mod test {
 			_ => panic!(),
 		}
 	}
+
+
+    #[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 expected = vec!["a0", "a1", "b0"];
+        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 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 expected = vec!["a0", "a1", "b0"];
+        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 expected = vec!["a0", "a1", "a2"];
+        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 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 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 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 expected:Vec<&str> = vec![];
+        assert_eq!(expected, result);
+
+        // test an empty vector amongst other filled vectors
+        let a = vec![
+            vec!["a0", "b0", "c0"],
+            vec![],
+            vec!["a1", "b1", "c1"],
+            vec!["a2", "b2", "c2"],
+        ];
+        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+
+        let expected = vec!["a0", "a1", "a2"];
+        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 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 expected = vec!["a0", "b0", "c0"];
+        assert_eq!(expected, result);
+
+        // test multiple empty vectors amongst multiple filled vectors
+        let a = vec![
+            vec![],
+            vec!["a1", "b1", "c1"],
+            vec![],
+            vec!["a3", "b3"],
+            vec![],
+        ];
+
+        let result = zip_nested_vectors(a, TEST_RESULT_SIZE).collect::<Vec<_>>();
+
+        let expected = vec!["a1", "a3", "b1"];
+        assert_eq!(expected, result);
+    }
 }