1use crate::Config;
28use alloc::{
29 boxed::Box,
30 collections::{btree_map::BTreeMap, btree_set::BTreeSet},
31};
32use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen};
33use core::{iter, marker::PhantomData};
34use frame_election_provider_support::ScoreProvider;
35use frame_support::{
36 defensive, ensure,
37 traits::{Defensive, DefensiveOption, Get},
38 CloneNoBound, DefaultNoBound, EqNoBound, PalletError, PartialEqNoBound, RuntimeDebugNoBound,
39};
40use scale_info::TypeInfo;
41use sp_runtime::traits::{Bounded, Zero};
42
43#[cfg(any(
44 test,
45 feature = "try-runtime",
46 feature = "fuzz",
47 feature = "std",
48 feature = "runtime-benchmarks"
49))]
50use alloc::vec::Vec;
51#[cfg(any(test, feature = "try-runtime", feature = "fuzz"))]
52use sp_runtime::TryRuntimeError;
53
54#[derive(
55 Debug,
56 PartialEq,
57 Eq,
58 Encode,
59 Decode,
60 DecodeWithMemTracking,
61 MaxEncodedLen,
62 TypeInfo,
63 PalletError,
64)]
65pub enum ListError {
66 Duplicate,
68 NotHeavier,
70 NotInSameBag,
72 NodeNotFound,
74 Locked,
76}
77
78#[cfg(test)]
79mod tests;
80
81pub fn notional_bag_for<T: Config<I>, I: 'static>(score: T::Score) -> T::Score {
89 let thresholds = T::BagThresholds::get();
90 let idx = thresholds.partition_point(|&threshold| score > threshold);
91 thresholds.get(idx).copied().unwrap_or_else(T::Score::max_value)
92}
93
94pub struct List<T: Config<I>, I: 'static = ()>(PhantomData<(T, I)>);
109
110impl<T: Config<I>, I: 'static> List<T, I> {
111 pub(crate) fn unsafe_clear() {
118 #[allow(deprecated)]
119 crate::ListBags::<T, I>::remove_all(None);
120 #[allow(deprecated)]
121 crate::ListNodes::<T, I>::remove_all();
122 }
123
124 pub fn unsafe_regenerate(
134 all: impl IntoIterator<Item = T::AccountId>,
135 score_of: Box<dyn Fn(&T::AccountId) -> Option<T::Score>>,
136 ) -> u32 {
137 Self::unsafe_clear();
141 Self::insert_many(all, score_of)
142 }
143
144 #[allow(dead_code)]
164 pub fn migrate(old_thresholds: &[T::Score]) -> u32 {
165 let new_thresholds = T::BagThresholds::get();
166 if new_thresholds == old_thresholds {
167 return 0
168 }
169
170 debug_assert!(
172 crate::ListBags::<T, I>::iter()
173 .all(|(threshold, _)| old_thresholds.contains(&threshold)),
174 "not all `bag_upper` currently in storage are members of `old_thresholds`",
175 );
176 debug_assert!(
177 crate::ListNodes::<T, I>::iter()
178 .all(|(_, node)| old_thresholds.contains(&node.bag_upper)),
179 "not all `node.bag_upper` currently in storage are members of `old_thresholds`",
180 );
181
182 let old_set: BTreeSet<_> = old_thresholds.iter().copied().collect();
183 let new_set: BTreeSet<_> = new_thresholds.iter().copied().collect();
184
185 let mut affected_accounts = BTreeSet::new();
187 let mut affected_old_bags = BTreeSet::new();
189
190 let new_bags = new_set.difference(&old_set).copied();
191 for inserted_bag in new_bags {
194 let affected_bag = {
195 let idx = old_thresholds.partition_point(|&threshold| inserted_bag > threshold);
197 old_thresholds.get(idx).copied().unwrap_or_else(T::Score::max_value)
198 };
199 if !affected_old_bags.insert(affected_bag) {
200 continue
203 }
204
205 if let Some(bag) = Bag::<T, I>::get(affected_bag) {
206 affected_accounts.extend(bag.iter().map(|node| node.id));
207 }
208 }
209
210 let removed_bags = old_set.difference(&new_set).copied();
211 for removed_bag in removed_bags.clone() {
213 if !affected_old_bags.insert(removed_bag) {
214 continue
215 }
216
217 if let Some(bag) = Bag::<T, I>::get(removed_bag) {
218 affected_accounts.extend(bag.iter().map(|node| node.id));
219 }
220 }
221
222 let num_affected = affected_accounts.len() as u32;
224 let score_of = T::ScoreProvider::score;
225 let _removed = Self::remove_many(&affected_accounts);
226 debug_assert_eq!(_removed, num_affected);
227 let _inserted = Self::insert_many(affected_accounts.into_iter(), score_of);
228 debug_assert_eq!(_inserted, num_affected);
229
230 for removed_bag in removed_bags {
237 debug_assert!(
238 !crate::ListNodes::<T, I>::iter().any(|(_, node)| node.bag_upper == removed_bag),
239 "no id should be present in a removed bag",
240 );
241 crate::ListBags::<T, I>::remove(removed_bag);
242 }
243
244 num_affected
245 }
246
247 pub(crate) fn contains(id: &T::AccountId) -> bool {
249 crate::ListNodes::<T, I>::contains_key(id)
250 }
251
252 pub fn get_score(id: &T::AccountId) -> Result<T::Score, ListError> {
254 Node::<T, I>::get(id).map(|node| node.score()).ok_or(ListError::NodeNotFound)
255 }
256
257 pub(crate) fn iter() -> impl Iterator<Item = Node<T, I>> {
262 let thresholds = T::BagThresholds::get();
269 let iter = thresholds.iter().copied();
270 let iter: Box<dyn Iterator<Item = T::Score>> = if thresholds.last() ==
271 Some(&T::Score::max_value())
272 {
273 Box::new(iter.rev())
275 } else {
276 Box::new(iter.chain(iter::once(T::Score::max_value())).rev())
278 };
279
280 iter.filter_map(Bag::get).flat_map(|bag| bag.iter())
281 }
282
283 pub(crate) fn iter_from(
287 start: &T::AccountId,
288 ) -> Result<impl Iterator<Item = Node<T, I>>, ListError> {
289 let start_node = Node::<T, I>::get(start).ok_or(ListError::NodeNotFound)?;
294 let start_node_upper = start_node.bag_upper;
295 let start_bag = core::iter::successors(start_node.next(), |prev| prev.next());
296
297 let thresholds = T::BagThresholds::get();
298 let idx = thresholds.partition_point(|&threshold| start_node_upper > threshold);
299 let leftover_bags = thresholds
300 .into_iter()
301 .take(idx)
302 .copied()
303 .rev()
304 .filter_map(Bag::get)
305 .flat_map(|bag| bag.iter());
306
307 crate::log!(
308 debug,
309 "starting to iterate from {:?}, who's bag is {:?}, and there are {:?} leftover bags",
310 &start,
311 start_node_upper,
312 idx
313 );
314 Ok(start_bag.chain(leftover_bags))
315 }
316
317 fn insert_many(
322 ids: impl IntoIterator<Item = T::AccountId>,
323 score_of: impl Fn(&T::AccountId) -> Option<T::Score>,
324 ) -> u32 {
325 let mut count = 0;
326 ids.into_iter().for_each(|v| {
327 if let Some(score) = score_of(&v) {
328 if Self::insert(v, score).is_ok() {
329 count += 1;
330 }
331 } else {
332 }
334 });
335
336 count
337 }
338
339 pub(crate) fn insert(id: T::AccountId, score: T::Score) -> Result<(), ListError> {
343 if Self::contains(&id) {
344 return Err(ListError::Duplicate)
345 }
346
347 let bag_score = notional_bag_for::<T, I>(score);
348 let mut bag = Bag::<T, I>::get_or_make(bag_score);
349 bag.insert_unchecked(id.clone(), score);
351
352 bag.put();
354
355 crate::log!(
356 trace,
357 "inserted {:?} with score {:?} into bag {:?}, new count is {}",
358 id,
359 score,
360 bag_score,
361 crate::ListNodes::<T, I>::count(),
362 );
363
364 Ok(())
365 }
366
367 pub(crate) fn remove(id: &T::AccountId) -> Result<(), ListError> {
369 if !Self::contains(id) {
370 return Err(ListError::NodeNotFound)
371 }
372 let _ = Self::remove_many(core::iter::once(id));
373 Ok(())
374 }
375
376 fn remove_many<'a>(ids: impl IntoIterator<Item = &'a T::AccountId>) -> u32 {
382 let mut bags = BTreeMap::new();
383 let mut count = 0;
384
385 for id in ids.into_iter() {
386 let node = match Node::<T, I>::get(id) {
387 Some(node) => node,
388 None => continue,
389 };
390 count += 1;
391
392 if !node.is_terminal() {
393 node.excise()
395 } else {
396 let bag = bags
398 .entry(node.bag_upper)
399 .or_insert_with(|| Bag::<T, I>::get_or_make(node.bag_upper));
400 bag.remove_node_unchecked(&node);
402 }
403
404 node.remove_from_storage_unchecked()
406 }
407
408 for (_, bag) in bags {
409 bag.put();
410 }
411
412 count
413 }
414
415 pub(crate) fn update_position_for(
427 mut node: Node<T, I>,
428 new_score: T::Score,
429 ) -> Option<(T::Score, T::Score)> {
430 node.score = new_score;
431 if node.is_misplaced(new_score) {
432 let old_bag_upper = node.bag_upper;
433
434 if !node.is_terminal() {
435 node.excise();
438 } else if let Some(mut bag) = Bag::<T, I>::get(node.bag_upper) {
439 bag.remove_node_unchecked(&node);
441 bag.put();
442 } else {
443 frame_support::defensive!(
444 "Node did not have a bag; BagsList is in an inconsistent state"
445 );
446 }
447
448 let new_bag_upper = notional_bag_for::<T, I>(new_score);
450 let mut bag = Bag::<T, I>::get_or_make(new_bag_upper);
451 bag.insert_node_unchecked(node);
454 bag.put();
455
456 Some((old_bag_upper, new_bag_upper))
457 } else {
458 node.put();
460 None
461 }
462 }
463
464 pub(crate) fn put_in_front_of(
467 lighter_id: &T::AccountId,
468 heavier_id: &T::AccountId,
469 ) -> Result<(), ListError> {
470 let lighter_node = Node::<T, I>::get(&lighter_id).ok_or(ListError::NodeNotFound)?;
471 let heavier_node = Node::<T, I>::get(&heavier_id).ok_or(ListError::NodeNotFound)?;
472
473 ensure!(lighter_node.bag_upper == heavier_node.bag_upper, ListError::NotInSameBag);
474
475 ensure!(
477 T::ScoreProvider::score(&heavier_id) > T::ScoreProvider::score(&lighter_id),
478 ListError::NotHeavier
479 );
480
481 let _ =
484 Self::remove(&heavier_id).defensive_proof("both nodes have been checked to exist; qed");
485
486 let lighter_node =
489 Node::<T, I>::get(lighter_id).defensive_ok_or_else(|| ListError::NodeNotFound)?;
490
491 Self::insert_at_unchecked(lighter_node, heavier_node);
494
495 Ok(())
496 }
497
498 fn insert_at_unchecked(mut at: Node<T, I>, mut node: Node<T, I>) {
505 node.prev = at.prev.clone();
507 if let Some(mut prev) = at.prev() {
508 prev.next = Some(node.id().clone());
509 prev.put()
510 }
511
512 node.next = Some(at.id().clone());
514 at.prev = Some(node.id().clone());
515
516 if node.is_terminal() {
517 let mut bag = Bag::<T, I>::get(at.bag_upper)
522 .expect("given nodes must always have a valid bag. qed.");
523
524 if node.prev == None {
525 bag.head = Some(node.id().clone())
526 }
527
528 bag.put()
529 };
530
531 at.put();
533 node.put();
534 }
535
536 #[cfg(any(test, feature = "try-runtime", feature = "fuzz"))]
546 pub(crate) fn do_try_state() -> Result<(), TryRuntimeError> {
547 let mut seen_in_list = BTreeSet::new();
548 ensure!(
549 Self::iter().map(|node| node.id).all(|id| seen_in_list.insert(id)),
550 "duplicate identified"
551 );
552
553 let iter_count = Self::iter().count() as u32;
554 let stored_count = crate::ListNodes::<T, I>::count();
555 let nodes_count = crate::ListNodes::<T, I>::iter().count() as u32;
556 ensure!(iter_count == stored_count, "iter_count != stored_count");
557 ensure!(stored_count == nodes_count, "stored_count != nodes_count");
558
559 crate::log!(trace, "count of nodes: {}", stored_count);
560
561 let active_bags = {
562 let thresholds = T::BagThresholds::get().iter().copied();
563 let thresholds: Vec<T::Score> =
564 if thresholds.clone().last() == Some(T::Score::max_value()) {
565 thresholds.collect()
567 } else {
568 thresholds.chain(iter::once(T::Score::max_value())).collect()
570 };
571 thresholds.into_iter().filter_map(|t| Bag::<T, I>::get(t))
572 };
573
574 let mut bags_map = BTreeMap::<T::Score, Vec<T::AccountId>>::new();
576
577 active_bags.clone().try_for_each(|b| {
578 bags_map.insert(
579 b.bag_upper,
580 b.iter().map(|n: Node<T, I>| n.id().clone()).collect::<Vec<_>>(),
581 );
582 b.do_try_state()
583 })?;
584
585 let nodes_in_bags_count =
586 active_bags.clone().fold(0u32, |acc, cur| acc + cur.iter().count() as u32);
587 ensure!(nodes_count == nodes_in_bags_count, "stored_count != nodes_in_bags_count");
588
589 crate::log!(trace, "count of active bags {}", active_bags.count());
590
591 for (_id, node) in crate::ListNodes::<T, I>::iter() {
594 let expected_bag = bags_map
596 .get(&node.bag_upper)
597 .ok_or("bag not found for the node in active bags")?;
598 frame_support::ensure!(expected_bag.contains(node.id()), "node not found in the bag");
599
600 node.do_try_state()?
602 }
603
604 Ok(())
605 }
606
607 #[cfg(any(feature = "std", feature = "runtime-benchmarks"))]
609 #[allow(dead_code)]
610 pub(crate) fn get_bags() -> Vec<(T::Score, Vec<T::AccountId>)> {
611 use frame_support::traits::Get as _;
612
613 let thresholds = T::BagThresholds::get();
614 let iter = thresholds.iter().copied();
615 let iter: Box<dyn Iterator<Item = T::Score>> = if thresholds.last() ==
616 Some(&T::Score::max_value())
617 {
618 Box::new(iter)
620 } else {
621 Box::new(iter.chain(core::iter::once(T::Score::max_value())))
623 };
624
625 iter.filter_map(|t| {
626 Bag::<T, I>::get(t)
627 .map(|bag| (t, bag.iter().map(|n| n.id().clone()).collect::<Vec<_>>()))
628 })
629 .collect::<Vec<_>>()
630 }
631}
632
633#[derive(
641 DefaultNoBound,
642 Encode,
643 Decode,
644 MaxEncodedLen,
645 TypeInfo,
646 RuntimeDebugNoBound,
647 CloneNoBound,
648 PartialEqNoBound,
649 EqNoBound,
650)]
651#[codec(mel_bound())]
652#[scale_info(skip_type_params(T, I))]
653pub struct Bag<T: Config<I>, I: 'static = ()> {
654 pub head: Option<T::AccountId>,
655 pub tail: Option<T::AccountId>,
656
657 #[codec(skip)]
658 pub bag_upper: T::Score,
659 #[codec(skip)]
660 pub _phantom: PhantomData<I>,
661}
662
663impl<T: Config<I>, I: 'static> Bag<T, I> {
664 #[cfg(test)]
665 pub(crate) fn new(
666 head: Option<T::AccountId>,
667 tail: Option<T::AccountId>,
668 bag_upper: T::Score,
669 ) -> Self {
670 Self { head, tail, bag_upper, _phantom: PhantomData }
671 }
672
673 pub(crate) fn get(bag_upper: T::Score) -> Option<Bag<T, I>> {
675 crate::ListBags::<T, I>::try_get(bag_upper).ok().map(|mut bag| {
676 bag.bag_upper = bag_upper;
677 bag
678 })
679 }
680
681 fn get_or_make(bag_upper: T::Score) -> Bag<T, I> {
684 Self::get(bag_upper).unwrap_or(Bag { bag_upper, ..Default::default() })
685 }
686
687 fn is_empty(&self) -> bool {
689 self.head.is_none() && self.tail.is_none()
690 }
691
692 fn put(self) {
694 if self.is_empty() {
695 crate::ListBags::<T, I>::remove(self.bag_upper);
696 } else {
697 crate::ListBags::<T, I>::insert(self.bag_upper, self);
698 }
699 }
700
701 fn head(&self) -> Option<Node<T, I>> {
703 self.head.as_ref().and_then(|id| Node::get(id))
704 }
705
706 fn tail(&self) -> Option<Node<T, I>> {
708 self.tail.as_ref().and_then(|id| Node::get(id))
709 }
710
711 pub(crate) fn iter(&self) -> impl Iterator<Item = Node<T, I>> {
713 core::iter::successors(self.head(), |prev| prev.next())
714 }
715
716 fn insert_unchecked(&mut self, id: T::AccountId, score: T::Score) {
724 self.insert_node_unchecked(Node::<T, I> {
728 id,
729 prev: None,
730 next: None,
731 bag_upper: Zero::zero(),
732 score,
733 _phantom: PhantomData,
734 });
735 }
736
737 fn insert_node_unchecked(&mut self, mut node: Node<T, I>) {
745 if let Some(tail) = &self.tail {
746 if *tail == node.id {
747 defensive!("system logic error: inserting a node who has the id of tail");
750 return
751 };
752 }
753
754 node.bag_upper = self.bag_upper;
757
758 let id = node.id.clone();
759 node.prev = self.tail.clone();
761 node.next = None;
762 node.put();
763
764 if let Some(mut old_tail) = self.tail() {
766 old_tail.next = Some(id.clone());
767 old_tail.put();
768 }
769 self.tail = Some(id.clone());
770
771 if self.head.is_none() {
775 self.head = Some(id);
776 debug_assert!(self.iter().count() == 1);
777 }
778 }
779
780 fn remove_node_unchecked(&mut self, node: &Node<T, I>) {
788 node.excise();
790
791 if self.tail.as_ref() == Some(&node.id) {
793 self.tail = node.prev.clone();
794 }
795 if self.head.as_ref() == Some(&node.id) {
796 self.head = node.next.clone();
797 }
798 }
799
800 #[cfg(any(test, feature = "try-runtime", feature = "fuzz"))]
809 fn do_try_state(&self) -> Result<(), TryRuntimeError> {
810 frame_support::ensure!(
811 self.head()
812 .map(|head| head.prev().is_none())
813 .unwrap_or_else(|| self.tail.is_none()),
816 "head has a prev"
817 );
818
819 frame_support::ensure!(
820 self.tail()
821 .map(|tail| tail.next().is_none())
822 .unwrap_or_else(|| self.head.is_none()),
825 "tail has a next"
826 );
827
828 let mut seen_in_bag = BTreeSet::new();
829 frame_support::ensure!(
830 self.iter()
831 .map(|node| node.id)
832 .all(|voter| seen_in_bag.insert(voter)),
834 "duplicate found in bag"
835 );
836
837 Ok(())
838 }
839
840 #[cfg(feature = "std")]
842 #[allow(dead_code)]
843 pub fn std_iter(&self) -> impl Iterator<Item = Node<T, I>> {
844 core::iter::successors(self.head(), |prev| prev.next())
845 }
846}
847
848#[derive(
850 Encode,
851 Decode,
852 MaxEncodedLen,
853 TypeInfo,
854 CloneNoBound,
855 PartialEqNoBound,
856 EqNoBound,
857 RuntimeDebugNoBound,
858)]
859#[codec(mel_bound())]
860#[scale_info(skip_type_params(T, I))]
861pub struct Node<T: Config<I>, I: 'static = ()> {
862 pub id: T::AccountId,
863 pub prev: Option<T::AccountId>,
864 pub next: Option<T::AccountId>,
865 pub bag_upper: T::Score,
866 pub score: T::Score,
867 #[codec(skip)]
868 pub _phantom: PhantomData<I>,
869}
870
871impl<T: Config<I>, I: 'static> Node<T, I> {
872 pub fn get(id: &T::AccountId) -> Option<Node<T, I>> {
874 crate::ListNodes::<T, I>::try_get(id).ok()
875 }
876
877 fn put(self) {
879 crate::ListNodes::<T, I>::insert(self.id.clone(), self);
880 }
881
882 fn excise(&self) {
887 if let Some(mut prev) = self.prev() {
889 prev.next = self.next.clone();
890 prev.put();
891 }
892 if let Some(mut next) = self.next() {
894 next.prev = self.prev.clone();
895 next.put();
896 }
897 }
898
899 fn remove_from_storage_unchecked(&self) {
903 crate::ListNodes::<T, I>::remove(&self.id)
904 }
905
906 fn prev(&self) -> Option<Node<T, I>> {
908 self.prev.as_ref().and_then(|id| Node::get(id))
909 }
910
911 fn next(&self) -> Option<Node<T, I>> {
913 self.next.as_ref().and_then(|id| Node::get(id))
914 }
915
916 pub fn is_misplaced(&self, current_score: T::Score) -> bool {
918 notional_bag_for::<T, I>(current_score) != self.bag_upper
919 }
920
921 fn is_terminal(&self) -> bool {
923 self.prev.is_none() || self.next.is_none()
924 }
925
926 pub(crate) fn id(&self) -> &T::AccountId {
928 &self.id
929 }
930
931 pub(crate) fn score(&self) -> T::Score {
933 self.score
934 }
935
936 #[cfg(feature = "std")]
938 #[allow(dead_code)]
939 pub fn std_id(&self) -> &T::AccountId {
940 &self.id
941 }
942
943 #[cfg(any(feature = "runtime-benchmarks", feature = "fuzz", test))]
944 pub fn set_score(&mut self, s: T::Score) {
945 self.score = s
946 }
947
948 #[cfg(feature = "runtime-benchmarks")]
950 #[allow(dead_code)]
951 pub fn bag_upper(&self) -> T::Score {
952 self.bag_upper
953 }
954
955 #[cfg(any(test, feature = "try-runtime", feature = "fuzz"))]
956 fn do_try_state(&self) -> Result<(), TryRuntimeError> {
957 let expected_bag = Bag::<T, I>::get(self.bag_upper).ok_or("bag not found for node")?;
958 let id = self.id();
959
960 let non_terminal_check = !self.is_terminal() &&
961 expected_bag.head.as_ref() != Some(id) &&
962 expected_bag.tail.as_ref() != Some(id);
963 let terminal_check =
964 expected_bag.head.as_ref() == Some(id) || expected_bag.tail.as_ref() == Some(id);
965 frame_support::ensure!(
966 non_terminal_check || terminal_check,
967 "a terminal node is neither its bag head or tail"
968 );
969
970 Ok(())
971 }
972}