add voting hooks to conviction-voting pallet

polkadot/runtime/rococo/src/governance/mod.rs

@@ -47,6 +47,7 @@ impl pallet_conviction_voting::Config for Runtime {
 	type MaxTurnout =
 		frame_support::traits::tokens::currency::ActiveIssuanceOf<Balances, Self::AccountId>;
 	type Polls = Referenda;
+	type VotingHooks = ();
 }
 
 parameter_types! {

polkadot/runtime/westend/src/governance/mod.rs

@@ -44,6 +44,7 @@ impl pallet_conviction_voting::Config for Runtime {
 	type MaxTurnout =
 		frame_support::traits::tokens::currency::ActiveIssuanceOf<Balances, Self::AccountId>;
 	type Polls = Referenda;
+	type VotingHooks = ();
 }
 
 parameter_types! {

substrate/bin/node/runtime/src/lib.rs

@@ -931,6 +931,7 @@ impl pallet_conviction_voting::Config for Runtime {
 	type MaxVotes = ConstU32<512>;
 	type MaxTurnout = frame_support::traits::TotalIssuanceOf<Balances, Self::AccountId>;
 	type Polls = Referenda;
+	type VotingHooks = ();
 }
 
 parameter_types! {

substrate/frame/conviction-voting/src/benchmarking.rs

@@ -73,6 +73,8 @@ benchmarks_instance_pallet! {
 		whitelist_account!(caller);
 		let account_vote = account_vote::<T, I>(100u32.into());
 
+		T::VotingHooks::on_vote_worst_case(&caller);
+
 		let (class, all_polls) = fill_voting::<T, I>();
 		let polls = &all_polls[&class];
 		let r = polls.len() - 1;
@@ -100,6 +102,8 @@ benchmarks_instance_pallet! {
 		whitelist_account!(caller);
 		let old_account_vote = account_vote::<T, I>(100u32.into());
 
+		T::VotingHooks::on_vote_worst_case(&caller);
+
 		let (class, all_polls) = fill_voting::<T, I>();
 		let polls = &all_polls[&class];
 		let r = polls.len();
@@ -128,6 +132,8 @@ benchmarks_instance_pallet! {
 		whitelist_account!(caller);
 		let old_account_vote = account_vote::<T, I>(100u32.into());
 
+		T::VotingHooks::on_remove_vote_worst_case(&caller);
+
 		let (class, all_polls) = fill_voting::<T, I>();
 		let polls = &all_polls[&class];
 		let r = polls.len();
@@ -157,6 +163,8 @@ benchmarks_instance_pallet! {
 		whitelist_account!(caller);
 		let old_account_vote = account_vote::<T, I>(100u32.into());
 
+		T::VotingHooks::on_remove_vote_worst_case(&caller);
+
 		let (class, all_polls) = fill_voting::<T, I>();
 		let polls = &all_polls[&class];
 		let r = polls.len();

substrate/frame/conviction-voting/src/lib.rs

@@ -44,13 +44,15 @@ use sp_runtime::{
 };
 
 mod conviction;
+pub mod traits;
 mod types;
 mod vote;
 pub mod weights;
 
 pub use self::{
 	conviction::Conviction,
 	pallet::*,
+	traits::VotingHooks,
 	types::{Delegations, Tally, UnvoteScope},
 	vote::{AccountVote, Casting, Delegating, Vote, Voting},
 	weights::WeightInfo,
@@ -137,6 +139,9 @@ pub mod pallet {
 		/// those successful voters are locked into the consequences that their votes entail.
 		#[pallet::constant]
 		type VoteLockingPeriod: Get<BlockNumberFor<Self>>;
+
+		/// Hooks are called when a new vote is registered or an existing vote is removed.
+		type VotingHooks: VotingHooks<Self::AccountId, PollIndexOf<Self, I>, BalanceOf<Self, I>>;
 	}
 
 	/// All voting for a particular voter in a particular voting class. We store the balance for the
@@ -433,6 +438,9 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 				// other votes are in place.
 				Self::extend_lock(who, &class, vote.balance());
 				Self::deposit_event(Event::Voted { who: who.clone(), vote });
+
+				// Call on_vote hook
+				T::VotingHooks::on_vote(who, poll_index, vote)?;
 				Ok(())
 			})
 		})
@@ -469,6 +477,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 							tally.reduce(approve, *delegations);
 						}
 						Self::deposit_event(Event::VoteRemoved { who: who.clone(), vote: v.1 });
+						T::VotingHooks::on_remove_vote(who, poll_index, Some(true));
 						Ok(())
 					},
 					PollStatus::Completed(end, approved) => {
@@ -484,10 +493,36 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 								);
 								prior.accumulate(unlock_at, balance)
 							}
+						} else if v.1.as_standard() == Some(!approved) {
+							// Unsuccessful vote, use special hook to lock the funds too in case of conviction.
+							if let Some(to_lock) =
+								T::VotingHooks::balance_locked_on_unsuccessful_vote(who, poll_index)
+							{
+								if let AccountVote::Standard { vote, .. } = v.1 {
+									let unlock_at = end.saturating_add(
+										T::VoteLockingPeriod::get()
+											.saturating_mul(vote.conviction.lock_periods().into()),
+									);
+									let now = frame_system::Pallet::<T>::block_number();
+									if now < unlock_at {
+										ensure!(
+											matches!(scope, UnvoteScope::Any),
+											Error::<T, I>::NoPermissionYet
+										);
+										prior.accumulate(unlock_at, to_lock)
+									}
+								}
+							}
 						}
+						// Call on_remove_vote hook
+						T::VotingHooks::on_remove_vote(who, poll_index, Some(false));
+						Ok(())
+					},
+					PollStatus::None => {
+						// Poll was cancelled.
+						T::VotingHooks::on_remove_vote(who, poll_index, None);
 						Ok(())
 					},
-					PollStatus::None => Ok(()), // Poll was cancelled.
 				})
 			} else {
 				Ok(())

substrate/frame/conviction-voting/src/tests.rs

@@ -17,6 +17,7 @@
 
 //! The crate's tests.
 
+use std::cell::RefCell;
 use std::collections::BTreeMap;
 
 use frame_support::{
@@ -154,6 +155,7 @@ impl Config for Test {
 	type WeightInfo = ();
 	type MaxTurnout = frame_support::traits::TotalIssuanceOf<Balances, Self::AccountId>;
 	type Polls = TestPolls;
+	type VotingHooks = HooksHandler;
 }
 
 pub fn new_test_ext() -> sp_io::TestExternalities {
@@ -846,3 +848,139 @@ fn errors_with_remove_vote_work() {
 		);
 	});
 }
+thread_local! {
+	static LAST_ON_VOTE_DATA: RefCell<Option<(u64, u8, AccountVote<u64>)>> = RefCell::new(None);
+	static LAST_ON_REMOVE_VOTE_DATA: RefCell<Option<(u64, u8, Option<bool>)>> = RefCell::new(None);
+	static LAST_LOCKED_IF_UNSUCCESSFUL_VOTE_DATA: RefCell<Option<(u64, u8)>> = RefCell::new(None);
+	static REMOVE_VOTE_LOCKED_AMOUNT: RefCell<Option<u64>> = RefCell::new(None);
+}
+
+pub struct HooksHandler;
+
+impl HooksHandler {
+	fn last_on_vote_data() -> Option<(u64, u8, AccountVote<u64>)> {
+		LAST_ON_VOTE_DATA.with(|data| data.borrow().clone())
+	}
+
+	fn last_on_remove_vote_data() -> Option<(u64, u8, Option<bool>)> {
+		LAST_ON_REMOVE_VOTE_DATA.with(|data| data.borrow().clone())
+	}
+
+	fn last_locked_if_unsuccessful_vote_data() -> Option<(u64, u8)> {
+		LAST_LOCKED_IF_UNSUCCESSFUL_VOTE_DATA.with(|data| data.borrow().clone())
+	}
+
+	fn reset() {
+		LAST_ON_VOTE_DATA.with(|data| *data.borrow_mut() = None);
+		LAST_ON_REMOVE_VOTE_DATA.with(|data| *data.borrow_mut() = None);
+		LAST_LOCKED_IF_UNSUCCESSFUL_VOTE_DATA.with(|data| *data.borrow_mut() = None);
+		REMOVE_VOTE_LOCKED_AMOUNT.with(|data| *data.borrow_mut() = None);
+	}
+
+	fn with_remove_locked_amount(v: u64) {
+		REMOVE_VOTE_LOCKED_AMOUNT.with(|data| *data.borrow_mut() = Some(v));
+	}
+}
+
+impl VotingHooks<u64, u8, u64> for HooksHandler {
+	fn on_vote(who: &u64, ref_index: u8, vote: AccountVote<u64>) -> DispatchResult {
+		LAST_ON_VOTE_DATA.with(|data| {
+			*data.borrow_mut() = Some((*who, ref_index, vote));
+		});
+		Ok(())
+	}
+
+	fn on_remove_vote(who: &u64, ref_index: u8, ongoing: Option<bool>) {
+		LAST_ON_REMOVE_VOTE_DATA.with(|data| {
+			*data.borrow_mut() = Some((*who, ref_index, ongoing));
+		});
+	}
+
+	fn balance_locked_on_unsuccessful_vote(who: &u64, ref_index: u8) -> Option<u64> {
+		LAST_LOCKED_IF_UNSUCCESSFUL_VOTE_DATA.with(|data| {
+			*data.borrow_mut() = Some((*who, ref_index));
+
+			REMOVE_VOTE_LOCKED_AMOUNT.with(|data| data.borrow().clone())
+		})
+	}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn on_vote_worst_case(_who: &u64) {}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn on_remove_vote_worst_case(_who: &u64) {}
+}
+
+#[test]
+fn voting_hooks_are_called_correctly() {
+	new_test_ext().execute_with(|| {
+		let c = class(3);
+
+		let usable_balance_1 = Balances::usable_balance(1);
+		dbg!(usable_balance_1);
+
+		// Voting
+		assert_ok!(Voting::vote(RuntimeOrigin::signed(1), 3, aye(1, 1)));
+		assert_eq!(
+			HooksHandler::last_on_vote_data(),
+			Some((
+				1,
+				3,
+				AccountVote::Standard {
+					vote: Vote { aye: true, conviction: Conviction::Locked1x },
+					balance: 1
+				}
+			))
+		);
+		assert_ok!(Voting::vote(RuntimeOrigin::signed(2), 3, nay(20, 2)));
+		assert_eq!(
+			HooksHandler::last_on_vote_data(),
+			Some((
+				2,
+				3,
+				AccountVote::Standard {
+					vote: Vote { aye: false, conviction: Conviction::Locked2x },
+					balance: 20
+				}
+			))
+		);
+		HooksHandler::reset();
+
+		// removing vote while ongoing
+		assert_ok!(Voting::vote(RuntimeOrigin::signed(3), 3, nay(20, 0)));
+		assert_eq!(
+			HooksHandler::last_on_vote_data(),
+			Some((
+				3,
+				3,
+				AccountVote::Standard {
+					vote: Vote { aye: false, conviction: Conviction::None },
+					balance: 20
+				}
+			))
+		);
+		assert_ok!(Voting::remove_vote(RuntimeOrigin::signed(3), Some(c), 3));
+		assert_eq!(HooksHandler::last_on_remove_vote_data(), Some((3, 3, Some(true))));
+		HooksHandler::reset();
+
+		Polls::set(vec![(3, Completed(3, false))].into_iter().collect());
+
+		// removing successful vote while completed
+		assert_ok!(Voting::remove_vote(RuntimeOrigin::signed(2), Some(c), 3));
+		assert_eq!(HooksHandler::last_on_remove_vote_data(), Some((2, 3, Some(false))));
+		assert_eq!(HooksHandler::last_locked_if_unsuccessful_vote_data(), None);
+
+		HooksHandler::reset();
+
+		HooksHandler::with_remove_locked_amount(5);
+
+		// removing unsuccessful vote when completed
+		assert_ok!(Voting::remove_vote(RuntimeOrigin::signed(1), Some(c), 3));
+		assert_eq!(HooksHandler::last_on_remove_vote_data(), Some((1, 3, Some(false))));
+		assert_eq!(HooksHandler::last_locked_if_unsuccessful_vote_data(), Some((1, 3)));
+
+		// Removing unsuccessful vote when completed should lock if given amount from the hook
+		assert_ok!(Voting::unlock(RuntimeOrigin::signed(1), c, 1));
+		assert_eq!(Balances::usable_balance(1), 5);
+	});
+}

substrate/frame/conviction-voting/src/traits.rs

@@ -0,0 +1,40 @@
+use crate::AccountVote;
+use frame_support::dispatch::DispatchResult;
+
+pub trait VotingHooks<AccountId, Index, Balance> {
+	// Called when vote is executed.
+	fn on_vote(who: &AccountId, ref_index: Index, vote: AccountVote<Balance>) -> DispatchResult;
+
+	// Called when removed vote is executed.
+	// is_finished indicates the state of the referendum = None if referendum is cancelled, Some(true) if referendum is ongoing and Some(false) when finished.
+	fn on_remove_vote(who: &AccountId, ref_index: Index, ongoing: Option<bool>);
+
+	// Called when removed vote is executed and voter lost the direction to possibly lock some balance.
+	// Can return an amount that should be locked for the conviction time.
+	fn balance_locked_on_unsuccessful_vote(who: &AccountId, ref_index: Index) -> Option<Balance>;
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn on_vote_worst_case(who: &AccountId);
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn on_remove_vote_worst_case(who: &AccountId);
+}
+
+// Default implementation for VotingHooks
+impl<A, I, B> VotingHooks<A, I, B> for () {
+	fn on_vote(_who: &A, _ref_index: I, _vote: AccountVote<B>) -> DispatchResult {
+		Ok(())
+	}
+
+	fn on_remove_vote(_who: &A, _ref_index: I, _ongoing: Option<bool>) {}
+
+	fn balance_locked_on_unsuccessful_vote(_who: &A, _ref_index: I) -> Option<B> {
+		None
+	}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn on_vote_worst_case(_who: &A) {}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn on_remove_vote_worst_case(_who: &A) {}
+}