From 9e8d7640bc8d565b022cd8bb988daae76a96845f Mon Sep 17 00:00:00 2001 From: kianenigma Date: Wed, 23 Mar 2022 21:53:38 +0000 Subject: [PATCH 1/4] add iter-from for bags-list --- frame/bags-list/src/lib.rs | 7 ++++++ frame/bags-list/src/list/mod.rs | 29 ++++++++++++++++++++++ frame/bags-list/src/tests.rs | 20 +++++++++++++++ frame/election-provider-support/src/lib.rs | 5 ++++ 4 files changed, 61 insertions(+) diff --git a/frame/bags-list/src/lib.rs b/frame/bags-list/src/lib.rs index c502245409fdb..f6c9cde1469ea 100644 --- a/frame/bags-list/src/lib.rs +++ b/frame/bags-list/src/lib.rs @@ -278,6 +278,13 @@ impl, I: 'static> SortedListProvider for Pallet Box::new(List::::iter().map(|n| n.id().clone())) } + fn iter_from( + start: &T::AccountId, + ) -> Result>, Self::Error> { + let iter = List::::iter_from(start)?; + Ok(Box::new(iter.map(|n| n.id().clone()))) + } + fn count() -> u32 { ListNodes::::count() } diff --git a/frame/bags-list/src/list/mod.rs b/frame/bags-list/src/list/mod.rs index 4921817c7e146..530e3502aa7c2 100644 --- a/frame/bags-list/src/list/mod.rs +++ b/frame/bags-list/src/list/mod.rs @@ -42,6 +42,8 @@ use sp_std::{ pub enum Error { /// A duplicate id has been detected. Duplicate, + /// Given node id was not found. + NodeNotFound, } #[cfg(test)] @@ -244,6 +246,33 @@ impl, I: 'static> List { iter.filter_map(Bag::get).flat_map(|bag| bag.iter()) } + /// Same as `iter`, but we start from a specific node. + /// + /// All items after this node are returned, excluding `start` itself. + pub(crate) fn iter_from( + start: &T::AccountId, + ) -> Result>, Error> { + // We chain two iterators: + // 1. from the given `start` till the end of the bag + // 2. all the bags that come after `start`'s bag. + + let start_node = Node::::get(start).ok_or(Error::NodeNotFound)?; + let start_node_upper = start_node.bag_upper; + let start_bag = sp_std::iter::successors(start_node.next(), |prev| prev.next()); + + let thresholds = T::BagThresholds::get(); + let idx = thresholds.partition_point(|&threshold| start_node_upper > threshold); + let leftover_bags = thresholds + .into_iter() + .take(idx) + .copied() + .rev() + .filter_map(Bag::get) + .flat_map(|bag| bag.iter()); + + Ok(start_bag.chain(leftover_bags)) + } + /// Insert several ids into the appropriate bags in the list. Continues with insertions /// if duplicates are detected. /// diff --git a/frame/bags-list/src/tests.rs b/frame/bags-list/src/tests.rs index 99396c9cbb3e3..f2ca7296e8a3b 100644 --- a/frame/bags-list/src/tests.rs +++ b/frame/bags-list/src/tests.rs @@ -458,6 +458,26 @@ mod sorted_list_provider { }); } + #[test] + fn iter_from_works() { + ExtBuilder::default().add_ids(vec![(5, 5), (6, 15)]).build_and_execute(|| { + // given + assert_eq!( + List::::get_bags(), + vec![(10, vec![1, 5]), (20, vec![6]), (1000, vec![2, 3, 4])] + ); + + assert_eq!(BagsList::iter_from(&2).unwrap().collect::>(), vec![3, 4, 6, 1, 5]); + assert_eq!(BagsList::iter_from(&3).unwrap().collect::>(), vec![4, 6, 1, 5]); + assert_eq!(BagsList::iter_from(&4).unwrap().collect::>(), vec![6, 1, 5]); + assert_eq!(BagsList::iter_from(&6).unwrap().collect::>(), vec![1, 5]); + assert_eq!(BagsList::iter_from(&1).unwrap().collect::>(), vec![5]); + assert!(BagsList::iter_from(&5).unwrap().collect::>().is_empty()); + + assert_storage_noop!(assert!(BagsList::iter_from(&8).is_err())); + }); + } + #[test] fn count_works() { ExtBuilder::default().build_and_execute(|| { diff --git a/frame/election-provider-support/src/lib.rs b/frame/election-provider-support/src/lib.rs index 2cc27472e8846..42fd003d8e4dd 100644 --- a/frame/election-provider-support/src/lib.rs +++ b/frame/election-provider-support/src/lib.rs @@ -432,6 +432,11 @@ pub trait SortedListProvider { /// An iterator over the list, which can have `take` called on it. fn iter() -> Box>; + /// Returns an iterator over the list, starting right after from the given voter. + /// + /// May return an error if `start` is invalid. + fn iter_from(start: &AccountId) -> Result>, Self::Error>; + /// The current count of ids in the list. fn count() -> u32; From 41837e24d38186af03bc4b0609c9f4f405c0f079 Mon Sep 17 00:00:00 2001 From: kianenigma Date: Wed, 23 Mar 2022 21:56:20 +0000 Subject: [PATCH 2/4] Fix --- frame/staking/src/pallet/impls.rs | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/frame/staking/src/pallet/impls.rs b/frame/staking/src/pallet/impls.rs index 9d5a3ed484184..90f19c6badd8f 100644 --- a/frame/staking/src/pallet/impls.rs +++ b/frame/staking/src/pallet/impls.rs @@ -1308,7 +1308,6 @@ impl SortedListProvider for UseNominatorsAndValidatorsM type Error = (); type Score = VoteWeight; - /// Returns iterator over voter list, which can have `take` called on it. fn iter() -> Box> { Box::new( Validators::::iter() @@ -1316,6 +1315,23 @@ impl SortedListProvider for UseNominatorsAndValidatorsM .chain(Nominators::::iter().map(|(n, _)| n)), ) } + fn iter_from( + start: &T::AccountId, + ) -> Result>, Self::Error> { + if Validators::::contains_key(start) { + let start_key = Validators::::hashed_key_for(start); + Ok(Box::new( + Validators::::iter_from(start_key) + .map(|(n, _)| n) + .chain(Nominators::::iter().map(|(x, _)| x)), + )) + } else if Nominators::::contains_key(start) { + let start_key = Nominators::::hashed_key_for(start); + Ok(Box::new(Nominators::::iter_from(start_key).map(|(n, _)| n))) + } else { + Err(()) + } + } fn count() -> u32 { Nominators::::count().saturating_add(Validators::::count()) } From e56a5e3ee43a4febdf82e704030a43ed0b8400bd Mon Sep 17 00:00:00 2001 From: Kian Paimani <5588131+kianenigma@users.noreply.github.com> Date: Mon, 28 Mar 2022 12:16:08 +0100 Subject: [PATCH 3/4] Apply suggestions from code review Co-authored-by: Oliver Tale-Yazdi --- frame/bags-list/src/tests.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/frame/bags-list/src/tests.rs b/frame/bags-list/src/tests.rs index f2ca7296e8a3b..7dc34e7a2f10e 100644 --- a/frame/bags-list/src/tests.rs +++ b/frame/bags-list/src/tests.rs @@ -473,6 +473,7 @@ mod sorted_list_provider { assert_eq!(BagsList::iter_from(&6).unwrap().collect::>(), vec![1, 5]); assert_eq!(BagsList::iter_from(&1).unwrap().collect::>(), vec![5]); assert!(BagsList::iter_from(&5).unwrap().collect::>().is_empty()); + assert!(BagsList::iter_from(&7).is_err()); assert_storage_noop!(assert!(BagsList::iter_from(&8).is_err())); }); From c3eb6af2f5e15d20e27f442834e56da38b57263a Mon Sep 17 00:00:00 2001 From: kianenigma Date: Thu, 31 Mar 2022 14:28:44 +0100 Subject: [PATCH 4/4] Fix --- frame/bags-list/src/lib.rs | 1 - frame/bags-list/src/list/mod.rs | 4 ++-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/frame/bags-list/src/lib.rs b/frame/bags-list/src/lib.rs index 45f3ca0615c50..94553433e230d 100644 --- a/frame/bags-list/src/lib.rs +++ b/frame/bags-list/src/lib.rs @@ -271,7 +271,6 @@ impl, I: 'static> Pallet { impl, I: 'static> SortedListProvider for Pallet { type Error = ListError; - type Score = T::Score; fn iter() -> Box> { diff --git a/frame/bags-list/src/list/mod.rs b/frame/bags-list/src/list/mod.rs index e30330d4539b0..db8c06a38d674 100644 --- a/frame/bags-list/src/list/mod.rs +++ b/frame/bags-list/src/list/mod.rs @@ -251,12 +251,12 @@ impl, I: 'static> List { /// All items after this node are returned, excluding `start` itself. pub(crate) fn iter_from( start: &T::AccountId, - ) -> Result>, Error> { + ) -> Result>, ListError> { // We chain two iterators: // 1. from the given `start` till the end of the bag // 2. all the bags that come after `start`'s bag. - let start_node = Node::::get(start).ok_or(Error::NodeNotFound)?; + let start_node = Node::::get(start).ok_or(ListError::NodeNotFound)?; let start_node_upper = start_node.bag_upper; let start_bag = sp_std::iter::successors(start_node.next(), |prev| prev.next());