use std::{iter, mem};
use fvm_ipld_encoding::tuple::*;
use fvm_shared4::clock::ChainEpoch;
use fvm_shared4::econ::TokenAmount;
use itertools::{EitherOrBoth, Itertools};
use num_traits::Zero;
use super::{QuantSpec, VestSpec};
#[derive(Debug, Serialize_tuple, Deserialize_tuple)]
pub struct VestingFund {
pub epoch: ChainEpoch,
pub amount: TokenAmount,
}
#[derive(Serialize_tuple, Deserialize_tuple, Default)]
pub struct VestingFunds {
pub funds: Vec<VestingFund>,
}
impl VestingFunds {
pub fn new() -> Self {
Default::default()
}
pub fn unlock_vested_funds(&mut self, current_epoch: ChainEpoch) -> TokenAmount {
let i = self
.funds
.iter()
.position(|fund| fund.epoch >= current_epoch)
.unwrap_or(self.funds.len());
self.funds.drain(..i).map(|f| f.amount).sum()
}
pub fn add_locked_funds(
&mut self,
current_epoch: ChainEpoch,
vesting_sum: &TokenAmount,
proving_period_start: ChainEpoch,
spec: &VestSpec,
) {
let vest_begin = current_epoch + spec.initial_delay; let mut vested_so_far = TokenAmount::zero();
let mut epoch = vest_begin;
let new_funds = iter::from_fn(|| {
if vested_so_far >= *vesting_sum {
return None;
}
epoch += spec.step_duration;
let vest_epoch = QuantSpec {
unit: spec.quantization,
offset: proving_period_start,
}
.quantize_up(epoch);
let elapsed = vest_epoch - vest_begin;
let target_vest = if elapsed < spec.vest_period {
(vesting_sum * elapsed).div_floor(spec.vest_period)
} else {
vesting_sum.clone()
};
let vest_this_time = &target_vest - &vested_so_far;
vested_so_far = target_vest;
Some(VestingFund {
epoch: vest_epoch,
amount: vest_this_time,
})
});
let funds_len = self.funds.len();
let old_funds = mem::replace(&mut self.funds, Vec::with_capacity(funds_len));
self.funds.extend(
old_funds
.into_iter()
.merge_join_by(new_funds, |a, b| a.epoch.cmp(&b.epoch))
.map(|item| match item {
EitherOrBoth::Left(a) => a,
EitherOrBoth::Right(b) => b,
EitherOrBoth::Both(a, b) => VestingFund {
epoch: a.epoch,
amount: a.amount + b.amount,
},
}),
);
}
pub fn unlock_unvested_funds(
&mut self,
current_epoch: ChainEpoch,
target: &TokenAmount,
) -> TokenAmount {
let mut amount_unlocked = TokenAmount::zero();
let mut last = None;
let mut start = 0;
for (i, vf) in self.funds.iter_mut().enumerate() {
if &amount_unlocked >= target {
break;
}
if vf.epoch >= current_epoch {
let unlock_amount = std::cmp::min(target - &amount_unlocked, vf.amount.clone());
amount_unlocked += &unlock_amount;
let new_amount = &vf.amount - &unlock_amount;
if new_amount.is_zero() {
last = Some(i);
} else {
vf.amount = new_amount;
}
} else {
start = i + 1;
}
}
if let Some(end) = last {
self.funds.drain(start..=end);
}
amount_unlocked
}
}