Changing sort_and_filter_channels return to ExactSizeIterator and

removing built up vector in `select_phantom_hints`

- refactoring `rotate_nested_vectors` from Matt Corallo's suggestion

Author: upjohnc

lightning-invoice/src/utils.rs

@@ -236,69 +236,80 @@ fn select_phantom_hints<L: Deref>(amt_msat: Option<u64>, phantom_route_hints: Ve
 where
 	L::Target: Logger,
 {
-	let mut phantom_hints: Vec<Vec<RouteHint>> = Vec::new();
+	let mut phantom_hints: Vec<_> = Vec::new();
 
 	for PhantomRouteHints { channels, phantom_scid, real_node_pubkey } in phantom_route_hints {
 		log_trace!(logger, "Generating phantom route hints for node {}",
 			log_pubkey!(real_node_pubkey));
-		let mut route_hints: Vec<RouteHint> = sort_and_filter_channels(channels, amt_msat, &logger).collect();
+		let route_hints = sort_and_filter_channels(channels, amt_msat, &logger);
 
 		// If we have any public channel, the route hints from `sort_and_filter_channels` will be
 		// empty. In that case we create a RouteHint on which we will push a single hop with the
 		// phantom route into the invoice, and let the sender find the path to the `real_node_pubkey`
 		// node by looking at our public channels.
-		if route_hints.is_empty() {
-			route_hints.push(RouteHint(vec![]))
-		}
-		for route_hint in &mut route_hints {
-			route_hint.0.push(RouteHintHop {
-				src_node_id: real_node_pubkey,
-				short_channel_id: phantom_scid,
-				fees: RoutingFees {
-					base_msat: 0,
-					proportional_millionths: 0,
-				},
-				cltv_expiry_delta: MIN_CLTV_EXPIRY_DELTA,
-				htlc_minimum_msat: None,
-				htlc_maximum_msat: None,});
-		}
+		let empty_route_hints = route_hints.len() == 0;
+		let mut have_pushed_empty = false;
+		let route_hints = route_hints
+			.chain(core::iter::from_fn(move || {
+				if empty_route_hints && !have_pushed_empty {
+					// set flag of having handled the empty route_hints and ensure empty vector
+					// returned only once
+					have_pushed_empty = true;
+					Some(RouteHint(Vec::new()))
+				} else {
+					None
+				}
+			}))
+			.map(move |mut hint| {
+				hint.0.push(RouteHintHop {
+					src_node_id: real_node_pubkey,
+					short_channel_id: phantom_scid,
+					fees: RoutingFees {
+						base_msat: 0,
+						proportional_millionths: 0,
+					},
+					cltv_expiry_delta: MIN_CLTV_EXPIRY_DELTA,
+					htlc_minimum_msat: None,
+					htlc_maximum_msat: None,
+				});
+				hint
+			});
 
 		phantom_hints.push(route_hints);
+
 	}
 
 	// We have one vector per real node involved in creating the phantom invoice. To distribute
 	// 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).
 	rotate_nested_vectors(phantom_hints)
-
 }
 
-// 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();
-
-	core::iter::from_fn(move || loop {
-		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;
-		};
+/// 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, I: Iterator<Item = T>>(mut vecs: Vec<I>) -> impl Iterator<Item = T> {
+	let mut idx = 0;
+
+	core::iter::from_fn(move || {
+		let mut exhausted_vectors = 0;
+		loop {
+			if vecs.is_empty() {
+				return None;
+			}
+			let next_idx = idx % vecs.len();
+			let hint_opt = vecs[next_idx].next();
+			idx += 1;
+			if let Some(hint) = hint_opt {
+				return Some(hint);
+			}
+			// exhausted_vectors increase when the "next_idx" vector is exhausted
+			exhausted_vectors += 1;
+			// return None when all of the nested vectors are exhausted
+			if exhausted_vectors > vecs.len() {
+				return None;
+			}
+		}
 	})
 }
 
@@ -588,7 +599,7 @@ fn sort_and_filter_channels<L: Deref>(
 	channels: Vec<ChannelDetails>,
 	min_inbound_capacity_msat: Option<u64>,
 	logger: &L,
-) -> impl Iterator<Item = RouteHint>
+) -> impl ExactSizeIterator<Item = RouteHint>
 where
 	L::Target: Logger,
 {
@@ -1908,94 +1919,94 @@ mod test {
 	#[test]
 	fn test_zip_nested_vectors() {
 		// two nested vectors
-		let a = vec![vec!["a0", "b0", "c0"], vec!["a1", "b1"]];
+		let a = vec![vec!["a0", "b0", "c0"].into_iter(), vec!["a1", "b1"].into_iter()];
 		let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
 		let expected = vec!["a0", "a1", "b0", "b1", "c0"];
 		assert_eq!(expected, result);
 
 		// test single nested vector
-		let a = vec![vec!["a0", "b0", "c0"]];
+		let a = vec![vec!["a0", "b0", "c0"].into_iter()];
 		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 a = vec![vec!["a0", "b0", "c0"].into_iter(), vec!["a1"].into_iter()];
 		let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
 		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 a = vec![vec!["a0"].into_iter(), vec!["a1", "b1", "c1"].into_iter(), vec!["a2"].into_iter()];
 		let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
 		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 a = vec![vec!["a0"].into_iter()];
 		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 = rotate_nested_vectors(a).collect::<Vec<_>>();
+		let a:Vec<std::vec::IntoIter<&str>> = vec![vec![].into_iter()];
+		let result = rotate_nested_vectors(a).collect::<Vec<&str>>();
 		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 = rotate_nested_vectors(a).collect::<Vec<_>>();
+		let a:Vec<std::vec::IntoIter<&str>>= vec![vec![].into_iter(), vec!["a1", "b1", "c1"].into_iter()];
+		let result = rotate_nested_vectors(a).collect::<Vec<&str>>();
 
 		let expected = vec!["a1", "b1", "c1"];
 		assert_eq!(expected, result);
 
 		// test two empty vectors
-		let a:Vec<Vec<&str>> = vec![vec![], vec![]];
-		let result = rotate_nested_vectors(a).collect::<Vec<_>>();
+		let a:Vec<std::vec::IntoIter<&str>> = vec![vec![].into_iter(), vec![].into_iter()];
+		let result = rotate_nested_vectors(a).collect::<Vec<&str>>();
 
 		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"],
+			vec!["a0", "b0", "c0"].into_iter(),
+			vec![].into_iter(),
+			vec!["a1", "b1", "c1"].into_iter(),
+			vec!["a2", "b2", "c2"].into_iter(),
 		];
 		let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
 		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 a = vec![vec![].into_iter(), vec!["a1", "b1", "c1"].into_iter(), vec![].into_iter()];
 		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 a = vec![vec!["a0", "b0", "c0"].into_iter(), vec![].into_iter()];
 		let result = rotate_nested_vectors(a).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![],
+			vec![].into_iter(),
+			vec!["a1", "b1", "c1"].into_iter(),
+			vec![].into_iter(),
+			vec!["a3", "b3"].into_iter(),
+			vec![].into_iter(),
 		];
 
 		let result = rotate_nested_vectors(a).collect::<Vec<_>>();
@@ -2005,7 +2016,7 @@ mod test {
 
 		// 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 a = vec![vec!["a0"].into_iter(), vec!["a1", "b1"].into_iter()];
 		let result = rotate_nested_vectors(a).collect::<Vec<_>>();
 
 		let expected = vec!["a0", "a1", "b1"];