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, DebugNoBound, DefaultNoBound, EqNoBound, PalletError, PartialEqNoBound,
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)>);
108
109impl<T: Config<I>, I: 'static> List<T, I> {
110 pub(crate) fn unsafe_clear() {
117 #[allow(deprecated)]
118 crate::ListBags::<T, I>::remove_all(None);
119 #[allow(deprecated)]
120 crate::ListNodes::<T, I>::remove_all();
121 }
122
123 pub fn unsafe_regenerate(
133 all: impl IntoIterator<Item = T::AccountId>,
134 score_of: Box<dyn Fn(&T::AccountId) -> Option<T::Score>>,
135 ) -> u32 {
136 Self::unsafe_clear();
140 Self::insert_many(all, score_of)
141 }
142
143 #[allow(dead_code)]
163 pub fn migrate(old_thresholds: &[T::Score]) -> u32 {
164 let new_thresholds = T::BagThresholds::get();
165 if new_thresholds == old_thresholds {
166 return 0;
167 }
168
169 debug_assert!(
171 crate::ListBags::<T, I>::iter()
172 .all(|(threshold, _)| old_thresholds.contains(&threshold)),
173 "not all `bag_upper` currently in storage are members of `old_thresholds`",
174 );
175 debug_assert!(
176 crate::ListNodes::<T, I>::iter()
177 .all(|(_, node)| old_thresholds.contains(&node.bag_upper)),
178 "not all `node.bag_upper` currently in storage are members of `old_thresholds`",
179 );
180
181 let old_set: BTreeSet<_> = old_thresholds.iter().copied().collect();
182 let new_set: BTreeSet<_> = new_thresholds.iter().copied().collect();
183
184 let mut affected_accounts = BTreeSet::new();
186 let mut affected_old_bags = BTreeSet::new();
188
189 let new_bags = new_set.difference(&old_set).copied();
190 for inserted_bag in new_bags {
193 let affected_bag = {
194 let idx = old_thresholds.partition_point(|&threshold| inserted_bag > threshold);
196 old_thresholds.get(idx).copied().unwrap_or_else(T::Score::max_value)
197 };
198 if !affected_old_bags.insert(affected_bag) {
199 continue;
202 }
203
204 if let Some(bag) = Bag::<T, I>::get(affected_bag) {
205 affected_accounts.extend(bag.iter().map(|node| node.id));
206 }
207 }
208
209 let removed_bags = old_set.difference(&new_set).copied();
210 for removed_bag in removed_bags.clone() {
212 if !affected_old_bags.insert(removed_bag) {
213 continue;
214 }
215
216 if let Some(bag) = Bag::<T, I>::get(removed_bag) {
217 affected_accounts.extend(bag.iter().map(|node| node.id));
218 }
219 }
220
221 let num_affected = affected_accounts.len() as u32;
223 let score_of = T::ScoreProvider::score;
224 let _removed = Self::remove_many(&affected_accounts);
225 debug_assert_eq!(_removed, num_affected);
226 let _inserted = Self::insert_many(affected_accounts.into_iter(), score_of);
227 debug_assert_eq!(_inserted, num_affected);
228
229 for removed_bag in removed_bags {
236 debug_assert!(
237 !crate::ListNodes::<T, I>::iter().any(|(_, node)| node.bag_upper == removed_bag),
238 "no id should be present in a removed bag",
239 );
240 crate::ListBags::<T, I>::remove(removed_bag);
241 }
242
243 num_affected
244 }
245
246 pub(crate) fn contains(id: &T::AccountId) -> bool {
248 crate::ListNodes::<T, I>::contains_key(id)
249 }
250
251 pub fn get_score(id: &T::AccountId) -> Result<T::Score, ListError> {
253 Node::<T, I>::get(id).map(|node| node.score()).ok_or(ListError::NodeNotFound)
254 }
255
256 pub(crate) fn iter() -> impl Iterator<Item = Node<T, I>> {
261 let thresholds = T::BagThresholds::get();
268 let iter = thresholds.iter().copied();
269 let iter: Box<dyn Iterator<Item = T::Score>> = if thresholds.last() ==
270 Some(&T::Score::max_value())
271 {
272 Box::new(iter.rev())
274 } else {
275 Box::new(iter.chain(iter::once(T::Score::max_value())).rev())
277 };
278
279 iter.filter_map(Bag::get).flat_map(|bag| bag.iter())
280 }
281
282 pub(crate) fn iter_from(
286 start: &T::AccountId,
287 ) -> Result<impl Iterator<Item = Node<T, I>>, ListError> {
288 let start_node = Node::<T, I>::get(start).ok_or(ListError::NodeNotFound)?;
293 let start_node_upper = start_node.bag_upper;
294 let start_bag = core::iter::successors(start_node.next(), |prev| prev.next());
295
296 let thresholds = T::BagThresholds::get();
297 let idx = thresholds.partition_point(|&threshold| start_node_upper > threshold);
298 let leftover_bags = thresholds
299 .into_iter()
300 .take(idx)
301 .copied()
302 .rev()
303 .filter_map(Bag::get)
304 .flat_map(|bag| bag.iter());
305
306 crate::log!(
307 debug,
308 "starting to iterate from {:?}, who's bag is {:?}, and there are {:?} leftover bags",
309 &start,
310 start_node_upper,
311 idx
312 );
313 Ok(start_bag.chain(leftover_bags))
314 }
315
316 fn insert_many(
321 ids: impl IntoIterator<Item = T::AccountId>,
322 score_of: impl Fn(&T::AccountId) -> Option<T::Score>,
323 ) -> u32 {
324 let mut count = 0;
325 ids.into_iter().for_each(|v| {
326 if let Some(score) = score_of(&v) {
327 if Self::insert(v, score).is_ok() {
328 count += 1;
329 }
330 } else {
331 }
333 });
334
335 count
336 }
337
338 pub(crate) fn insert(id: T::AccountId, score: T::Score) -> Result<(), ListError> {
342 if Self::contains(&id) {
343 return Err(ListError::Duplicate);
344 }
345
346 let bag_score = notional_bag_for::<T, I>(score);
347 let mut bag = Bag::<T, I>::get_or_make(bag_score);
348 bag.insert_unchecked(id.clone(), score);
350
351 bag.put();
353
354 crate::log!(
355 trace,
356 "inserted {:?} with score {:?} into bag {:?}, new count is {}",
357 id,
358 score,
359 bag_score,
360 crate::ListNodes::<T, I>::count(),
361 );
362
363 Ok(())
364 }
365
366 pub(crate) fn remove(id: &T::AccountId) -> Result<(), ListError> {
368 if !Self::contains(id) {
369 return Err(ListError::NodeNotFound);
370 }
371 let _ = Self::remove_many(core::iter::once(id));
372 Ok(())
373 }
374
375 fn remove_many<'a>(ids: impl IntoIterator<Item = &'a T::AccountId>) -> u32 {
381 let mut bags = BTreeMap::new();
382 let mut count = 0;
383
384 for id in ids.into_iter() {
385 let node = match Node::<T, I>::get(id) {
386 Some(node) => node,
387 None => continue,
388 };
389 count += 1;
390
391 if !node.is_terminal() {
392 node.excise()
394 } else {
395 let bag = bags
397 .entry(node.bag_upper)
398 .or_insert_with(|| Bag::<T, I>::get_or_make(node.bag_upper));
399 bag.remove_node_unchecked(&node);
401 }
402
403 node.remove_from_storage_unchecked()
405 }
406
407 for (_, bag) in bags {
408 bag.put();
409 }
410
411 count
412 }
413
414 pub(crate) fn update_position_for(
426 mut node: Node<T, I>,
427 new_score: T::Score,
428 ) -> Option<(T::Score, T::Score)> {
429 node.score = new_score;
430 if node.is_misplaced(new_score) {
431 let old_bag_upper = node.bag_upper;
432
433 if !node.is_terminal() {
434 node.excise();
437 } else if let Some(mut bag) = Bag::<T, I>::get(node.bag_upper) {
438 bag.remove_node_unchecked(&node);
440 bag.put();
441 } else {
442 frame_support::defensive!(
443 "Node did not have a bag; BagsList is in an inconsistent state"
444 );
445 }
446
447 let new_bag_upper = notional_bag_for::<T, I>(new_score);
449 let mut bag = Bag::<T, I>::get_or_make(new_bag_upper);
450 bag.insert_node_unchecked(node);
453 bag.put();
454
455 Some((old_bag_upper, new_bag_upper))
456 } else {
457 node.put();
459 None
460 }
461 }
462
463 pub(crate) fn put_in_front_of(
466 lighter_id: &T::AccountId,
467 heavier_id: &T::AccountId,
468 ) -> Result<(), ListError> {
469 let lighter_node = Node::<T, I>::get(&lighter_id).ok_or(ListError::NodeNotFound)?;
470 let heavier_node = Node::<T, I>::get(&heavier_id).ok_or(ListError::NodeNotFound)?;
471
472 ensure!(lighter_node.bag_upper == heavier_node.bag_upper, ListError::NotInSameBag);
473
474 ensure!(
476 T::ScoreProvider::score(&heavier_id) > T::ScoreProvider::score(&lighter_id),
477 ListError::NotHeavier
478 );
479
480 let _ =
483 Self::remove(&heavier_id).defensive_proof("both nodes have been checked to exist; qed");
484
485 let lighter_node =
488 Node::<T, I>::get(lighter_id).defensive_ok_or_else(|| ListError::NodeNotFound)?;
489
490 Self::insert_at_unchecked(lighter_node, heavier_node);
493
494 Ok(())
495 }
496
497 fn insert_at_unchecked(mut at: Node<T, I>, mut node: Node<T, I>) {
504 node.prev = at.prev.clone();
506 if let Some(mut prev) = at.prev() {
507 prev.next = Some(node.id().clone());
508 prev.put()
509 }
510
511 node.next = Some(at.id().clone());
513 at.prev = Some(node.id().clone());
514
515 if node.is_terminal() {
516 let mut bag = Bag::<T, I>::get(at.bag_upper)
521 .expect("given nodes must always have a valid bag. qed.");
522
523 if node.prev == None {
524 bag.head = Some(node.id().clone())
525 }
526
527 bag.put()
528 };
529
530 at.put();
532 node.put();
533 }
534
535 #[cfg(any(test, feature = "try-runtime", feature = "fuzz"))]
545 pub(crate) fn do_try_state() -> Result<(), TryRuntimeError> {
546 let mut seen_in_list = BTreeSet::new();
547 ensure!(
548 Self::iter().map(|node| node.id).all(|id| seen_in_list.insert(id)),
549 "duplicate identified"
550 );
551
552 let iter_count = Self::iter().count() as u32;
553 let stored_count = crate::ListNodes::<T, I>::count();
554 let nodes_count = crate::ListNodes::<T, I>::iter().count() as u32;
555 ensure!(iter_count == stored_count, "iter_count != stored_count");
556 ensure!(stored_count == nodes_count, "stored_count != nodes_count");
557
558 crate::log!(trace, "count of nodes: {}", stored_count);
559
560 let active_bags = {
561 let thresholds = T::BagThresholds::get().iter().copied();
562 let thresholds: Vec<T::Score> =
563 if thresholds.clone().last() == Some(T::Score::max_value()) {
564 thresholds.collect()
566 } else {
567 thresholds.chain(iter::once(T::Score::max_value())).collect()
569 };
570 thresholds.into_iter().filter_map(|t| Bag::<T, I>::get(t))
571 };
572
573 let mut bags_map = BTreeMap::<T::Score, Vec<T::AccountId>>::new();
575
576 active_bags.clone().try_for_each(|b| {
577 bags_map.insert(
578 b.bag_upper,
579 b.iter().map(|n: Node<T, I>| n.id().clone()).collect::<Vec<_>>(),
580 );
581 b.do_try_state()
582 })?;
583
584 let nodes_in_bags_count =
585 active_bags.clone().fold(0u32, |acc, cur| acc + cur.iter().count() as u32);
586 ensure!(nodes_count == nodes_in_bags_count, "stored_count != nodes_in_bags_count");
587
588 crate::log!(trace, "count of active bags {}", active_bags.count());
589
590 for (_id, node) in crate::ListNodes::<T, I>::iter() {
593 let expected_bag = bags_map
595 .get(&node.bag_upper)
596 .ok_or("bag not found for the node in active bags")?;
597 frame_support::ensure!(expected_bag.contains(node.id()), "node not found in the bag");
598
599 node.do_try_state()?
601 }
602
603 Ok(())
604 }
605
606 #[cfg(any(feature = "std", feature = "runtime-benchmarks"))]
608 #[allow(dead_code)]
609 pub(crate) fn get_bags() -> Vec<(T::Score, Vec<T::AccountId>)> {
610 use frame_support::traits::Get as _;
611
612 let thresholds = T::BagThresholds::get();
613 let iter = thresholds.iter().copied();
614 let iter: Box<dyn Iterator<Item = T::Score>> = if thresholds.last() ==
615 Some(&T::Score::max_value())
616 {
617 Box::new(iter)
619 } else {
620 Box::new(iter.chain(core::iter::once(T::Score::max_value())))
622 };
623
624 iter.filter_map(|t| {
625 Bag::<T, I>::get(t)
626 .map(|bag| (t, bag.iter().map(|n| n.id().clone()).collect::<Vec<_>>()))
627 })
628 .collect::<Vec<_>>()
629 }
630}
631
632#[derive(
640 DefaultNoBound,
641 Encode,
642 Decode,
643 MaxEncodedLen,
644 TypeInfo,
645 DebugNoBound,
646 CloneNoBound,
647 PartialEqNoBound,
648 EqNoBound,
649)]
650#[codec(mel_bound())]
651#[scale_info(skip_type_params(T, I))]
652pub struct Bag<T: Config<I>, I: 'static = ()> {
653 pub head: Option<T::AccountId>,
654 pub tail: Option<T::AccountId>,
655
656 #[codec(skip)]
657 pub bag_upper: T::Score,
658 #[codec(skip)]
659 pub _phantom: PhantomData<I>,
660}
661
662impl<T: Config<I>, I: 'static> Bag<T, I> {
663 #[cfg(test)]
664 pub(crate) fn new(
665 head: Option<T::AccountId>,
666 tail: Option<T::AccountId>,
667 bag_upper: T::Score,
668 ) -> Self {
669 Self { head, tail, bag_upper, _phantom: PhantomData }
670 }
671
672 pub(crate) fn get(bag_upper: T::Score) -> Option<Bag<T, I>> {
674 crate::ListBags::<T, I>::try_get(bag_upper).ok().map(|mut bag| {
675 bag.bag_upper = bag_upper;
676 bag
677 })
678 }
679
680 fn get_or_make(bag_upper: T::Score) -> Bag<T, I> {
683 Self::get(bag_upper).unwrap_or(Bag { bag_upper, ..Default::default() })
684 }
685
686 fn is_empty(&self) -> bool {
688 self.head.is_none() && self.tail.is_none()
689 }
690
691 fn put(self) {
693 if self.is_empty() {
694 crate::ListBags::<T, I>::remove(self.bag_upper);
695 } else {
696 crate::ListBags::<T, I>::insert(self.bag_upper, self);
697 }
698 }
699
700 fn head(&self) -> Option<Node<T, I>> {
702 self.head.as_ref().and_then(|id| Node::get(id))
703 }
704
705 fn tail(&self) -> Option<Node<T, I>> {
707 self.tail.as_ref().and_then(|id| Node::get(id))
708 }
709
710 pub(crate) fn iter(&self) -> impl Iterator<Item = Node<T, I>> {
712 core::iter::successors(self.head(), |prev| prev.next())
713 }
714
715 fn insert_unchecked(&mut self, id: T::AccountId, score: T::Score) {
723 self.insert_node_unchecked(Node::<T, I> {
727 id,
728 prev: None,
729 next: None,
730 bag_upper: Zero::zero(),
731 score,
732 _phantom: PhantomData,
733 });
734 }
735
736 fn insert_node_unchecked(&mut self, mut node: Node<T, I>) {
744 if let Some(tail) = &self.tail {
745 if *tail == node.id {
746 defensive!("system logic error: inserting a node who has the id of tail");
749 return;
750 };
751 }
752
753 node.bag_upper = self.bag_upper;
756
757 let id = node.id.clone();
758 node.prev = self.tail.clone();
760 node.next = None;
761 node.put();
762
763 if let Some(mut old_tail) = self.tail() {
765 old_tail.next = Some(id.clone());
766 old_tail.put();
767 }
768 self.tail = Some(id.clone());
769
770 if self.head.is_none() {
774 self.head = Some(id);
775 debug_assert!(self.iter().count() == 1);
776 }
777 }
778
779 fn remove_node_unchecked(&mut self, node: &Node<T, I>) {
787 node.excise();
789
790 if self.tail.as_ref() == Some(&node.id) {
792 self.tail = node.prev.clone();
793 }
794 if self.head.as_ref() == Some(&node.id) {
795 self.head = node.next.clone();
796 }
797 }
798
799 #[cfg(any(test, feature = "try-runtime", feature = "fuzz"))]
808 fn do_try_state(&self) -> Result<(), TryRuntimeError> {
809 frame_support::ensure!(
810 self.head()
811 .map(|head| head.prev().is_none())
812 .unwrap_or_else(|| self.tail.is_none()),
815 "head has a prev"
816 );
817
818 frame_support::ensure!(
819 self.tail()
820 .map(|tail| tail.next().is_none())
821 .unwrap_or_else(|| self.head.is_none()),
824 "tail has a next"
825 );
826
827 let mut seen_in_bag = BTreeSet::new();
828 frame_support::ensure!(
829 self.iter()
830 .map(|node| node.id)
831 .all(|voter| seen_in_bag.insert(voter)),
833 "duplicate found in bag"
834 );
835
836 Ok(())
837 }
838
839 #[cfg(feature = "std")]
841 #[allow(dead_code)]
842 pub fn std_iter(&self) -> impl Iterator<Item = Node<T, I>> {
843 core::iter::successors(self.head(), |prev| prev.next())
844 }
845}
846
847#[derive(
849 Encode, Decode, MaxEncodedLen, TypeInfo, CloneNoBound, PartialEqNoBound, EqNoBound, DebugNoBound,
850)]
851#[codec(mel_bound())]
852#[scale_info(skip_type_params(T, I))]
853pub struct Node<T: Config<I>, I: 'static = ()> {
854 pub id: T::AccountId,
855 pub prev: Option<T::AccountId>,
856 pub next: Option<T::AccountId>,
857 pub bag_upper: T::Score,
858 pub score: T::Score,
859 #[codec(skip)]
860 pub _phantom: PhantomData<I>,
861}
862
863impl<T: Config<I>, I: 'static> Node<T, I> {
864 pub fn get(id: &T::AccountId) -> Option<Node<T, I>> {
866 crate::ListNodes::<T, I>::try_get(id).ok()
867 }
868
869 fn put(self) {
871 crate::ListNodes::<T, I>::insert(self.id.clone(), self);
872 }
873
874 fn excise(&self) {
879 if let Some(mut prev) = self.prev() {
881 prev.next = self.next.clone();
882 prev.put();
883 }
884 if let Some(mut next) = self.next() {
886 next.prev = self.prev.clone();
887 next.put();
888 }
889 }
890
891 fn remove_from_storage_unchecked(&self) {
895 crate::ListNodes::<T, I>::remove(&self.id)
896 }
897
898 fn prev(&self) -> Option<Node<T, I>> {
900 self.prev.as_ref().and_then(|id| Node::get(id))
901 }
902
903 fn next(&self) -> Option<Node<T, I>> {
905 self.next.as_ref().and_then(|id| Node::get(id))
906 }
907
908 pub fn is_misplaced(&self, current_score: T::Score) -> bool {
910 notional_bag_for::<T, I>(current_score) != self.bag_upper
911 }
912
913 fn is_terminal(&self) -> bool {
915 self.prev.is_none() || self.next.is_none()
916 }
917
918 pub(crate) fn id(&self) -> &T::AccountId {
920 &self.id
921 }
922
923 pub(crate) fn score(&self) -> T::Score {
925 self.score
926 }
927
928 #[cfg(feature = "std")]
930 #[allow(dead_code)]
931 pub fn std_id(&self) -> &T::AccountId {
932 &self.id
933 }
934
935 #[cfg(any(feature = "runtime-benchmarks", feature = "fuzz", test))]
936 pub fn set_score(&mut self, s: T::Score) {
937 self.score = s
938 }
939
940 #[cfg(feature = "runtime-benchmarks")]
942 #[allow(dead_code)]
943 pub fn bag_upper(&self) -> T::Score {
944 self.bag_upper
945 }
946
947 #[cfg(any(test, feature = "try-runtime", feature = "fuzz"))]
948 fn do_try_state(&self) -> Result<(), TryRuntimeError> {
949 let expected_bag = Bag::<T, I>::get(self.bag_upper).ok_or("bag not found for node")?;
950 let id = self.id();
951
952 let non_terminal_check = !self.is_terminal() &&
953 expected_bag.head.as_ref() != Some(id) &&
954 expected_bag.tail.as_ref() != Some(id);
955 let terminal_check =
956 expected_bag.head.as_ref() == Some(id) || expected_bag.tail.as_ref() == Some(id);
957 frame_support::ensure!(
958 non_terminal_check || terminal_check,
959 "a terminal node is neither its bag head or tail"
960 );
961
962 Ok(())
963 }
964}