From 108e2b57d99f44efadcee7b8e1ec0f1bda0e7ef2 Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Thu, 12 Oct 2023 00:26:35 -0400
Subject: [PATCH] Add claim_deallocation to the staking pallet

---
 substrate/staking/pallet/src/lib.rs        | 19 ++++++++++++++++---
 substrate/validator-sets/pallet/src/lib.rs |  1 +
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/substrate/staking/pallet/src/lib.rs b/substrate/staking/pallet/src/lib.rs
index 32bdfb73..83aceb8d 100644
--- a/substrate/staking/pallet/src/lib.rs
+++ b/substrate/staking/pallet/src/lib.rs
@@ -13,12 +13,13 @@ pub mod pallet {
 
   use serai_primitives::{NetworkId, Amount, PublicKey};
 
-  use validator_sets_pallet::{Config as VsConfig, Pallet as VsPallet};
+  use validator_sets_pallet::{primitives::Session, Config as VsConfig, Pallet as VsPallet};
   use pallet_session::{Config as SessionConfig, SessionManager};
 
   #[pallet::error]
   pub enum Error<T> {
     StakeUnavilable,
+    NoDeallocation,
   }
 
   // TODO: Event
@@ -76,7 +77,6 @@ pub mod pallet {
       })
     }
 
-    #[allow(unused)] // TODO
     fn deallocate_internal(account: &T::AccountId, amount: u64) -> Result<(), Error<T>> {
       Allocated::<T>::try_mutate(account, |allocated| {
         if *allocated < amount {
@@ -150,7 +150,20 @@ pub mod pallet {
       Ok(())
     }
 
-    // TODO: Add a function to reclaim deallocated funds
+    #[pallet::call_index(4)]
+    #[pallet::weight((0, DispatchClass::Operational))] // TODO
+    pub fn claim_deallocation(
+      origin: OriginFor<T>,
+      network: NetworkId,
+      session: Session,
+    ) -> DispatchResult {
+      let account = ensure_signed(origin)?;
+      let Some(amount) = VsPallet::<T>::take_deallocatable_amount(network, session, account) else {
+        Err(Error::<T>::NoDeallocation)?
+      };
+      Self::deallocate_internal(&account, amount.0)?;
+      Ok(())
+    }
   }
 
   // Call order is end_session(i - 1) -> start_session(i) -> new_session(i + 1)
diff --git a/substrate/validator-sets/pallet/src/lib.rs b/substrate/validator-sets/pallet/src/lib.rs
index 77dd3117..82c5f531 100644
--- a/substrate/validator-sets/pallet/src/lib.rs
+++ b/substrate/validator-sets/pallet/src/lib.rs
@@ -163,6 +163,7 @@ pub mod pallet {
   #[pallet::getter(fn keys)]
   pub type Keys<T: Config> = StorageMap<_, Twox64Concat, ValidatorSet, KeyPair, OptionQuery>;
 
+  // TODO: Expand
   #[pallet::event]
   #[pallet::generate_deposit(pub(super) fn deposit_event)]
   pub enum Event<T: Config> {