Tidy Democracy (paritytech#10867)

* add test

* Assorted refactorings

* complete test

* saturating math

* final check

* use `default`

Co-authored-by: Gav Wood <gavin@parity.io>

frame/democracy/src/lib.rs

@@ -849,7 +849,12 @@ pub mod pallet {
 
 			<NextExternal<T>>::kill();
 			let now = <frame_system::Pallet<T>>::block_number();
-			Self::inject_referendum(now + voting_period, proposal_hash, threshold, delay);
+			Self::inject_referendum(
+				now.saturating_add(voting_period),
+				proposal_hash,
+				threshold,
+				delay,
+			);
 			Ok(())
 		}
 
@@ -878,7 +883,8 @@ pub mod pallet {
 				existing_vetoers.binary_search(&who).err().ok_or(Error::<T>::AlreadyVetoed)?;
 
 			existing_vetoers.insert(insert_position, who.clone());
-			let until = <frame_system::Pallet<T>>::block_number() + T::CooloffPeriod::get();
+			let until =
+				<frame_system::Pallet<T>>::block_number().saturating_add(T::CooloffPeriod::get());
 			<Blacklist<T>>::insert(&proposal_hash, (until, existing_vetoers));
 
 			Self::deposit_event(Event::<T>::Vetoed { who, proposal_hash, until });
@@ -1096,7 +1102,10 @@ pub mod pallet {
 			let now = <frame_system::Pallet<T>>::block_number();
 			let (voting, enactment) = (T::VotingPeriod::get(), T::EnactmentPeriod::get());
 			let additional = if who == provider { Zero::zero() } else { enactment };
-			ensure!(now >= since + voting + additional, Error::<T>::TooEarly);
+			ensure!(
+				now >= since.saturating_add(voting).saturating_add(additional),
+				Error::<T>::TooEarly
+			);
 			ensure!(expiry.map_or(true, |e| now > e), Error::<T>::Imminent);
 
 			let res =
@@ -1289,7 +1298,7 @@ impl<T: Config> Pallet<T> {
 	/// Get the amount locked in support of `proposal`; `None` if proposal isn't a valid proposal
 	/// index.
 	pub fn backing_for(proposal: PropIndex) -> Option<BalanceOf<T>> {
-		Self::deposit_of(proposal).map(|(l, d)| d * (l.len() as u32).into())
+		Self::deposit_of(proposal).map(|(l, d)| d.saturating_mul((l.len() as u32).into()))
 	}
 
 	/// Get all referenda ready for tally at block `n`.
@@ -1325,7 +1334,7 @@ impl<T: Config> Pallet<T> {
 		delay: T::BlockNumber,
 	) -> ReferendumIndex {
 		<Pallet<T>>::inject_referendum(
-			<frame_system::Pallet<T>>::block_number() + T::VotingPeriod::get(),
+			<frame_system::Pallet<T>>::block_number().saturating_add(T::VotingPeriod::get()),
 			proposal_hash,
 			threshold,
 			delay,
@@ -1431,7 +1440,9 @@ impl<T: Config> Pallet<T> {
 					},
 					Some(ReferendumInfo::Finished { end, approved }) => {
 						if let Some((lock_periods, balance)) = votes[i].1.locked_if(approved) {
-							let unlock_at = end + T::VoteLockingPeriod::get() * lock_periods.into();
+							let unlock_at = end.saturating_add(
+								T::VoteLockingPeriod::get().saturating_mul(lock_periods.into()),
+							);
 							let now = frame_system::Pallet::<T>::block_number();
 							if now < unlock_at {
 								ensure!(
@@ -1520,9 +1531,16 @@ impl<T: Config> Pallet<T> {
 			};
 			sp_std::mem::swap(&mut old, voting);
 			match old {
-				Voting::Delegating { balance, target, conviction, delegations, prior, .. } => {
+				Voting::Delegating {
+					balance, target, conviction, delegations, mut prior, ..
+				} => {
 					// remove any delegation votes to our current target.
 					Self::reduce_upstream_delegation(&target, conviction.votes(balance));
+					let now = frame_system::Pallet::<T>::block_number();
+					let lock_periods = conviction.lock_periods().into();
+					let unlock_block = now
+						.saturating_add(T::VoteLockingPeriod::get().saturating_mul(lock_periods));
+					prior.accumulate(unlock_block, balance);
 					voting.set_common(delegations, prior);
 				},
 				Voting::Direct { votes, delegations, prior } => {
@@ -1555,7 +1573,9 @@ impl<T: Config> Pallet<T> {
 						Self::reduce_upstream_delegation(&target, conviction.votes(balance));
 					let now = frame_system::Pallet::<T>::block_number();
 					let lock_periods = conviction.lock_periods().into();
-					prior.accumulate(now + T::VoteLockingPeriod::get() * lock_periods, balance);
+					let unlock_block = now
+						.saturating_add(T::VoteLockingPeriod::get().saturating_mul(lock_periods));
+					prior.accumulate(unlock_block, balance);
 					voting.set_common(delegations, prior);
 
 					Ok(votes)
@@ -1614,7 +1634,7 @@ impl<T: Config> Pallet<T> {
 			LastTabledWasExternal::<T>::put(true);
 			Self::deposit_event(Event::<T>::ExternalTabled);
 			Self::inject_referendum(
-				now + T::VotingPeriod::get(),
+				now.saturating_add(T::VotingPeriod::get()),
 				proposal,
 				threshold,
 				T::EnactmentPeriod::get(),
@@ -1646,7 +1666,7 @@ impl<T: Config> Pallet<T> {
 					depositors,
 				});
 				Self::inject_referendum(
-					now + T::VotingPeriod::get(),
+					now.saturating_add(T::VotingPeriod::get()),
 					proposal,
 					VoteThreshold::SuperMajorityApprove,
 					T::EnactmentPeriod::get(),
@@ -1700,7 +1720,7 @@ impl<T: Config> Pallet<T> {
 			if status.delay.is_zero() {
 				let _ = Self::do_enact_proposal(status.proposal_hash, index);
 			} else {
-				let when = now + status.delay;
+				let when = now.saturating_add(status.delay);
 				// Note that we need the preimage now.
 				Preimages::<T>::mutate_exists(
 					&status.proposal_hash,

frame/democracy/src/tests/delegation.rs

@@ -158,7 +158,7 @@ fn conviction_should_be_honored_in_delegation() {
 	// If transactor voted, delegated vote is overwritten.
 	new_test_ext().execute_with(|| {
 		let r = begin_referendum();
-		// Delegate, undelegate and vote.
+		// Delegate and vote.
 		assert_ok!(Democracy::delegate(Origin::signed(2), 1, Conviction::Locked6x, 20));
 		assert_ok!(Democracy::vote(Origin::signed(1), r, aye(1)));
 		// Delegated vote is huge.
@@ -177,3 +177,42 @@ fn split_vote_delegation_should_be_ignored() {
 		assert_eq!(tally(r), Tally { ayes: 1, nays: 0, turnout: 10 });
 	});
 }
+
+#[test]
+fn redelegation_keeps_lock() {
+	// If transactor voted, delegated vote is overwritten.
+	new_test_ext().execute_with(|| {
+		let r = begin_referendum();
+		// Delegate and vote.
+		assert_ok!(Democracy::delegate(Origin::signed(2), 1, Conviction::Locked6x, 20));
+		assert_ok!(Democracy::vote(Origin::signed(1), r, aye(1)));
+		// Delegated vote is huge.
+		assert_eq!(tally(r), Tally { ayes: 121, nays: 0, turnout: 30 });
+
+		let mut prior_lock = vote::PriorLock::default();
+
+		// Locked balance of delegator exists
+		assert_eq!(VotingOf::<Test>::get(2).locked_balance(), 20);
+		assert_eq!(VotingOf::<Test>::get(2).prior(), &prior_lock);
+
+		// Delegate someone else at a lower conviction and amount
+		assert_ok!(Democracy::delegate(Origin::signed(2), 3, Conviction::None, 10));
+
+		// 6x prior should appear w/ locked balance.
+		prior_lock.accumulate(98, 20);
+		assert_eq!(VotingOf::<Test>::get(2).prior(), &prior_lock);
+		assert_eq!(VotingOf::<Test>::get(2).locked_balance(), 20);
+		// Unlock shouldn't work
+		assert_ok!(Democracy::unlock(Origin::signed(2), 2));
+		assert_eq!(VotingOf::<Test>::get(2).prior(), &prior_lock);
+		assert_eq!(VotingOf::<Test>::get(2).locked_balance(), 20);
+
+		fast_forward_to(100);
+
+		// Now unlock can remove the prior lock and reduce the locked amount.
+		assert_eq!(VotingOf::<Test>::get(2).prior(), &prior_lock);
+		assert_ok!(Democracy::unlock(Origin::signed(2), 2));
+		assert_eq!(VotingOf::<Test>::get(2).prior(), &vote::PriorLock::default());
+		assert_eq!(VotingOf::<Test>::get(2).locked_balance(), 10);
+	});
+}

frame/democracy/src/vote.rs

@@ -183,7 +183,7 @@ impl<Balance: Saturating + Ord + Zero + Copy, BlockNumber: Ord + Copy + Zero, Ac
 		match self {
 			Voting::Direct { votes, prior, .. } =>
 				votes.iter().map(|i| i.1.balance()).fold(prior.locked(), |a, i| a.max(i)),
-			Voting::Delegating { balance, .. } => *balance,
+			Voting::Delegating { balance, prior, .. } => *balance.max(&prior.locked()),
 		}
 	}
 
@@ -199,4 +199,11 @@ impl<Balance: Saturating + Ord + Zero + Copy, BlockNumber: Ord + Copy + Zero, Ac
 		*d = delegations;
 		*p = prior;
 	}
+
+	pub fn prior(&self) -> &PriorLock<BlockNumber, Balance> {
+		match self {
+			Voting::Direct { prior, .. } => prior,
+			Voting::Delegating { prior, .. } => prior,
+		}
+	}
 }