From 23e47ee73dadf4fe817ef535d093fe8a6109eb7a Mon Sep 17 00:00:00 2001
From: kingsleydon <10992364+kingsleydon@users.noreply.github.com>
Date: Wed, 11 Dec 2024 18:57:32 +0800
Subject: [PATCH] ensure invest pools when vault withdraw from stake pools

---
 pallets/phala/src/compute/stake_pool_v2.rs | 30 +++++++++++++++++-----
 pallets/phala/src/compute/vault.rs         |  2 +-
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/pallets/phala/src/compute/stake_pool_v2.rs b/pallets/phala/src/compute/stake_pool_v2.rs
index f0e4fbbae..b1dc4fd2e 100644
--- a/pallets/phala/src/compute/stake_pool_v2.rs
+++ b/pallets/phala/src/compute/stake_pool_v2.rs
@@ -780,13 +780,6 @@ pub mod pallet {
 				who.clone(),
 				pool_info.basepool.pid,
 			)?;
-			let nft_id = maybe_nft_id.ok_or(Error::<T>::NoNftToWithdraw)?;
-			// The nft instance must be wrote to Nft storage at the end of the function
-			// this nft's property shouldn't be accessed or wrote again from storage before set_nft_attr
-			// is called. Or the property of the nft will be overwrote incorrectly.
-			let mut nft_guard =
-				base_pool::Pallet::<T>::get_nft_attr_guard(pool_info.basepool.cid, nft_id)?;
-			let nft = &mut nft_guard.attr;
 			let in_queue_shares = match pool_info
 				.basepool
 				.withdraw_queue
@@ -803,10 +796,33 @@ pub mod pallet {
 				}
 				None => Zero::zero(),
 			};
+			ensure!(maybe_nft_id.is_some() || in_queue_shares > Zero::zero(), Error::<T>::NoNftToWithdraw);
+			let nft_id = match maybe_nft_id {
+				Some(nft_id) => nft_id,
+				// An nft is necessary to initiate a smaller withdrawal
+				None => base_pool::Pallet::<T>::mint_nft(
+					pool_info.basepool.cid,
+					who.clone(),
+					Zero::zero(),
+					pool_info.basepool.pid,
+				)?,
+			};
+			// The nft instance must be wrote to Nft storage at the end of the function
+			// this nft's property shouldn't be accessed or wrote again from storage before set_nft_attr
+			// is called. Or the property of the nft will be overwrote incorrectly.
+			let mut nft_guard =
+				base_pool::Pallet::<T>::get_nft_attr_guard(pool_info.basepool.cid, nft_id)?;
+			let nft = &mut nft_guard.attr;
 			ensure!(
 				base_pool::is_nondust_balance(shares) && (shares <= nft.shares + in_queue_shares),
 				Error::<T>::InvalidWithdrawalAmount
 			);
+			if let Some(vault_pid) = as_vault {
+				let mut vault_info = ensure_vault::<T>(vault_pid)?;
+				if !vault_info.invest_pools.contains(&pid) {
+					vault_info.invest_pools.push(pid);
+				}
+			}
 			base_pool::Pallet::<T>::try_withdraw(
 				&mut pool_info.basepool,
 				nft,
diff --git a/pallets/phala/src/compute/vault.rs b/pallets/phala/src/compute/vault.rs
index bf0003af4..a8fca5cca 100644
--- a/pallets/phala/src/compute/vault.rs
+++ b/pallets/phala/src/compute/vault.rs
@@ -429,7 +429,7 @@ pub mod pallet {
 					if base_pool::is_nondust_balance(property.shares) {
 						true
 					} else {
-						if (!base_pool::is_nondust_balance(total_shares)) {
+						if !base_pool::is_nondust_balance(total_shares) {
 							let _ = base_pool::Pallet::<T>::burn_nft(
 								&base_pool::pallet_id::<T::AccountId>(),
 								stake_pool.basepool.cid,