use alloy_primitives::{Address, U256};
use serde::{Deserialize, Serialize};
use crate::defi::tick_map::full_math::{FullMath, Q128};
#[cfg_attr(
feature = "python",
pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.model", from_py_object)
)]
#[cfg_attr(
feature = "python",
pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.model")
)]
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PoolPosition {
pub owner: Address,
pub tick_lower: i32,
pub tick_upper: i32,
pub liquidity: u128,
pub fee_growth_inside_0_last: U256,
pub fee_growth_inside_1_last: U256,
pub tokens_owed_0: u128,
pub tokens_owed_1: u128,
pub total_amount0_deposited: U256,
pub total_amount1_deposited: U256,
pub total_amount0_collected: u128,
pub total_amount1_collected: u128,
}
impl PoolPosition {
#[must_use]
pub fn new(owner: Address, tick_lower: i32, tick_upper: i32, liquidity: i128) -> Self {
Self {
owner,
tick_lower,
tick_upper,
liquidity: liquidity.unsigned_abs(),
fee_growth_inside_0_last: U256::ZERO,
fee_growth_inside_1_last: U256::ZERO,
tokens_owed_0: 0,
tokens_owed_1: 0,
total_amount0_deposited: U256::ZERO,
total_amount1_deposited: U256::ZERO,
total_amount0_collected: 0,
total_amount1_collected: 0,
}
}
#[must_use]
pub fn get_position_key(owner: &Address, tick_lower: i32, tick_upper: i32) -> String {
format!("{owner}:{tick_lower}:{tick_upper}")
}
pub fn update_liquidity(&mut self, liquidity_delta: i128) {
if liquidity_delta < 0 {
self.liquidity = self.liquidity.saturating_sub((-liquidity_delta) as u128);
} else {
self.liquidity = self.liquidity.saturating_add(liquidity_delta as u128);
}
}
pub fn update_fees(&mut self, fee_growth_inside_0: U256, fee_growth_inside_1: U256) {
if self.liquidity > 0 {
let fee_delta_0 = fee_growth_inside_0.saturating_sub(self.fee_growth_inside_0_last);
let fee_delta_1 = fee_growth_inside_1.saturating_sub(self.fee_growth_inside_1_last);
let tokens_owed_0_full =
FullMath::mul_div(fee_delta_0, U256::from(self.liquidity), Q128)
.unwrap_or(U256::ZERO);
let tokens_owed_1_full =
FullMath::mul_div(fee_delta_1, U256::from(self.liquidity), Q128)
.unwrap_or(U256::ZERO);
self.tokens_owed_0 = self
.tokens_owed_0
.wrapping_add(FullMath::truncate_to_u128(tokens_owed_0_full));
self.tokens_owed_1 = self
.tokens_owed_1
.wrapping_add(FullMath::truncate_to_u128(tokens_owed_1_full));
}
self.fee_growth_inside_0_last = fee_growth_inside_0;
self.fee_growth_inside_1_last = fee_growth_inside_1;
}
pub fn collect_fees(&mut self, amount0: u128, amount1: u128) {
let collect_amount_0 = amount0.min(self.tokens_owed_0);
let collect_amount_1 = amount1.min(self.tokens_owed_1);
self.tokens_owed_0 -= collect_amount_0;
self.tokens_owed_1 -= collect_amount_1;
self.total_amount0_collected += collect_amount_0;
self.total_amount1_collected += collect_amount_1;
}
pub fn update_amounts(&mut self, liquidity_delta: i128, amount0: U256, amount1: U256) {
if liquidity_delta > 0 {
self.total_amount0_deposited += amount0;
self.total_amount1_deposited += amount1;
} else {
self.tokens_owed_0 = self
.tokens_owed_0
.wrapping_add(FullMath::truncate_to_u128(amount0));
self.tokens_owed_1 = self
.tokens_owed_1
.wrapping_add(FullMath::truncate_to_u128(amount1));
}
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.liquidity == 0 && self.tokens_owed_0 == 0 && self.tokens_owed_1 == 0
}
}
#[cfg(test)]
mod tests {
use alloy_primitives::address;
use rstest::rstest;
use super::*;
#[rstest]
fn test_new_position() {
let owner = address!("1234567890123456789012345678901234567890");
let tick_lower = -100;
let tick_upper = 100;
let liquidity = 1000i128;
let position = PoolPosition::new(owner, tick_lower, tick_upper, liquidity);
assert_eq!(position.owner, owner);
assert_eq!(position.tick_lower, tick_lower);
assert_eq!(position.tick_upper, tick_upper);
assert_eq!(position.liquidity, liquidity as u128);
assert_eq!(position.fee_growth_inside_0_last, U256::ZERO);
assert_eq!(position.fee_growth_inside_1_last, U256::ZERO);
assert_eq!(position.tokens_owed_0, 0);
assert_eq!(position.tokens_owed_1, 0);
}
#[rstest]
fn test_get_position_key() {
let owner = address!("1234567890123456789012345678901234567890");
let tick_lower = -100;
let tick_upper = 100;
let key = PoolPosition::get_position_key(&owner, tick_lower, tick_upper);
let expected = format!("{owner:?}:{tick_lower}:{tick_upper}");
assert_eq!(key, expected);
}
#[rstest]
fn test_update_liquidity_positive() {
let owner = address!("1234567890123456789012345678901234567890");
let mut position = PoolPosition::new(owner, -100, 100, 1000);
position.update_liquidity(500);
assert_eq!(position.liquidity, 1500);
}
#[rstest]
fn test_update_liquidity_negative() {
let owner = address!("1234567890123456789012345678901234567890");
let mut position = PoolPosition::new(owner, -100, 100, 1000);
position.update_liquidity(-300);
assert_eq!(position.liquidity, 700);
}
#[rstest]
fn test_update_liquidity_negative_saturating() {
let owner = address!("1234567890123456789012345678901234567890");
let mut position = PoolPosition::new(owner, -100, 100, 1000);
position.update_liquidity(-2000); assert_eq!(position.liquidity, 0);
}
#[rstest]
fn test_update_fees() {
let owner = address!("1234567890123456789012345678901234567890");
let mut position = PoolPosition::new(owner, -100, 100, 1000);
let fee_growth_inside_0 = U256::from(100);
let fee_growth_inside_1 = U256::from(200);
position.update_fees(fee_growth_inside_0, fee_growth_inside_1);
assert_eq!(position.fee_growth_inside_0_last, fee_growth_inside_0);
assert_eq!(position.fee_growth_inside_1_last, fee_growth_inside_1);
}
#[rstest]
fn test_collect_fees() {
let owner = address!("1234567890123456789012345678901234567890");
let mut position = PoolPosition::new(owner, -100, 100, 1000);
position.tokens_owed_0 = 100;
position.tokens_owed_1 = 200;
position.collect_fees(50, 150);
assert_eq!(position.total_amount0_collected, 50);
assert_eq!(position.total_amount1_collected, 150);
assert_eq!(position.tokens_owed_0, 50);
assert_eq!(position.tokens_owed_1, 50);
}
#[rstest]
fn test_collect_fees_more_than_owed() {
let owner = address!("1234567890123456789012345678901234567890");
let mut position = PoolPosition::new(owner, -100, 100, 1000);
position.tokens_owed_0 = 100;
position.tokens_owed_1 = 200;
position.collect_fees(150, 300);
assert_eq!(position.total_amount0_collected, 100); assert_eq!(position.total_amount1_collected, 200);
assert_eq!(position.tokens_owed_0, 0);
assert_eq!(position.tokens_owed_1, 0);
}
#[rstest]
fn test_is_empty() {
let owner = address!("1234567890123456789012345678901234567890");
let mut position = PoolPosition::new(owner, -100, 100, 0);
assert!(position.is_empty());
position.liquidity = 100;
assert!(!position.is_empty());
position.liquidity = 0;
position.tokens_owed_0 = 50;
assert!(!position.is_empty());
position.tokens_owed_0 = 0;
position.tokens_owed_1 = 25;
assert!(!position.is_empty());
position.tokens_owed_1 = 0;
assert!(position.is_empty());
}
}