Adding try_state hook for Vesting pallet

prdoc/pr_2847.prdoc

@@ -0,0 +1,9 @@
+title: Adding `try-state` hook to vesting pallet
+
+doc:
+  - audience: Runtime User
+    description: |
+      At the blocks before vesting begins, the locked amount of a vesting schedule must be equal to the product of the duration and the release per block amount when the locked amount is divisible by the per block amount, else the final vesting block should be equal to the unvested amount. This should also hold true during vesting schedules.
+
+crates:
+- name: pallet-vesting

substrate/frame/vesting/src/lib.rs

@@ -78,6 +78,9 @@ use sp_runtime::{
 };
 use sp_std::{fmt::Debug, marker::PhantomData, prelude::*};
 
+#[cfg(any(feature = "try-runtime", test))]
+use sp_runtime::{SaturatedConversion, TryRuntimeError};
+
 pub use pallet::*;
 pub use vesting_info::*;
 pub use weights::WeightInfo;
@@ -196,6 +199,11 @@ pub mod pallet {
 		fn integrity_test() {
 			assert!(T::MAX_VESTING_SCHEDULES > 0, "`MaxVestingSchedules` must ge greater than 0");
 		}
+
+		#[cfg(feature = "try-runtime")]
+		fn try_state(_n: BlockNumberFor<T>) -> Result<(), TryRuntimeError> {
+			Self::do_try_state()
+		}
 	}
 
 	/// Information regarding the vesting of a given account.
@@ -678,6 +686,149 @@ impl<T: Config> Pallet<T> {
 	}
 }
 
+/// Ensure the correctness of the state of this pallet.
+///
+/// The following expectations must always apply.
+///
+/// ## Expectations:
+///
+/// `handle_before_schedule_starts`:
+/// * the locked amount of a vesting schedule must be equal to the product of the duration
+///   (`schedules_left` - `starting_block`) and the per block amount when the locked amount is
+///   divisible by the per block amount.
+/// * However, If the locked amount is not divisible by the per block amount, the final vesting
+///   block (`schedules_left` - 1), the unvested amount should be equal to the remainder.
+///
+/// `handle_during_schedule`:
+/// * the amount `still_vesting` must be equal to the product of the remaining blocks to vest
+///   (`schedules_left` - `current_block`) and per block amount when the locked amount is divisible
+///   by the per block amount.
+/// * However, If the locked amount is not divisible by the per block amount, then at the final
+///   vesting block of the current schedule (`schedules_left` - 1), the unvested amount should be
+///   equal to the remainder.
+#[cfg(any(feature = "try-runtime", test))]
+impl<T: Config> Pallet<T> {
+	pub fn do_try_state() -> Result<(), TryRuntimeError> {
+		for who in Vesting::<T>::iter_keys() {
+			if let Some(infos) = Vesting::<T>::get(who.clone()) {
+				// Accumulate the total locked
+				let mut total_locked_amount: BalanceOf<T> = Zero::zero();
+
+				for info in infos.iter() {
+					let schedules_left: BalanceOf<T> =
+						info.ending_block_as_balance::<T::BlockNumberToBalance>();
+					let starting_block = T::BlockNumberToBalance::convert(info.starting_block());
+					let current_block_to_balance = T::BlockNumberToBalance::convert(
+						T::BlockNumberProvider::current_block_number(),
+					);
+
+					if current_block_to_balance < starting_block {
+						// handle the case when vesting has not started for this schedule
+						match Self::handle_before_schedule_starts(
+							info,
+							starting_block,
+							schedules_left,
+						) {
+							Ok(schedule_locked_amount) =>
+								total_locked_amount += schedule_locked_amount,
+							Err(e) => return Err(e),
+						}
+					} else {
+						match Self::handle_during_schedule(
+							info,
+							current_block_to_balance,
+							schedules_left,
+						) {
+							Ok(schedule_locked_amount) =>
+								total_locked_amount += schedule_locked_amount,
+							Err(e) => return Err(e),
+						}
+					}
+				}
+
+				if let Some(vesting_balance) = Self::vesting_balance(&who) {
+					ensure!(
+						vesting_balance == total_locked_amount,
+						TryRuntimeError::Other("inconsistent locked amount")
+					);
+				}
+			}
+		}
+		Ok(())
+	}
+
+	fn handle_before_schedule_starts(
+		info: &VestingInfo<BalanceOf<T>, BlockNumberFor<T>>,
+		starting_block: BalanceOf<T>,
+		schedules_left: BalanceOf<T>,
+	) -> Result<BalanceOf<T>, TryRuntimeError> {
+		let count = schedules_left.saturating_sub(starting_block);
+
+		let still_vesting = info
+			.locked_at::<T::BlockNumberToBalance>(T::BlockNumberProvider::current_block_number());
+
+		if (info.locked() % info.per_block()).is_zero() {
+			ensure!(
+                still_vesting == (count * info.per_block()),
+                TryRuntimeError::Other("Before schedule starts, the vesting balance should be equal to the total per block releases")
+            );
+			return Ok(still_vesting)
+		} else {
+			let re = info.locked() % info.per_block();
+
+			let final_vest_block =
+				schedules_left.saturating_sub(One::one()).saturated_into::<u64>();
+
+			let final_vest_amount =
+				info.locked_at::<T::BlockNumberToBalance>(final_vest_block.saturated_into());
+
+			ensure!(final_vest_amount == re, TryRuntimeError::Other("Before schedule starts, the final vest amount should be equal to the remainder"));
+
+			let no_schedules_left = schedules_left.saturated_into::<u64>();
+
+			let no_locks =
+				info.locked_at::<T::BlockNumberToBalance>(no_schedules_left.saturated_into());
+
+			ensure!(
+				no_locks == Zero::zero(),
+				TryRuntimeError::Other("After all schedules, all amounts should be unlocked")
+			);
+			return Ok(still_vesting)
+		}
+	}
+
+	fn handle_during_schedule(
+		info: &VestingInfo<BalanceOf<T>, BlockNumberFor<T>>,
+		current_block_to_balance: BalanceOf<T>,
+		schedules_left: BalanceOf<T>,
+	) -> Result<BalanceOf<T>, TryRuntimeError> {
+		let still_vesting = info
+			.locked_at::<T::BlockNumberToBalance>(T::BlockNumberProvider::current_block_number());
+
+		if (info.locked() % info.per_block()).is_zero() {
+			ensure!(
+                still_vesting == (schedules_left.saturating_sub(current_block_to_balance) * info.per_block()),
+                TryRuntimeError::Other("during schedules, the vesting balance should be equal to the total per block releases")
+            );
+			return Ok(still_vesting)
+		} else {
+			let re = info.locked() % info.per_block();
+
+			if current_block_to_balance == schedules_left.saturating_sub(One::one()) {
+				ensure!(still_vesting == re, TryRuntimeError::Other("At the final vesting block, the vesting balance should be equal to the remainder"));
+			}
+
+			if current_block_to_balance == schedules_left {
+				ensure!(
+					still_vesting == Zero::zero(),
+					TryRuntimeError::Other("Schedule ended, no more vesting balance")
+				);
+			}
+			return Ok(still_vesting)
+		}
+	}
+}
+
 impl<T: Config> VestingSchedule<T::AccountId> for Pallet<T>
 where
 	BalanceOf<T>: MaybeSerializeDeserialize + Debug,

substrate/frame/vesting/src/mock.rs

@@ -150,7 +150,16 @@ impl ExtBuilder {
 			.assimilate_storage(&mut t)
 			.unwrap();
 		let mut ext = sp_io::TestExternalities::new(t);
-		ext.execute_with(|| System::set_block_number(1));
+		ext.execute_with(|| {
+			System::set_block_number(1);
+		});
 		ext
 	}
+
+	pub fn build_and_execute(self, test: impl FnOnce() -> ()) {
+		self.build().execute_with(|| {
+			test();
+			Vesting::do_try_state().unwrap();
+		})
+	}
 }