Lazy reaping (#4895)

* Squash and rebase from gav-lazy-reaping

* Bump version

* Bump runtime again

* Docs.

* Remove old functions

* Update frame/balances/src/lib.rs

Co-Authored-By: Shawn Tabrizi <shawntabrizi@gmail.com>

* Update frame/contracts/src/lib.rs

Co-Authored-By: Shawn Tabrizi <shawntabrizi@gmail.com>

* Warnings

* Bump runtime version

* Update frame/democracy/src/lib.rs

Co-Authored-By: Shawn Tabrizi <shawntabrizi@gmail.com>

* Update frame/system/src/lib.rs

* Clean up OnReapAccount

* Use frame_support debug

* Bump spec

* Renames and fix

* Fix

* Fix rename

* Fix

* Increase time for test

Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>
Co-authored-by: Benjamin Kampmann <ben.kampmann@googlemail.com>

substrate/frame/staking/src/lib.rs

@@ -259,6 +259,7 @@ use codec::{HasCompact, Encode, Decode};
 use frame_support::{
 	decl_module, decl_event, decl_storage, ensure, decl_error,
 	weights::SimpleDispatchInfo,
+	dispatch::DispatchResult,
 	traits::{
 		Currency, LockIdentifier, LockableCurrency,
 		WithdrawReasons, OnUnbalanced, Imbalance, Get, Time
@@ -282,7 +283,6 @@ use sp_runtime::{Serialize, Deserialize};
 use frame_system::{self as system, ensure_signed, ensure_root};
 
 use sp_phragmen::ExtendedBalance;
-use frame_support::traits::OnReapAccount;
 
 const DEFAULT_MINIMUM_VALIDATOR_COUNT: u32 = 4;
 const MAX_NOMINATIONS: usize = 16;
@@ -831,6 +831,8 @@ decl_error! {
 		NoMoreChunks,
 		/// Can not rebond without unlocking chunks.
 		NoUnlockChunk,
+		/// Attempting to target a stash that still has funds.
+		FundedTarget,
 	}
 }
 
@@ -900,6 +902,8 @@ decl_module! {
 			<Bonded<T>>::insert(&stash, &controller);
 			<Payee<T>>::insert(&stash, payee);
 
+			system::Module::<T>::inc_ref(&stash);
+
 			let stash_balance = T::Currency::free_balance(&stash);
 			let value = value.min(stash_balance);
 			let item = StakingLedger { stash, total: value, active: value, unlocking: vec![] };
@@ -1013,10 +1017,10 @@ decl_module! {
 				// portion to fall below existential deposit + will have no more unlocking chunks
 				// left. We can now safely remove this.
 				let stash = ledger.stash;
+				// remove all staking-related information.
+				Self::kill_stash(&stash)?;
 				// remove the lock.
 				T::Currency::remove_lock(STAKING_ID, &stash);
-				// remove all staking-related information.
-				Self::kill_stash(&stash);
 			} else {
 				// This was the consequence of a partial unbond. just update the ledger and move on.
 				Self::update_ledger(&controller, &ledger);
@@ -1187,10 +1191,11 @@ decl_module! {
 		fn force_unstake(origin, stash: T::AccountId) {
 			ensure_root(origin)?;
 
+			// remove all staking-related information.
+			Self::kill_stash(&stash)?;
+
 			// remove the lock.
 			T::Currency::remove_lock(STAKING_ID, &stash);
-			// remove all staking-related information.
-			Self::kill_stash(&stash);
 		}
 
 		/// Force there to be a new era at the end of sessions indefinitely.
@@ -1254,9 +1259,22 @@ decl_module! {
 			);
 
 			let ledger = ledger.rebond(value);
-
 			Self::update_ledger(&controller, &ledger);
 		}
+
+		/// Remove all data structure concerning a staker/stash once its balance is zero.
+		/// This is essentially equivalent to `withdraw_unbonded` except it can be called by anyone
+		/// and the target `stash` must have no funds left.
+		///
+		/// This can be called from any origin.
+		///
+		/// - `stash`: The stash account to reap. Its balance must be zero.
+		fn reap_stash(_origin, stash: T::AccountId) {
+			Self::ensure_storage_upgraded();
+			ensure!(T::Currency::total_balance(&stash).is_zero(), Error::<T>::FundedTarget);
+			Self::kill_stash(&stash)?;
+			T::Currency::remove_lock(STAKING_ID, &stash);
+		}
 	}
 }
 
@@ -1585,18 +1603,22 @@ impl<T: Trait> Module<T> {
 	///
 	/// Assumes storage is upgraded before calling.
 	///
-	/// This is called :
-	/// - Immediately when an account's balance falls below existential deposit.
+	/// This is called:
 	/// - after a `withdraw_unbond()` call that frees all of a stash's bonded balance.
-	fn kill_stash(stash: &T::AccountId) {
-		if let Some(controller) = <Bonded<T>>::take(stash) {
-			<Ledger<T>>::remove(&controller);
-		}
+	/// - through `reap_stash()` if the balance has fallen to zero (through slashing).
+	fn kill_stash(stash: &T::AccountId) -> DispatchResult {
+		let controller = Bonded::<T>::take(stash).ok_or(Error::<T>::NotStash)?;
+		<Ledger<T>>::remove(&controller);
+
 		<Payee<T>>::remove(stash);
 		<Validators<T>>::remove(stash);
 		<Nominators<T>>::remove(stash);
 
 		slashing::clear_stash_metadata::<T>(stash);
+
+		system::Module::<T>::dec_ref(stash);
+
+		Ok(())
 	}
 
 	/// Add reward points to validators using their stash account ID.
@@ -1676,13 +1698,6 @@ impl<T: Trait> SessionManager<T::AccountId, Exposure<T::AccountId, BalanceOf<T>>
 	}
 }
 
-impl<T: Trait> OnReapAccount<T::AccountId> for Module<T> {
-	fn on_reap_account(stash: &T::AccountId) {
-		Self::ensure_storage_upgraded();
-		Self::kill_stash(stash);
-	}
-}
-
 /// Add reward points to block authors:
 /// * 20 points to the block producer for producing a (non-uncle) block in the relay chain,
 /// * 2 points to the block producer for each reference to a previously unreferenced uncle, and