@@ -18,6 +18,7 @@ use lightning::util::logger::Logger;
1818use secp256k1:: PublicKey ;
1919use core:: ops:: Deref ;
2020use core:: time:: Duration ;
21+ use core:: iter:: Iterator ;
2122
2223/// Utility to create an invoice that can be paid to one of multiple nodes, or a "phantom invoice."
2324/// See [`PhantomKeysManager`] for more information on phantom node payments.
@@ -292,6 +293,33 @@ where
292293 }
293294}
294295
296+ /// Draw items iteratively from multiple iterators. The items are retrieved by index and
297+ /// rotates through the iterators - first the zero index then the first index then second index, etc.
298+ fn rotate_through_iterators < T , I : Iterator < Item = T > > ( mut vecs : Vec < I > ) -> impl Iterator < Item = T > {
299+ let mut iterations = 0 ;
300+
301+ core:: iter:: from_fn ( move || {
302+ let mut exhausted_iterators = 0 ;
303+ loop {
304+ if vecs. is_empty ( ) {
305+ return None ;
306+ }
307+ let next_idx = iterations % vecs. len ( ) ;
308+ iterations += 1 ;
309+ if let Some ( item) = vecs[ next_idx] . next ( ) {
310+ return Some ( item) ;
311+ }
312+ // exhausted_vectors increase when the "next_idx" vector is exhausted
313+ exhausted_iterators += 1 ;
314+ // The check for exhausted iterators gets reset to 0 after each yield of `Some()`
315+ // The loop will return None when all of the nested iterators are exhausted
316+ if exhausted_iterators == vecs. len ( ) {
317+ return None ;
318+ }
319+ }
320+ } )
321+ }
322+
295323#[ cfg( feature = "std" ) ]
296324/// Utility to construct an invoice. Generally, unless you want to do something like a custom
297325/// cltv_expiry, this is what you should be using to create an invoice. The reason being, this
@@ -777,7 +805,7 @@ mod test {
777805 use lightning:: routing:: router:: { PaymentParameters , RouteParameters } ;
778806 use lightning:: util:: test_utils;
779807 use lightning:: util:: config:: UserConfig ;
780- use crate :: utils:: create_invoice_from_channelmanager_and_duration_since_epoch;
808+ use crate :: utils:: { create_invoice_from_channelmanager_and_duration_since_epoch, rotate_through_iterators } ;
781809 use std:: collections:: HashSet ;
782810
783811 #[ test]
@@ -1886,4 +1914,111 @@ mod test {
18861914 _ => panic ! ( ) ,
18871915 }
18881916 }
1917+
1918+ #[ test]
1919+ fn test_rotate_through_iterators ( ) {
1920+ // two nested vectors
1921+ let a = vec ! [ vec![ "a0" , "b0" , "c0" ] . into_iter( ) , vec![ "a1" , "b1" ] . into_iter( ) ] ;
1922+ let result = rotate_through_iterators ( a) . collect :: < Vec < _ > > ( ) ;
1923+
1924+ let expected = vec ! [ "a0" , "a1" , "b0" , "b1" , "c0" ] ;
1925+ assert_eq ! ( expected, result) ;
1926+
1927+ // test single nested vector
1928+ let a = vec ! [ vec![ "a0" , "b0" , "c0" ] . into_iter( ) ] ;
1929+ let result = rotate_through_iterators ( a) . collect :: < Vec < _ > > ( ) ;
1930+
1931+ let expected = vec ! [ "a0" , "b0" , "c0" ] ;
1932+ assert_eq ! ( expected, result) ;
1933+
1934+ // test second vector with only one element
1935+ let a = vec ! [ vec![ "a0" , "b0" , "c0" ] . into_iter( ) , vec![ "a1" ] . into_iter( ) ] ;
1936+ let result = rotate_through_iterators ( a) . collect :: < Vec < _ > > ( ) ;
1937+
1938+ let expected = vec ! [ "a0" , "a1" , "b0" , "c0" ] ;
1939+ assert_eq ! ( expected, result) ;
1940+
1941+ // test three nestend vectors
1942+ let a = vec ! [ vec![ "a0" ] . into_iter( ) , vec![ "a1" , "b1" , "c1" ] . into_iter( ) , vec![ "a2" ] . into_iter( ) ] ;
1943+ let result = rotate_through_iterators ( a) . collect :: < Vec < _ > > ( ) ;
1944+
1945+ let expected = vec ! [ "a0" , "a1" , "a2" , "b1" , "c1" ] ;
1946+ assert_eq ! ( expected, result) ;
1947+
1948+ // test single nested vector with a single value
1949+ let a = vec ! [ vec![ "a0" ] . into_iter( ) ] ;
1950+ let result = rotate_through_iterators ( a) . collect :: < Vec < _ > > ( ) ;
1951+
1952+ let expected = vec ! [ "a0" ] ;
1953+ assert_eq ! ( expected, result) ;
1954+
1955+ // test single empty nested vector
1956+ let a: Vec < std:: vec:: IntoIter < & str > > = vec ! [ vec![ ] . into_iter( ) ] ;
1957+ let result = rotate_through_iterators ( a) . collect :: < Vec < & str > > ( ) ;
1958+ let expected: Vec < & str > = vec ! [ ] ;
1959+
1960+ assert_eq ! ( expected, result) ;
1961+
1962+ // test first nested vector is empty
1963+ let a: Vec < std:: vec:: IntoIter < & str > > = vec ! [ vec![ ] . into_iter( ) , vec![ "a1" , "b1" , "c1" ] . into_iter( ) ] ;
1964+ let result = rotate_through_iterators ( a) . collect :: < Vec < & str > > ( ) ;
1965+
1966+ let expected = vec ! [ "a1" , "b1" , "c1" ] ;
1967+ assert_eq ! ( expected, result) ;
1968+
1969+ // test two empty vectors
1970+ let a: Vec < std:: vec:: IntoIter < & str > > = vec ! [ vec![ ] . into_iter( ) , vec![ ] . into_iter( ) ] ;
1971+ let result = rotate_through_iterators ( a) . collect :: < Vec < & str > > ( ) ;
1972+
1973+ let expected: Vec < & str > = vec ! [ ] ;
1974+ assert_eq ! ( expected, result) ;
1975+
1976+ // test an empty vector amongst other filled vectors
1977+ let a = vec ! [
1978+ vec![ "a0" , "b0" , "c0" ] . into_iter( ) ,
1979+ vec![ ] . into_iter( ) ,
1980+ vec![ "a1" , "b1" , "c1" ] . into_iter( ) ,
1981+ vec![ "a2" , "b2" , "c2" ] . into_iter( ) ,
1982+ ] ;
1983+ let result = rotate_through_iterators ( a) . collect :: < Vec < _ > > ( ) ;
1984+
1985+ let expected = vec ! [ "a0" , "a1" , "a2" , "b0" , "b1" , "b2" , "c0" , "c1" , "c2" ] ;
1986+ assert_eq ! ( expected, result) ;
1987+
1988+ // test a filled vector between two empty vectors
1989+ let a = vec ! [ vec![ ] . into_iter( ) , vec![ "a1" , "b1" , "c1" ] . into_iter( ) , vec![ ] . into_iter( ) ] ;
1990+ let result = rotate_through_iterators ( a) . collect :: < Vec < _ > > ( ) ;
1991+
1992+ let expected = vec ! [ "a1" , "b1" , "c1" ] ;
1993+ assert_eq ! ( expected, result) ;
1994+
1995+ // test an empty vector at the end of the vectors
1996+ let a = vec ! [ vec![ "a0" , "b0" , "c0" ] . into_iter( ) , vec![ ] . into_iter( ) ] ;
1997+ let result = rotate_through_iterators ( a) . collect :: < Vec < _ > > ( ) ;
1998+
1999+ let expected = vec ! [ "a0" , "b0" , "c0" ] ;
2000+ assert_eq ! ( expected, result) ;
2001+
2002+ // test multiple empty vectors amongst multiple filled vectors
2003+ let a = vec ! [
2004+ vec![ ] . into_iter( ) ,
2005+ vec![ "a1" , "b1" , "c1" ] . into_iter( ) ,
2006+ vec![ ] . into_iter( ) ,
2007+ vec![ "a3" , "b3" ] . into_iter( ) ,
2008+ vec![ ] . into_iter( ) ,
2009+ ] ;
2010+
2011+ let result = rotate_through_iterators ( a) . collect :: < Vec < _ > > ( ) ;
2012+
2013+ let expected = vec ! [ "a1" , "a3" , "b1" , "b3" , "c1" ] ;
2014+ assert_eq ! ( expected, result) ;
2015+
2016+ // test one element in the first nested vectore and two elements in the second nested
2017+ // vector
2018+ let a = vec ! [ vec![ "a0" ] . into_iter( ) , vec![ "a1" , "b1" ] . into_iter( ) ] ;
2019+ let result = rotate_through_iterators ( a) . collect :: < Vec < _ > > ( ) ;
2020+
2021+ let expected = vec ! [ "a0" , "a1" , "b1" ] ;
2022+ assert_eq ! ( expected, result) ;
2023+ }
18892024}
0 commit comments