hylo-core 2.0.1

use std::cmp::Ordering;

use anchor_lang::prelude::{
  borsh, AnchorDeserialize, AnchorSerialize, InitSpace,
};
use fix::prelude::*;
use serde::{Deserialize, Serialize};

use crate::error::CoreError;
use crate::error::CoreError::{
  WithdrawalLimitArithmetic, WithdrawalLimitExceededForEpoch,
  WithdrawalLimitInvalidEpoch, WithdrawalLimitValidation,
};
use crate::virtual_stablecoin::VirtualStablecoin;

/// Per-epoch withdrawal window, reset lazily on epoch rollover.
#[derive(
  Debug,
  Clone,
  Copy,
  AnchorSerialize,
  AnchorDeserialize,
  InitSpace,
  Serialize,
  Deserialize,
  PartialEq,
  Eq,
)]
pub struct WithdrawalLimiter {
  pub limit: UFixValue64,
  withdrawal_ledger: VirtualStablecoin,
  epoch: u64,
}

impl WithdrawalLimiter {
  #[cfg(test)]
  fn new(limit: UFixValue64, epoch: u64) -> WithdrawalLimiter {
    WithdrawalLimiter {
      limit,
      withdrawal_ledger: VirtualStablecoin::new(),
      epoch,
    }
  }

  /// Converts configured limit to `UFix64`.
  ///
  /// # Errors
  /// * Numeric conversion
  pub fn limit(&self) -> Result<UFix64<N6>, CoreError> {
    Ok(self.limit.try_into()?)
  }

  /// Sets new withdrawal limit, resetting ledger and epoch.
  ///
  /// # Errors
  /// * Numeric conversion
  /// * New limit zero
  /// * Ledger epoch greater than current epoch
  pub fn update_limit(
    &mut self,
    new_limit_raw: UFixValue64,
    current_epoch: u64,
  ) -> Result<(), CoreError> {
    let new_limit: UFix64<N6> = new_limit_raw.try_into()?;
    if current_epoch < self.epoch {
      Err(WithdrawalLimitInvalidEpoch)
    } else if new_limit > UFix64::zero() {
      self.limit = new_limit_raw;
      self.withdrawal_ledger = VirtualStablecoin::new();
      self.epoch = current_epoch;
      Ok(())
    } else {
      Err(WithdrawalLimitValidation)
    }
  }

  /// Registers a withdrawal against the ledger, resetting it on epoch
  /// rollover.
  ///
  /// # Errors
  /// * Numeric conversion
  /// * Overflow while totaling withdrawals
  /// * Withdrawal exceeds limit for epoch
  /// * Ledger epoch greater than current epoch
  pub fn register_withdrawal(
    &mut self,
    withdrawal: UFix64<N6>,
    current_epoch: u64,
  ) -> Result<(), CoreError> {
    match current_epoch.cmp(&self.epoch) {
      Ordering::Less => Err(WithdrawalLimitInvalidEpoch),
      Ordering::Equal => self.register_within_limit(withdrawal),
      Ordering::Greater => {
        self.epoch = current_epoch;
        self.withdrawal_ledger = VirtualStablecoin::new();
        self.register_within_limit(withdrawal)
      }
    }
  }

  /// Adds withdrawal to ledger if projected total stays within limit.
  fn register_within_limit(
    &mut self,
    withdrawal: UFix64<N6>,
  ) -> Result<(), CoreError> {
    let projected = self
      .withdrawal_ledger
      .supply()?
      .checked_add(&withdrawal)
      .ok_or(WithdrawalLimitArithmetic)?;
    if projected <= self.limit()? {
      self.withdrawal_ledger.mint(withdrawal)
    } else {
      Err(WithdrawalLimitExceededForEpoch)
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use crate::error::CoreError::{
    WithdrawalLimitArithmetic, WithdrawalLimitExceededForEpoch,
    WithdrawalLimitInvalidEpoch, WithdrawalLimitValidation,
  };

  const EPOCH: u64 = 10;

  fn limiter() -> WithdrawalLimiter {
    WithdrawalLimiter::new(UFixValue64::new(1_000_000, -6), EPOCH)
  }

  #[test]
  fn update_limit_nonzero() -> Result<(), CoreError> {
    let mut limiter = limiter();
    let new_limit = UFixValue64::new(2_000_000, -6);
    limiter.update_limit(new_limit, EPOCH)?;
    assert_eq!(limiter.limit()?, UFix64::constant(2_000_000));
    Ok(())
  }

  #[test]
  fn reject_update_zero_limit() {
    let mut limiter = limiter();
    let result = limiter.update_limit(UFixValue64::new(0, -6), EPOCH);
    assert_eq!(result.err(), Some(WithdrawalLimitValidation));
  }

  #[test]
  fn update_limit_resets_ledger_and_epoch() -> Result<(), CoreError> {
    let mut limiter = limiter();
    limiter.register_withdrawal(UFix64::constant(1_000_000), EPOCH)?;
    limiter.update_limit(UFixValue64::new(1_000_000, -6), EPOCH)?;
    assert_eq!(limiter.epoch, EPOCH);
    assert_eq!(limiter.withdrawal_ledger.supply()?, UFix64::zero());
    limiter.register_withdrawal(UFix64::constant(1_000_000), EPOCH)?;
    Ok(())
  }

  #[test]
  fn accept_withdrawal_under_limit() -> Result<(), CoreError> {
    let mut limiter = limiter();
    let amount = UFix64::constant(400_000);
    limiter.register_withdrawal(amount, EPOCH)?;
    assert_eq!(limiter.withdrawal_ledger.supply()?, amount);
    Ok(())
  }

  #[test]
  fn accept_cumulative_at_limit() -> Result<(), CoreError> {
    let mut limiter = limiter();
    limiter.register_withdrawal(UFix64::constant(400_000), EPOCH)?;
    limiter.register_withdrawal(UFix64::constant(600_000), EPOCH)?;
    assert_eq!(limiter.withdrawal_ledger.supply()?, limiter.limit()?);
    Ok(())
  }

  #[test]
  fn reject_cumulative_over_limit() -> Result<(), CoreError> {
    let mut limiter = limiter();
    let amount = UFix64::constant(400_000);
    limiter.register_withdrawal(amount, EPOCH)?;
    let result = limiter.register_withdrawal(UFix64::constant(600_001), EPOCH);
    assert_eq!(result.err(), Some(WithdrawalLimitExceededForEpoch));
    assert_eq!(limiter.withdrawal_ledger.supply()?, amount);
    Ok(())
  }

  #[test]
  fn reject_on_ledger_overflow() -> Result<(), CoreError> {
    let mut limiter = limiter();
    limiter.register_withdrawal(UFix64::constant(1), EPOCH)?;
    let result = limiter.register_withdrawal(UFix64::constant(u64::MAX), EPOCH);
    assert_eq!(result.err(), Some(WithdrawalLimitArithmetic));
    Ok(())
  }

  #[test]
  fn rollover_resets_ledger() -> Result<(), CoreError> {
    let mut limiter = limiter();
    limiter.register_withdrawal(UFix64::constant(1_000_000), EPOCH)?;
    let amount = UFix64::constant(700_000);
    limiter.register_withdrawal(amount, EPOCH + 1)?;
    assert_eq!(limiter.epoch, EPOCH + 1);
    assert_eq!(limiter.withdrawal_ledger.supply()?, amount);
    Ok(())
  }

  #[test]
  fn reject_rollover_over_limit() -> Result<(), CoreError> {
    let mut limiter = limiter();
    limiter.register_withdrawal(UFix64::constant(500_000), EPOCH)?;
    let result =
      limiter.register_withdrawal(UFix64::constant(1_000_001), EPOCH + 1);
    assert_eq!(result.err(), Some(WithdrawalLimitExceededForEpoch));
    Ok(())
  }

  #[test]
  fn reject_update_stale_epoch() {
    let mut limiter = limiter();
    let result =
      limiter.update_limit(UFixValue64::new(2_000_000, -6), EPOCH - 1);
    assert_eq!(result.err(), Some(WithdrawalLimitInvalidEpoch));
  }

  #[test]
  fn reject_stale_epoch() {
    let mut limiter = limiter();
    let result = limiter.register_withdrawal(UFix64::constant(1), EPOCH - 1);
    assert_eq!(result.err(), Some(WithdrawalLimitInvalidEpoch));
  }
}