cleanup the storage (#560)

Author: zjb0807

vesting/src/lib.rs

@@ -283,6 +283,8 @@ impl<T: Config> Pallet<T> {
 	fn do_claim(who: &T::AccountId) -> BalanceOf<T> {
 		let locked = Self::locked_balance(who);
 		if locked.is_zero() {
+			// cleanup the storage and unlock the fund
+			<VestingSchedules<T>>::remove(who);
 			T::Currency::remove_lock(VESTING_LOCK_ID, who);
 		} else {
 			T::Currency::set_lock(VESTING_LOCK_ID, who, locked, WithdrawReasons::all());
@@ -331,6 +333,13 @@ impl<T: Config> Pallet<T> {
 			.try_into()
 			.map_err(|_| Error::<T>::MaxVestingSchedulesExceeded)?;
 
+		// empty vesting schedules cleanup the storage and unlock the fund
+		if bounded_schedules.len().is_zero() {
+			<VestingSchedules<T>>::remove(who);
+			T::Currency::remove_lock(VESTING_LOCK_ID, who);
+			return Ok(());
+		}
+
 		let total_amount = bounded_schedules
 			.iter()
 			.try_fold::<_, _, Result<BalanceOf<T>, DispatchError>>(Zero::zero(), |acc_amount, schedule| {

vesting/src/tests.rs

@@ -218,13 +218,15 @@ fn claim_works() {
 		assert!(PalletBalances::transfer(Origin::signed(BOB), ALICE, 10).is_err());
 		// unlocked after claiming
 		assert_ok!(Vesting::claim(Origin::signed(BOB)));
+		assert_eq!(VestingSchedules::<Runtime>::contains_key(BOB), true);
 		assert_ok!(PalletBalances::transfer(Origin::signed(BOB), ALICE, 10));
 		// more are still locked
 		assert!(PalletBalances::transfer(Origin::signed(BOB), ALICE, 1).is_err());
 
 		MockBlockNumberProvider::set(21);
 		// claim more
 		assert_ok!(Vesting::claim(Origin::signed(BOB)));
+		assert_eq!(VestingSchedules::<Runtime>::contains_key(BOB), false);
 		assert_ok!(PalletBalances::transfer(Origin::signed(BOB), ALICE, 10));
 		// all used up
 		assert_eq!(PalletBalances::free_balance(BOB), 0);
@@ -264,6 +266,20 @@ fn update_vesting_schedules_works() {
 		MockBlockNumberProvider::set(21);
 		assert_ok!(Vesting::claim(Origin::signed(BOB)));
 		assert_ok!(PalletBalances::transfer(Origin::signed(BOB), ALICE, 10));
+
+		// empty vesting schedules cleanup the storage and unlock the fund
+		assert_eq!(VestingSchedules::<Runtime>::contains_key(BOB), true);
+		assert_eq!(
+			PalletBalances::locks(&BOB).get(0),
+			Some(&BalanceLock {
+				id: VESTING_LOCK_ID,
+				amount: 10u64,
+				reasons: Reasons::All,
+			})
+		);
+		assert_ok!(Vesting::update_vesting_schedules(Origin::root(), BOB, vec![]));
+		assert_eq!(VestingSchedules::<Runtime>::contains_key(BOB), false);
+		assert_eq!(PalletBalances::locks(&BOB), vec![]);
 	});
 }