1#![cfg_attr(not(feature = "std"), no_std)]
47
48mod benchmarking;
49
50#[cfg(test)]
51mod mock;
52#[cfg(test)]
53mod tests;
54mod vesting_info;
55
56pub mod migrations;
57pub mod weights;
58
59extern crate alloc;
60
61use alloc::vec::Vec;
62use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen};
63use core::{fmt::Debug, marker::PhantomData};
64use pezframe_support::{
65 dispatch::DispatchResult,
66 ensure,
67 storage::bounded_vec::BoundedVec,
68 traits::{
69 Currency, ExistenceRequirement, Get, LockIdentifier, LockableCurrency, VestedTransfer,
70 VestingSchedule, WithdrawReasons,
71 },
72 weights::Weight,
73};
74use pezframe_system::pezpallet_prelude::BlockNumberFor;
75use pezsp_runtime::{
76 traits::{
77 AtLeast32BitUnsigned, BlockNumberProvider, Bounded, Convert, MaybeSerializeDeserialize,
78 One, Saturating, StaticLookup, Zero,
79 },
80 DispatchError, RuntimeDebug,
81};
82use scale_info::TypeInfo;
83
84pub use pezpallet::*;
85pub use vesting_info::*;
86pub use weights::WeightInfo;
87
88type BalanceOf<T> =
89 <<T as Config>::Currency as Currency<<T as pezframe_system::Config>::AccountId>>::Balance;
90type MaxLocksOf<T> = <<T as Config>::Currency as LockableCurrency<
91 <T as pezframe_system::Config>::AccountId,
92>>::MaxLocks;
93type AccountIdLookupOf<T> = <<T as pezframe_system::Config>::Lookup as StaticLookup>::Source;
94
95const VESTING_ID: LockIdentifier = *b"vesting ";
96
97#[derive(Encode, Decode, Clone, Copy, PartialEq, Eq, RuntimeDebug, MaxEncodedLen, TypeInfo)]
100pub enum Releases {
101 V0,
102 V1,
103}
104
105impl Default for Releases {
106 fn default() -> Self {
107 Releases::V0
108 }
109}
110
111#[derive(Clone, Copy)]
113enum VestingAction {
114 Passive,
116 Remove { index: usize },
118 Merge { index1: usize, index2: usize },
120}
121
122impl VestingAction {
123 fn should_remove(&self, index: usize) -> bool {
125 match self {
126 Self::Passive => false,
127 Self::Remove { index: index1 } => *index1 == index,
128 Self::Merge { index1, index2 } => *index1 == index || *index2 == index,
129 }
130 }
131
132 fn pick_schedules<T: Config>(
134 &self,
135 schedules: Vec<VestingInfo<BalanceOf<T>, BlockNumberFor<T>>>,
136 ) -> impl Iterator<Item = VestingInfo<BalanceOf<T>, BlockNumberFor<T>>> + '_ {
137 schedules.into_iter().enumerate().filter_map(move |(index, schedule)| {
138 if self.should_remove(index) {
139 None
140 } else {
141 Some(schedule)
142 }
143 })
144 }
145}
146
147pub struct MaxVestingSchedulesGet<T>(PhantomData<T>);
149impl<T: Config> Get<u32> for MaxVestingSchedulesGet<T> {
150 fn get() -> u32 {
151 T::MAX_VESTING_SCHEDULES
152 }
153}
154
155#[pezframe_support::pezpallet]
156pub mod pezpallet {
157 use super::*;
158 use pezframe_support::pezpallet_prelude::*;
159 use pezframe_system::pezpallet_prelude::*;
160
161 #[pezpallet::config]
162 pub trait Config: pezframe_system::Config {
163 #[allow(deprecated)]
165 type RuntimeEvent: From<Event<Self>>
166 + IsType<<Self as pezframe_system::Config>::RuntimeEvent>;
167
168 type Currency: LockableCurrency<Self::AccountId>;
170
171 type BlockNumberToBalance: Convert<BlockNumberFor<Self>, BalanceOf<Self>>;
173
174 #[pezpallet::constant]
176 type MinVestedTransfer: Get<BalanceOf<Self>>;
177
178 type WeightInfo: WeightInfo;
180
181 type UnvestedFundsAllowedWithdrawReasons: Get<WithdrawReasons>;
184
185 type BlockNumberProvider: BlockNumberProvider<BlockNumber = BlockNumberFor<Self>>;
208
209 const MAX_VESTING_SCHEDULES: u32;
211 }
212
213 #[pezpallet::extra_constants]
214 impl<T: Config> Pezpallet<T> {
215 #[pezpallet::constant_name(MaxVestingSchedules)]
216 fn max_vesting_schedules() -> u32 {
217 T::MAX_VESTING_SCHEDULES
218 }
219 }
220
221 #[pezpallet::hooks]
222 impl<T: Config> Hooks<BlockNumberFor<T>> for Pezpallet<T> {
223 fn integrity_test() {
224 assert!(T::MAX_VESTING_SCHEDULES > 0, "`MaxVestingSchedules` must be greater than 0");
225 }
226 }
227
228 #[pezpallet::storage]
230 pub type Vesting<T: Config> = StorageMap<
231 _,
232 Blake2_128Concat,
233 T::AccountId,
234 BoundedVec<VestingInfo<BalanceOf<T>, BlockNumberFor<T>>, MaxVestingSchedulesGet<T>>,
235 >;
236
237 #[pezpallet::storage]
241 pub type StorageVersion<T: Config> = StorageValue<_, Releases, ValueQuery>;
242
243 #[pezpallet::pezpallet]
244 pub struct Pezpallet<T>(_);
245
246 #[pezpallet::genesis_config]
247 #[derive(pezframe_support::DefaultNoBound)]
248 pub struct GenesisConfig<T: Config> {
249 pub vesting: Vec<(T::AccountId, BlockNumberFor<T>, BlockNumberFor<T>, BalanceOf<T>)>,
250 }
251
252 #[pezpallet::genesis_build]
253 impl<T: Config> BuildGenesisConfig for GenesisConfig<T> {
254 fn build(&self) {
255 use pezsp_runtime::traits::Saturating;
256
257 StorageVersion::<T>::put(Releases::V1);
259
260 for &(ref who, begin, length, liquid) in self.vesting.iter() {
266 let balance = T::Currency::free_balance(who);
267 assert!(!balance.is_zero(), "Currencies must be init'd before vesting");
268 let locked = balance.saturating_sub(liquid);
270 let length_as_balance = T::BlockNumberToBalance::convert(length);
271 let per_block = locked / length_as_balance.max(pezsp_runtime::traits::One::one());
272 let vesting_info = VestingInfo::new(locked, per_block, begin);
273 if !vesting_info.is_valid() {
274 panic!("Invalid VestingInfo params at genesis")
275 };
276
277 Vesting::<T>::try_append(who, vesting_info)
278 .expect("Too many vesting schedules at genesis.");
279
280 let reasons =
281 WithdrawReasons::except(T::UnvestedFundsAllowedWithdrawReasons::get());
282
283 T::Currency::set_lock(VESTING_ID, who, locked, reasons);
284 }
285 }
286 }
287
288 #[pezpallet::event]
289 #[pezpallet::generate_deposit(pub(super) fn deposit_event)]
290 pub enum Event<T: Config> {
291 VestingCreated { account: T::AccountId, schedule_index: u32 },
293 VestingUpdated { account: T::AccountId, unvested: BalanceOf<T> },
296 VestingCompleted { account: T::AccountId },
298 }
299
300 #[pezpallet::error]
302 pub enum Error<T> {
303 NotVesting,
305 AtMaxVestingSchedules,
308 AmountLow,
310 ScheduleIndexOutOfBounds,
312 InvalidScheduleParams,
314 }
315
316 #[pezpallet::call]
317 impl<T: Config> Pezpallet<T> {
318 #[pezpallet::call_index(0)]
328 #[pezpallet::weight(T::WeightInfo::vest_locked(MaxLocksOf::<T>::get(), T::MAX_VESTING_SCHEDULES)
329 .max(T::WeightInfo::vest_unlocked(MaxLocksOf::<T>::get(), T::MAX_VESTING_SCHEDULES))
330 )]
331 pub fn vest(origin: OriginFor<T>) -> DispatchResult {
332 let who = ensure_signed(origin)?;
333 Self::do_vest(who)
334 }
335
336 #[pezpallet::call_index(1)]
348 #[pezpallet::weight(T::WeightInfo::vest_other_locked(MaxLocksOf::<T>::get(), T::MAX_VESTING_SCHEDULES)
349 .max(T::WeightInfo::vest_other_unlocked(MaxLocksOf::<T>::get(), T::MAX_VESTING_SCHEDULES))
350 )]
351 pub fn vest_other(origin: OriginFor<T>, target: AccountIdLookupOf<T>) -> DispatchResult {
352 ensure_signed(origin)?;
353 let who = T::Lookup::lookup(target)?;
354 Self::do_vest(who)
355 }
356
357 #[pezpallet::call_index(2)]
371 #[pezpallet::weight(
372 T::WeightInfo::vested_transfer(MaxLocksOf::<T>::get(), T::MAX_VESTING_SCHEDULES)
373 )]
374 pub fn vested_transfer(
375 origin: OriginFor<T>,
376 target: AccountIdLookupOf<T>,
377 schedule: VestingInfo<BalanceOf<T>, BlockNumberFor<T>>,
378 ) -> DispatchResult {
379 let transactor = ensure_signed(origin)?;
380 let target = T::Lookup::lookup(target)?;
381 Self::do_vested_transfer(&transactor, &target, schedule)
382 }
383
384 #[pezpallet::call_index(3)]
399 #[pezpallet::weight(
400 T::WeightInfo::force_vested_transfer(MaxLocksOf::<T>::get(), T::MAX_VESTING_SCHEDULES)
401 )]
402 pub fn force_vested_transfer(
403 origin: OriginFor<T>,
404 source: AccountIdLookupOf<T>,
405 target: AccountIdLookupOf<T>,
406 schedule: VestingInfo<BalanceOf<T>, BlockNumberFor<T>>,
407 ) -> DispatchResult {
408 ensure_root(origin)?;
409 let target = T::Lookup::lookup(target)?;
410 let source = T::Lookup::lookup(source)?;
411 Self::do_vested_transfer(&source, &target, schedule)
412 }
413
414 #[pezpallet::call_index(4)]
436 #[pezpallet::weight(
437 T::WeightInfo::not_unlocking_merge_schedules(MaxLocksOf::<T>::get(), T::MAX_VESTING_SCHEDULES)
438 .max(T::WeightInfo::unlocking_merge_schedules(MaxLocksOf::<T>::get(), T::MAX_VESTING_SCHEDULES))
439 )]
440 pub fn merge_schedules(
441 origin: OriginFor<T>,
442 schedule1_index: u32,
443 schedule2_index: u32,
444 ) -> DispatchResult {
445 let who = ensure_signed(origin)?;
446 if schedule1_index == schedule2_index {
447 return Ok(());
448 };
449 let schedule1_index = schedule1_index as usize;
450 let schedule2_index = schedule2_index as usize;
451
452 let schedules = Vesting::<T>::get(&who).ok_or(Error::<T>::NotVesting)?;
453 let merge_action =
454 VestingAction::Merge { index1: schedule1_index, index2: schedule2_index };
455
456 let (schedules, locked_now) = Self::exec_action(schedules.to_vec(), merge_action)?;
457
458 Self::write_vesting(&who, schedules)?;
459 Self::write_lock(&who, locked_now);
460
461 Ok(())
462 }
463
464 #[pezpallet::call_index(5)]
471 #[pezpallet::weight(
472 T::WeightInfo::force_remove_vesting_schedule(MaxLocksOf::<T>::get(), T::MAX_VESTING_SCHEDULES)
473 )]
474 pub fn force_remove_vesting_schedule(
475 origin: OriginFor<T>,
476 target: <T::Lookup as StaticLookup>::Source,
477 schedule_index: u32,
478 ) -> DispatchResultWithPostInfo {
479 ensure_root(origin)?;
480 let who = T::Lookup::lookup(target)?;
481
482 let schedules_count = Vesting::<T>::decode_len(&who).unwrap_or_default();
483 ensure!(schedule_index < schedules_count as u32, Error::<T>::InvalidScheduleParams);
484
485 Self::remove_vesting_schedule(&who, schedule_index)?;
486
487 Ok(Some(T::WeightInfo::force_remove_vesting_schedule(
488 MaxLocksOf::<T>::get(),
489 schedules_count as u32,
490 ))
491 .into())
492 }
493 }
494}
495
496impl<T: Config> Pezpallet<T> {
497 pub fn vesting(
499 account: T::AccountId,
500 ) -> Option<BoundedVec<VestingInfo<BalanceOf<T>, BlockNumberFor<T>>, MaxVestingSchedulesGet<T>>>
501 {
502 Vesting::<T>::get(account)
503 }
504
505 fn merge_vesting_info(
508 now: BlockNumberFor<T>,
509 schedule1: VestingInfo<BalanceOf<T>, BlockNumberFor<T>>,
510 schedule2: VestingInfo<BalanceOf<T>, BlockNumberFor<T>>,
511 ) -> Option<VestingInfo<BalanceOf<T>, BlockNumberFor<T>>> {
512 let schedule1_ending_block = schedule1.ending_block_as_balance::<T::BlockNumberToBalance>();
513 let schedule2_ending_block = schedule2.ending_block_as_balance::<T::BlockNumberToBalance>();
514 let now_as_balance = T::BlockNumberToBalance::convert(now);
515
516 match (schedule1_ending_block <= now_as_balance, schedule2_ending_block <= now_as_balance) {
518 (true, true) => return None,
520 (true, false) => return Some(schedule2),
523 (false, true) => return Some(schedule1),
524 _ => {},
526 }
527
528 let locked = schedule1
529 .locked_at::<T::BlockNumberToBalance>(now)
530 .saturating_add(schedule2.locked_at::<T::BlockNumberToBalance>(now));
531 debug_assert!(
534 !locked.is_zero(),
535 "merge_vesting_info validation checks failed to catch a locked of 0"
536 );
537
538 let ending_block = schedule1_ending_block.max(schedule2_ending_block);
539 let starting_block = now.max(schedule1.starting_block()).max(schedule2.starting_block());
540
541 let per_block = {
542 let duration = ending_block
543 .saturating_sub(T::BlockNumberToBalance::convert(starting_block))
544 .max(One::one());
545 (locked / duration).max(One::one())
546 };
547
548 let schedule = VestingInfo::new(locked, per_block, starting_block);
549 debug_assert!(schedule.is_valid(), "merge_vesting_info schedule validation check failed");
550
551 Some(schedule)
552 }
553
554 fn do_vested_transfer(
556 source: &T::AccountId,
557 target: &T::AccountId,
558 schedule: VestingInfo<BalanceOf<T>, BlockNumberFor<T>>,
559 ) -> DispatchResult {
560 ensure!(schedule.locked() >= T::MinVestedTransfer::get(), Error::<T>::AmountLow);
562 if !schedule.is_valid() {
563 return Err(Error::<T>::InvalidScheduleParams.into());
564 };
565
566 Self::can_add_vesting_schedule(
568 target,
569 schedule.locked(),
570 schedule.per_block(),
571 schedule.starting_block(),
572 )?;
573
574 T::Currency::transfer(source, target, schedule.locked(), ExistenceRequirement::AllowDeath)?;
575
576 let res = Self::add_vesting_schedule(
580 target,
581 schedule.locked(),
582 schedule.per_block(),
583 schedule.starting_block(),
584 );
585 debug_assert!(res.is_ok(), "Failed to add a schedule when we had to succeed.");
586
587 Ok(())
588 }
589
590 fn report_schedule_updates(
601 schedules: Vec<VestingInfo<BalanceOf<T>, BlockNumberFor<T>>>,
602 action: VestingAction,
603 ) -> (Vec<VestingInfo<BalanceOf<T>, BlockNumberFor<T>>>, BalanceOf<T>) {
604 let now = T::BlockNumberProvider::current_block_number();
605
606 let mut total_locked_now: BalanceOf<T> = Zero::zero();
607 let filtered_schedules = action
608 .pick_schedules::<T>(schedules)
609 .filter(|schedule| {
610 let locked_now = schedule.locked_at::<T::BlockNumberToBalance>(now);
611 let keep = !locked_now.is_zero();
612 if keep {
613 total_locked_now = total_locked_now.saturating_add(locked_now);
614 }
615 keep
616 })
617 .collect::<Vec<_>>();
618
619 (filtered_schedules, total_locked_now)
620 }
621
622 fn write_lock(who: &T::AccountId, total_locked_now: BalanceOf<T>) {
624 if total_locked_now.is_zero() {
625 T::Currency::remove_lock(VESTING_ID, who);
626 Self::deposit_event(Event::<T>::VestingCompleted { account: who.clone() });
627 } else {
628 let reasons = WithdrawReasons::except(T::UnvestedFundsAllowedWithdrawReasons::get());
629 T::Currency::set_lock(VESTING_ID, who, total_locked_now, reasons);
630 Self::deposit_event(Event::<T>::VestingUpdated {
631 account: who.clone(),
632 unvested: total_locked_now,
633 });
634 };
635 }
636
637 fn write_vesting(
639 who: &T::AccountId,
640 schedules: Vec<VestingInfo<BalanceOf<T>, BlockNumberFor<T>>>,
641 ) -> Result<(), DispatchError> {
642 let schedules: BoundedVec<
643 VestingInfo<BalanceOf<T>, BlockNumberFor<T>>,
644 MaxVestingSchedulesGet<T>,
645 > = schedules.try_into().map_err(|_| Error::<T>::AtMaxVestingSchedules)?;
646
647 if schedules.len() == 0 {
648 Vesting::<T>::remove(&who);
649 } else {
650 Vesting::<T>::insert(who, schedules)
651 }
652
653 Ok(())
654 }
655
656 fn do_vest(who: T::AccountId) -> DispatchResult {
658 let schedules = Vesting::<T>::get(&who).ok_or(Error::<T>::NotVesting)?;
659
660 let (schedules, locked_now) =
661 Self::exec_action(schedules.to_vec(), VestingAction::Passive)?;
662
663 Self::write_vesting(&who, schedules)?;
664 Self::write_lock(&who, locked_now);
665
666 Ok(())
667 }
668
669 fn exec_action(
672 schedules: Vec<VestingInfo<BalanceOf<T>, BlockNumberFor<T>>>,
673 action: VestingAction,
674 ) -> Result<(Vec<VestingInfo<BalanceOf<T>, BlockNumberFor<T>>>, BalanceOf<T>), DispatchError> {
675 let (schedules, locked_now) = match action {
676 VestingAction::Merge { index1: idx1, index2: idx2 } => {
677 let schedule1 = *schedules.get(idx1).ok_or(Error::<T>::ScheduleIndexOutOfBounds)?;
680 let schedule2 = *schedules.get(idx2).ok_or(Error::<T>::ScheduleIndexOutOfBounds)?;
681
682 let (mut schedules, mut locked_now) =
686 Self::report_schedule_updates(schedules.to_vec(), action);
687
688 let now = T::BlockNumberProvider::current_block_number();
689 if let Some(new_schedule) = Self::merge_vesting_info(now, schedule1, schedule2) {
690 schedules.push(new_schedule);
693 let new_schedule_locked =
695 new_schedule.locked_at::<T::BlockNumberToBalance>(now);
696 locked_now = locked_now.saturating_add(new_schedule_locked);
698 } (schedules, locked_now)
701 },
702 _ => Self::report_schedule_updates(schedules.to_vec(), action),
703 };
704
705 debug_assert!(
706 locked_now > Zero::zero() && schedules.len() > 0
707 || locked_now == Zero::zero() && schedules.len() == 0
708 );
709
710 Ok((schedules, locked_now))
711 }
712}
713
714impl<T: Config> VestingSchedule<T::AccountId> for Pezpallet<T>
715where
716 BalanceOf<T>: MaybeSerializeDeserialize + Debug,
717{
718 type Currency = T::Currency;
719 type Moment = BlockNumberFor<T>;
720
721 fn vesting_balance(who: &T::AccountId) -> Option<BalanceOf<T>> {
723 if let Some(v) = Vesting::<T>::get(who) {
724 let now = T::BlockNumberProvider::current_block_number();
725 let total_locked_now = v.iter().fold(Zero::zero(), |total, schedule| {
726 schedule.locked_at::<T::BlockNumberToBalance>(now).saturating_add(total)
727 });
728 Some(T::Currency::free_balance(who).min(total_locked_now))
729 } else {
730 None
731 }
732 }
733
734 fn add_vesting_schedule(
747 who: &T::AccountId,
748 locked: BalanceOf<T>,
749 per_block: BalanceOf<T>,
750 starting_block: BlockNumberFor<T>,
751 ) -> DispatchResult {
752 if locked.is_zero() {
753 return Ok(());
754 }
755
756 let vesting_schedule = VestingInfo::new(locked, per_block, starting_block);
757 if !vesting_schedule.is_valid() {
759 return Err(Error::<T>::InvalidScheduleParams.into());
760 };
761
762 let mut schedules = Vesting::<T>::get(who).unwrap_or_default();
763
764 ensure!(schedules.try_push(vesting_schedule).is_ok(), Error::<T>::AtMaxVestingSchedules);
767
768 debug_assert!(schedules.len() > 0, "schedules cannot be empty after insertion");
769 let schedule_index = schedules.len() - 1;
770 Self::deposit_event(Event::<T>::VestingCreated {
771 account: who.clone(),
772 schedule_index: schedule_index as u32,
773 });
774
775 let (schedules, locked_now) =
776 Self::exec_action(schedules.to_vec(), VestingAction::Passive)?;
777
778 Self::write_vesting(who, schedules)?;
779 Self::write_lock(who, locked_now);
780
781 Ok(())
782 }
783
784 fn can_add_vesting_schedule(
787 who: &T::AccountId,
788 locked: BalanceOf<T>,
789 per_block: BalanceOf<T>,
790 starting_block: BlockNumberFor<T>,
791 ) -> DispatchResult {
792 if !VestingInfo::new(locked, per_block, starting_block).is_valid() {
794 return Err(Error::<T>::InvalidScheduleParams.into());
795 }
796
797 ensure!(
798 (Vesting::<T>::decode_len(who).unwrap_or_default() as u32) < T::MAX_VESTING_SCHEDULES,
799 Error::<T>::AtMaxVestingSchedules
800 );
801
802 Ok(())
803 }
804
805 fn remove_vesting_schedule(who: &T::AccountId, schedule_index: u32) -> DispatchResult {
807 let schedules = Vesting::<T>::get(who).ok_or(Error::<T>::NotVesting)?;
808 let remove_action = VestingAction::Remove { index: schedule_index as usize };
809
810 let (schedules, locked_now) = Self::exec_action(schedules.to_vec(), remove_action)?;
811
812 Self::write_vesting(who, schedules)?;
813 Self::write_lock(who, locked_now);
814 Ok(())
815 }
816}
817
818impl<T: Config> VestedTransfer<T::AccountId> for Pezpallet<T>
821where
822 BalanceOf<T>: MaybeSerializeDeserialize + Debug,
823{
824 type Currency = T::Currency;
825 type Moment = BlockNumberFor<T>;
826
827 fn vested_transfer(
828 source: &T::AccountId,
829 target: &T::AccountId,
830 locked: BalanceOf<T>,
831 per_block: BalanceOf<T>,
832 starting_block: BlockNumberFor<T>,
833 ) -> DispatchResult {
834 use pezframe_support::storage::{with_transaction, TransactionOutcome};
835 let schedule = VestingInfo::new(locked, per_block, starting_block);
836 with_transaction(|| -> TransactionOutcome<DispatchResult> {
837 let result = Self::do_vested_transfer(source, target, schedule);
838
839 match &result {
840 Ok(()) => TransactionOutcome::Commit(result),
841 _ => TransactionOutcome::Rollback(result),
842 }
843 })
844 }
845}