Adding try_state hook(tips & vesting)

frame/tips/src/lib.rs

@@ -68,6 +68,7 @@ use sp_std::prelude::*;
 
 use codec::{Decode, Encode};
 use frame_support::{
+	ensure,
 	traits::{
 		ContainsLengthBound, Currency, EnsureOrigin, ExistenceRequirement::KeepAlive, Get,
 		OnUnbalanced, ReservableCurrency, SortedMembers,
@@ -76,6 +77,9 @@ use frame_support::{
 };
 use frame_system::pallet_prelude::BlockNumberFor;
 
+#[cfg(any(feature = "try-runtime", test))]
+use sp_runtime::TryRuntimeError;
+
 pub use pallet::*;
 pub use weights::WeightInfo;
 
@@ -453,6 +457,14 @@ pub mod pallet {
 			Ok(())
 		}
 	}
+
+	#[pallet::hooks]
+	impl<T: Config<I>, I: 'static> Hooks<BlockNumberFor<T>> for Pallet<T, I> {
+		#[cfg(feature = "try-runtime")]
+		fn try_state(_n: BlockNumberFor<T>) -> Result<(), TryRuntimeError> {
+			Self::do_try_state()
+		}
+	}
 }
 
 impl<T: Config<I>, I: 'static> Pallet<T, I> {
@@ -599,4 +611,22 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 			Tips::<T, I>::insert(hash, new_tip)
 		}
 	}
+
+	/// Ensure the correctness of the state of this pallet.
+	///
+	/// This should be valid before or after each state transition of this pallet.
+	///
+	/// ## Invariants:
+	///
+	/// * There should be a corresponding `OpenTip.reason` for each key in `Reasons`.
+	#[cfg(any(feature = "try-runtime", test))]
+	pub fn do_try_state() -> Result<(), TryRuntimeError> {
+		let reasons = Reasons::<T, I>::iter_keys().collect::<Vec<_>>();
+
+		for tip in Tips::<T, I>::iter_keys() {
+			let tip_reason = Tips::<T, I>::get(&tip).unwrap().reason;
+			ensure!(reasons.contains(&tip_reason), TryRuntimeError::Other("reasons don't match"));
+		}
+		Ok(())
+	}
 }

frame/tips/src/tests.rs

@@ -231,6 +231,7 @@ fn tip_new_cannot_be_used_twice() {
 			Tips::tip_new(RuntimeOrigin::signed(11), b"awesome.dot".to_vec(), 3, 10),
 			Error::<Test>::AlreadyKnown
 		);
+		Tips::do_try_state().unwrap();
 	});
 }
 
@@ -254,6 +255,7 @@ fn report_awesome_and_tip_works() {
 		assert_ok!(Tips::tip(RuntimeOrigin::signed(12), h, 10));
 		assert_noop!(Tips::tip(RuntimeOrigin::signed(9), h, 10), BadOrigin);
 		System::set_block_number(2);
+		Tips::do_try_state().unwrap();
 		assert_ok!(Tips::close_tip(RuntimeOrigin::signed(100), h.into()));
 		assert_eq!(Balances::reserved_balance(0), 0);
 		assert_eq!(Balances::free_balance(0), 102);
@@ -273,6 +275,7 @@ fn report_awesome_from_beneficiary_and_tip_works() {
 		assert_ok!(Tips::tip(RuntimeOrigin::signed(11), h, 10));
 		assert_ok!(Tips::tip(RuntimeOrigin::signed(12), h, 10));
 		System::set_block_number(2);
+		Tips::do_try_state().unwrap();
 		assert_ok!(Tips::close_tip(RuntimeOrigin::signed(100), h.into()));
 		assert_eq!(Balances::reserved_balance(0), 0);
 		assert_eq!(Balances::free_balance(0), 110);
@@ -304,6 +307,8 @@ fn close_tip_works() {
 		assert_noop!(Tips::close_tip(RuntimeOrigin::signed(0), h.into()), Error::<Test>::Premature);
 
 		System::set_block_number(2);
+		Tips::do_try_state().unwrap();
+
 		assert_noop!(Tips::close_tip(RuntimeOrigin::none(), h.into()), BadOrigin);
 		assert_ok!(Tips::close_tip(RuntimeOrigin::signed(0), h.into()));
 		assert_eq!(Balances::free_balance(3), 10);
@@ -339,6 +344,7 @@ fn slash_tip_works() {
 		assert_noop!(Tips::slash_tip(RuntimeOrigin::signed(0), h), BadOrigin);
 
 		// can remove from root.
+		Tips::do_try_state().unwrap();
 		assert_ok!(Tips::slash_tip(RuntimeOrigin::root(), h));
 		assert_eq!(last_event(), TipEvent::TipSlashed { tip_hash: h, finder: 0, deposit: 12 });
 
@@ -358,9 +364,15 @@ fn retract_tip_works() {
 		assert_ok!(Tips::tip(RuntimeOrigin::signed(10), h, 10));
 		assert_ok!(Tips::tip(RuntimeOrigin::signed(11), h, 10));
 		assert_ok!(Tips::tip(RuntimeOrigin::signed(12), h, 10));
+		Tips::do_try_state().unwrap();
+
 		assert_noop!(Tips::retract_tip(RuntimeOrigin::signed(10), h), Error::<Test>::NotFinder);
 		assert_ok!(Tips::retract_tip(RuntimeOrigin::signed(0), h));
+		Tips::do_try_state().unwrap();
+
 		System::set_block_number(2);
+		Tips::do_try_state().unwrap();
+
 		assert_noop!(
 			Tips::close_tip(RuntimeOrigin::signed(0), h.into()),
 			Error::<Test>::UnknownTip
@@ -372,9 +384,15 @@ fn retract_tip_works() {
 		let h = tip_hash();
 		assert_ok!(Tips::tip(RuntimeOrigin::signed(11), h, 10));
 		assert_ok!(Tips::tip(RuntimeOrigin::signed(12), h, 10));
+		Tips::do_try_state().unwrap();
+
 		assert_noop!(Tips::retract_tip(RuntimeOrigin::signed(0), h), Error::<Test>::NotFinder);
 		assert_ok!(Tips::retract_tip(RuntimeOrigin::signed(10), h));
+		Tips::do_try_state().unwrap();
+
 		System::set_block_number(2);
+		Tips::do_try_state().unwrap();
+
 		assert_noop!(
 			Tips::close_tip(RuntimeOrigin::signed(10), h.into()),
 			Error::<Test>::UnknownTip
@@ -390,7 +408,11 @@ fn tip_median_calculation_works() {
 		let h = tip_hash();
 		assert_ok!(Tips::tip(RuntimeOrigin::signed(11), h, 10));
 		assert_ok!(Tips::tip(RuntimeOrigin::signed(12), h, 1000000));
+		Tips::do_try_state().unwrap();
+
 		System::set_block_number(2);
+		Tips::do_try_state().unwrap();
+
 		assert_ok!(Tips::close_tip(RuntimeOrigin::signed(0), h.into()));
 		assert_eq!(Balances::free_balance(3), 10);
 	});
@@ -409,7 +431,11 @@ fn tip_changing_works() {
 		assert_ok!(Tips::tip(RuntimeOrigin::signed(12), h, 1000));
 		assert_ok!(Tips::tip(RuntimeOrigin::signed(11), h, 100));
 		assert_ok!(Tips::tip(RuntimeOrigin::signed(10), h, 10));
+		Tips::do_try_state().unwrap();
+
 		System::set_block_number(2);
+		Tips::do_try_state().unwrap();
+
 		assert_ok!(Tips::close_tip(RuntimeOrigin::signed(0), h.into()));
 		assert_eq!(Balances::free_balance(3), 10);
 	});
@@ -603,8 +629,10 @@ fn report_awesome_and_tip_works_second_instance() {
 		assert_ok!(Tips1::tip(RuntimeOrigin::signed(11), h, 10));
 		assert_ok!(Tips1::tip(RuntimeOrigin::signed(12), h, 10));
 		assert_noop!(Tips1::tip(RuntimeOrigin::signed(9), h, 10), BadOrigin);
+		Tips::do_try_state().unwrap();
 
 		System::set_block_number(2);
+		Tips::do_try_state().unwrap();
 
 		assert_ok!(Tips1::close_tip(RuntimeOrigin::signed(100), h.into()));
 		// Treasury 1 unchanged

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;
@@ -193,6 +196,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.
@@ -645,6 +653,119 @@ 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 blo
+/// (`schedules_left` - 1), the vesting balance 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 vesting balance 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 (_, d) in Vesting::<T>::iter() {
+			let infos = d.to_vec();
+
+			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(<frame_system::Pallet<T>>::block_number());
+
+				if current_block_to_balance < starting_block {
+					Self::handle_before_schedule_starts(info, starting_block, schedules_left)?;
+				} else {
+					Self::handle_during_schedule(info, current_block_to_balance, schedules_left)?;
+				}
+			}
+		}
+		Ok(())
+	}
+
+	fn handle_before_schedule_starts(
+		info: &VestingInfo<BalanceOf<T>, BlockNumberFor<T>>,
+		starting_block: BalanceOf<T>,
+		schedules_left: BalanceOf<T>,
+	) -> Result<(), TryRuntimeError> {
+		let count = schedules_left.saturating_sub(starting_block);
+
+		if (info.locked() % info.per_block()).is_zero() {
+			ensure!(
+                info.locked_at::<T::BlockNumberToBalance>(frame_system::Pallet::<T>::block_number()) == (count * info.per_block()),
+                TryRuntimeError::Other("Before schedule starts, the vesting balance should be equal to the total per block releases")
+            );
+		} 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")
+			);
+		}
+
+		Ok(())
+	}
+
+	fn handle_during_schedule(
+		info: &VestingInfo<BalanceOf<T>, BlockNumberFor<T>>,
+		current_block_to_balance: BalanceOf<T>,
+		schedules_left: BalanceOf<T>,
+	) -> Result<(), TryRuntimeError> {
+		let current_block = frame_system::Pallet::<T>::block_number();
+
+		let still_vesting = info.locked_at::<T::BlockNumberToBalance>(current_block);
+
+		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")
+            );
+		} 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")
+				);
+			}
+		}
+		Ok(())
+	}
+}
+
 impl<T: Config> VestingSchedule<T::AccountId> for Pallet<T>
 where
 	BalanceOf<T>: MaybeSerializeDeserialize + Debug,

frame/vesting/src/mock.rs

@@ -148,7 +148,11 @@ 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(|| {
+			// sp_tracing::try_init_simple(); // Add this line
+			System::set_block_number(1);
+			Vesting::do_try_state().unwrap();
+		});
 		ext
 	}
 }