feat: implement migration to v2 and enhance boosted supply calculations

pallets/bb-bnc/src/lib.rs

@@ -82,7 +82,7 @@ const BB_LOCK_ID: LockIdentifier = *b"bbbnclck";
 const MARKUP_LOCK_ID: LockIdentifier = *b"bbbncmkp";
 pub const BB_BNC_SYSTEM_POOL_ID: PoolId = u32::MAX;
 pub type PositionId = u128;
-const STORAGE_VERSION: StorageVersion = StorageVersion::new(1);
+const STORAGE_VERSION: StorageVersion = StorageVersion::new(2);
 
 #[derive(
 	Clone,
@@ -452,10 +452,22 @@ pub mod pallet {
 		RewardTooSmallToLock,
 	}
 
-	/// Total supply of locked tokens
+	/// Raw total amount of locked bbBNC principal.
+	///
+	/// This excludes markup and is used for lock accounting invariants with
+	/// `Locked` and `UserLocked`. Reward distribution must not use this as the
+	/// denominator when markup is enabled.
 	#[pallet::storage]
 	pub type Supply<T: Config> = StorageValue<_, BalanceOf<T>, ValueQuery>;
 
+	/// Current total amount of locked bbBNC after applying user markup.
+	///
+	/// This is the O(1) source of truth for writing global bbBNC point
+	/// snapshots. `PointHistory.amount` records this value at checkpoints so
+	/// `total_supply` uses the same boosted unit as user balances.
+	#[pallet::storage]
+	pub type BoostedSupply<T: Config> = StorageValue<_, BalanceOf<T>, ValueQuery>;
+
 	/// Configurations
 	#[pallet::storage]
 	pub type BbConfigs<T: Config> =
@@ -758,7 +770,7 @@ pub mod pallet {
 
 	#[cfg(any(test, feature = "try-runtime"))]
 	impl<T: Config> Pallet<T> {
-		fn do_try_state() -> Result<(), sp_runtime::TryRuntimeError> {
+		pub(crate) fn do_try_state() -> Result<(), sp_runtime::TryRuntimeError> {
 			for (owner, positions) in UserPositions::<T>::iter() {
 				let mut computed_user_locked = BalanceOf::<T>::zero();
 				for position in positions.iter() {
@@ -789,6 +801,22 @@ pub mod pallet {
 				sp_runtime::TryRuntimeError::Other("UserLocked total mismatch with Supply")
 			);
 
+			let mut boosted_supply_sum = BalanceOf::<T>::zero();
+			for position in 0..Position::<T>::get() {
+				let epoch = UserPointEpoch::<T>::get(position);
+				if epoch.is_zero() {
+					continue;
+				}
+				boosted_supply_sum = boosted_supply_sum
+					.saturating_add(UserPointHistory::<T>::get(position, epoch).amount);
+			}
+			frame_support::ensure!(
+				boosted_supply_sum == BoostedSupply::<T>::get(),
+				sp_runtime::TryRuntimeError::Other(
+					"UserPointHistory total mismatch with BoostedSupply"
+				)
+			);
+
 			for (end_block, position, _) in ExpiringPositions::<T>::iter() {
 				let lock = Locked::<T>::get(position);
 				frame_support::ensure!(
@@ -1170,13 +1198,14 @@ pub mod pallet {
 	}
 
 	impl<T: Config> Pallet<T> {
-		pub fn checkpoint(
+		pub(crate) fn checkpoint(
 			who: &AccountIdOf<T>,
 			position: PositionId,
 			old_locked: LockedBalance<BalanceOf<T>, BlockNumberFor<T>>,
 			new_locked: LockedBalance<BalanceOf<T>, BlockNumberFor<T>>,
 		) -> DispatchResult {
 			Self::update_reward_all(who)?;
+			Self::update_boosted_supply(&old_locked, &new_locked)?;
 
 			let mut u_old = Point::<BalanceOf<T>, BlockNumberFor<T>>::default();
 			let mut u_new = Point::<BalanceOf<T>, BlockNumberFor<T>>::default();
@@ -1245,7 +1274,7 @@ pub mod pallet {
 			if g_epoch > U256::zero() {
 				last_point = PointHistory::<T>::get(g_epoch);
 			} else {
-				last_point.amount = Supply::<T>::get();
+				last_point.amount = BoostedSupply::<T>::get();
 			}
 			let mut last_checkpoint = last_point.block;
 			let mut t_i: BlockNumberFor<T> = last_checkpoint
@@ -1299,7 +1328,7 @@ pub mod pallet {
 
 				// Fill for the current block, if applicable
 				if t_i == current_block_number {
-					last_point.amount = Supply::<T>::get();
+					last_point.amount = BoostedSupply::<T>::get();
 					break;
 				} else {
 					PointHistory::<T>::insert(g_epoch, last_point);
@@ -1360,6 +1389,31 @@ pub mod pallet {
 			Ok(())
 		}
 
+		fn update_boosted_supply(
+			old_locked: &LockedBalance<BalanceOf<T>, BlockNumberFor<T>>,
+			new_locked: &LockedBalance<BalanceOf<T>, BlockNumberFor<T>>,
+		) -> DispatchResult {
+			// `old_locked` and `new_locked` are expected to use markup-adjusted
+			// amounts. Keep raw principal accounting in `Supply`.
+			let old_amount = old_locked.amount;
+			let new_amount = new_locked.amount;
+			if new_amount >= old_amount {
+				BoostedSupply::<T>::try_mutate(|supply| -> DispatchResult {
+					*supply = supply
+						.checked_add(new_amount.saturating_sub(old_amount))
+						.ok_or(ArithmeticError::Overflow)?;
+					Ok(())
+				})
+			} else {
+				BoostedSupply::<T>::try_mutate(|supply| -> DispatchResult {
+					*supply = supply
+						.checked_sub(old_amount.saturating_sub(new_amount))
+						.ok_or(ArithmeticError::Underflow)?;
+					Ok(())
+				})
+			}
+		}
+
 		pub fn deposit_for_inner(
 			who: &AccountIdOf<T>,
 			position: PositionId,
@@ -1399,12 +1453,11 @@ pub mod pallet {
 				Self::set_ve_locked(who, new_locked_balance)?;
 			}
 
-			Self::markup_calc(
+			Self::checkpoint_with_current_markup(
 				who,
 				position,
 				old_locked.clone(),
 				locked.clone(),
-				UserMarkupInfos::<T>::get(who).as_ref(),
 			)?;
 
 			Self::deposit_event(Event::Minted {
@@ -1577,28 +1630,65 @@ pub mod pallet {
 			Ok(balance)
 		}
 
-		pub fn markup_calc(
+		pub(crate) fn markup_calc(
 			who: &AccountIdOf<T>,
 			position: PositionId,
-			mut old_locked: LockedBalance<BalanceOf<T>, BlockNumberFor<T>>,
-			mut new_locked: LockedBalance<BalanceOf<T>, BlockNumberFor<T>>,
+			old_locked: LockedBalance<BalanceOf<T>, BlockNumberFor<T>>,
+			new_locked: LockedBalance<BalanceOf<T>, BlockNumberFor<T>>,
 			user_markup_info: Option<&UserMarkupInfo>,
 		) -> DispatchResult {
 			if let Some(info) = user_markup_info {
-				old_locked.amount = info
-					.old_markup_coefficient
-					.checked_mul_int(old_locked.amount)
-					.and_then(|x| x.checked_add(old_locked.amount))
-					.ok_or(ArithmeticError::Overflow)?;
-				new_locked.amount = info
-					.markup_coefficient
-					.checked_mul_int(new_locked.amount)
-					.and_then(|x| x.checked_add(new_locked.amount))
-					.ok_or(ArithmeticError::Overflow)?;
+				Self::checkpoint_with_markup_coefficients(
+					who,
+					position,
+					old_locked,
+					new_locked,
+					info.old_markup_coefficient,
+					info.markup_coefficient,
+				)
+			} else {
+				Self::checkpoint(who, position, old_locked, new_locked)
 			}
+		}
 
-			Self::checkpoint(who, position, old_locked.clone(), new_locked.clone())?;
-			Ok(())
+		fn checkpoint_with_current_markup(
+			who: &AccountIdOf<T>,
+			position: PositionId,
+			old_locked: LockedBalance<BalanceOf<T>, BlockNumberFor<T>>,
+			new_locked: LockedBalance<BalanceOf<T>, BlockNumberFor<T>>,
+		) -> DispatchResult {
+			let current_markup_coefficient = UserMarkupInfos::<T>::get(who)
+				.map(|info| info.markup_coefficient)
+				.unwrap_or_else(FixedU128::zero);
+
+			Self::checkpoint_with_markup_coefficients(
+				who,
+				position,
+				old_locked,
+				new_locked,
+				current_markup_coefficient,
+				current_markup_coefficient,
+			)
+		}
+
+		fn checkpoint_with_markup_coefficients(
+			who: &AccountIdOf<T>,
+			position: PositionId,
+			mut old_locked: LockedBalance<BalanceOf<T>, BlockNumberFor<T>>,
+			mut new_locked: LockedBalance<BalanceOf<T>, BlockNumberFor<T>>,
+			old_markup_coefficient: FixedU128,
+			new_markup_coefficient: FixedU128,
+		) -> DispatchResult {
+			old_locked.amount = old_markup_coefficient
+				.checked_mul_int(old_locked.amount)
+				.and_then(|x| x.checked_add(old_locked.amount))
+				.ok_or(ArithmeticError::Overflow)?;
+			new_locked.amount = new_markup_coefficient
+				.checked_mul_int(new_locked.amount)
+				.and_then(|x| x.checked_add(new_locked.amount))
+				.ok_or(ArithmeticError::Overflow)?;
+
+			Self::checkpoint(who, position, old_locked, new_locked)
 		}
 
 		pub fn deposit_markup_inner(
@@ -2034,7 +2124,7 @@ pub mod pallet {
 				}
 			}
 
-			Self::checkpoint(who, position, old_locked, locked.clone())?;
+			Self::checkpoint_with_current_markup(who, position, old_locked, locked.clone())?;
 
 			T::FarmingInfo::refresh_gauge_pool(who)?;
 

pallets/bb-bnc/src/migrations/mod.rs

@@ -17,3 +17,4 @@
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
 pub mod v1;
+pub mod v2;

pallets/bb-bnc/src/migrations/v1.rs

@@ -164,3 +164,68 @@ impl<T: Config> OnRuntimeUpgrade for MigrateToV1<T> {
 		Ok(())
 	}
 }
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use crate::{mock::*, Pallet as BbBNCPallet};
+	use frame_support::traits::{OnRuntimeUpgrade, StorageVersion};
+	use sp_core::ConstU32;
+
+	#[test]
+	fn skips_when_in_code_version_is_newer() {
+		ExtBuilder::default().build().execute_with(|| {
+			StorageVersion::new(0).put::<BbBNCPallet<Runtime>>();
+			v0::ExpiringPositions::<Runtime>::insert(
+				10,
+				BoundedVec::<PositionId, ConstU32<100>>::try_from(vec![1, 2, 3]).unwrap(),
+			);
+			v0::ExpiringPositions::<Runtime>::insert(
+				20,
+				BoundedVec::<PositionId, ConstU32<100>>::try_from(vec![4, 5]).unwrap(),
+			);
+
+			let _ = MigrateToV1::<Runtime>::on_runtime_upgrade();
+
+			assert_eq!(BbBNCPallet::<Runtime>::on_chain_storage_version(), 0);
+			assert!(frame_support::storage::unhashed::exists(
+				&v0::ExpiringPositions::<Runtime>::hashed_key_for(10)
+			));
+			assert!(frame_support::storage::unhashed::exists(
+				&v0::ExpiringPositions::<Runtime>::hashed_key_for(20)
+			));
+			assert_eq!(ExpiringPositions::<Runtime>::iter().count(), 0);
+
+			let _ = MigrateToV1::<Runtime>::on_runtime_upgrade();
+
+			assert_eq!(BbBNCPallet::<Runtime>::on_chain_storage_version(), 0);
+			assert_eq!(ExpiringPositions::<Runtime>::iter().count(), 0);
+		});
+	}
+
+	#[test]
+	fn empty_storage_skips_when_in_code_version_is_newer() {
+		ExtBuilder::default().build().execute_with(|| {
+			StorageVersion::new(0).put::<BbBNCPallet<Runtime>>();
+
+			let _ = MigrateToV1::<Runtime>::on_runtime_upgrade();
+
+			assert_eq!(BbBNCPallet::<Runtime>::on_chain_storage_version(), 0);
+			assert_eq!(ExpiringPositions::<Runtime>::iter().count(), 0);
+		});
+	}
+
+	#[cfg(feature = "try-runtime")]
+	#[test]
+	fn try_runtime_rejects_when_in_code_version_is_newer() {
+		ExtBuilder::default().build().execute_with(|| {
+			StorageVersion::new(0).put::<BbBNCPallet<Runtime>>();
+			v0::ExpiringPositions::<Runtime>::insert(
+				10,
+				BoundedVec::<PositionId, ConstU32<100>>::try_from(vec![1, 2, 3]).unwrap(),
+			);
+
+			assert!(MigrateToV1::<Runtime>::pre_upgrade().is_err());
+		});
+	}
+}