Skip to main content

balancer_maths_rust/pools/quantamm/
quantamm_pool.rs

1//! QuantAmm pool implementation
2
3use crate::common::errors::PoolError;
4use crate::common::maths::mul_down_fixed;
5use crate::common::pool_base::PoolBase;
6use crate::common::types::{Rounding, SwapParams};
7use crate::pools::quantamm::quantamm_data::QuantAmmState;
8use crate::pools::quantamm::quantamm_math::{
9    calculate_block_normalised_weight, get_first_four_weights_and_multipliers,
10    get_second_four_weights_and_multipliers,
11};
12use crate::pools::weighted::weighted_math::{MAX_INVARIANT_RATIO, MIN_INVARIANT_RATIO, *};
13use alloy_primitives::{I256, U256};
14
15/// QuantAmm pool implementation
16pub struct QuantAmmPool {
17    /// Current normalized weights (scaled 18) based on time interpolation
18    normalized_weights: Vec<U256>,
19    /// Pool state
20    state: QuantAmmState,
21}
22
23impl QuantAmmPool {
24    /// Create a new QuantAmm pool
25    pub fn new(state: QuantAmmState) -> Result<Self, PoolError> {
26        let (first_weights, first_multipliers) = get_first_four_weights_and_multipliers(
27            &state.base.tokens,
28            &state.mutable.first_four_weights_and_multipliers,
29        );
30
31        let (second_weights, second_multipliers) = get_second_four_weights_and_multipliers(
32            &state.base.tokens,
33            &state.mutable.second_four_weights_and_multipliers,
34        );
35
36        let base_weights = [first_weights, second_weights].concat();
37        let multipliers = [first_multipliers, second_multipliers].concat();
38
39        if base_weights.is_empty() {
40            return Err(PoolError::InvalidSwapParameters);
41        }
42
43        // Calculate current normalized weights based on time interpolation
44        let normalized_weights = Self::calculate_normalized_weights(
45            &base_weights,
46            &multipliers,
47            &state.mutable.last_update_time,
48            &state.mutable.last_interop_time,
49            &state.mutable.current_timestamp,
50        )?;
51
52        Ok(Self {
53            normalized_weights,
54            state,
55        })
56    }
57
58    /// Calculate normalized weights based on time interpolation
59    fn calculate_normalized_weights(
60        base_weights: &[I256],
61        multipliers: &[I256],
62        last_update_time: &U256,
63        last_interop_time: &U256,
64        current_timestamp: &U256,
65    ) -> Result<Vec<U256>, PoolError> {
66        let mut multiplier_time = *current_timestamp;
67
68        if current_timestamp >= last_interop_time {
69            multiplier_time = *last_interop_time;
70        }
71
72        let time_since_last_update = multiplier_time
73            .checked_sub(*last_update_time)
74            .ok_or(PoolError::TimestampBeforeLastUpdate)?;
75
76        let mut normalized_weights = Vec::with_capacity(base_weights.len());
77
78        for i in 0..base_weights.len() {
79            let normalized_weight = calculate_block_normalised_weight(
80                &base_weights[i],
81                &multipliers[i],
82                &time_since_last_update,
83            )?;
84            normalized_weights.push(normalized_weight);
85        }
86
87        Ok(normalized_weights)
88    }
89
90    /// Get normalized weights for a specific token pair
91    fn get_normalized_weight_pair(
92        &self,
93        index_in: usize,
94        index_out: usize,
95    ) -> Result<(U256, U256), PoolError> {
96        if index_in >= self.normalized_weights.len() || index_out >= self.normalized_weights.len() {
97            return Err(PoolError::InvalidTokenIndex);
98        }
99
100        let token_in_weight = self.normalized_weights[index_in];
101        let token_out_weight = self.normalized_weights[index_out];
102
103        Ok((token_in_weight, token_out_weight))
104    }
105
106    /// Check if trade size exceeds max trade size ratio
107    fn check_max_trade_size(
108        &self,
109        amount_scaled_18: &U256,
110        balance_scaled_18: &U256,
111    ) -> Result<(), PoolError> {
112        let max_amount = mul_down_fixed(
113            balance_scaled_18,
114            &self.state.immutable.max_trade_size_ratio,
115        )?;
116
117        if amount_scaled_18 > &max_amount {
118            return Err(PoolError::InvalidSwapParameters);
119        }
120
121        Ok(())
122    }
123}
124
125impl PoolBase for QuantAmmPool {
126    fn on_swap(&self, swap_params: &SwapParams) -> Result<U256, PoolError> {
127        let token_in_index = swap_params.token_in_index;
128        let token_out_index = swap_params.token_out_index;
129
130        if token_in_index >= self.normalized_weights.len()
131            || token_out_index >= self.normalized_weights.len()
132        {
133            return Err(PoolError::InvalidTokenIndex);
134        }
135
136        let balance_in = &swap_params.balances_live_scaled_18[token_in_index];
137        let balance_out = &swap_params.balances_live_scaled_18[token_out_index];
138        let amount_scaled_18 = &swap_params.amount_scaled_18;
139
140        let (weight_in, weight_out) =
141            self.get_normalized_weight_pair(token_in_index, token_out_index)?;
142
143        match swap_params.swap_kind {
144            crate::common::types::SwapKind::GivenIn => {
145                // Check max trade size ratio for input
146                self.check_max_trade_size(amount_scaled_18, balance_in)?;
147
148                let amount_out_scaled_18 = compute_out_given_exact_in(
149                    balance_in,
150                    &weight_in,
151                    balance_out,
152                    &weight_out,
153                    amount_scaled_18,
154                )?;
155
156                // Check max trade size ratio for output
157                self.check_max_trade_size(&amount_out_scaled_18, balance_out)?;
158
159                Ok(amount_out_scaled_18)
160            }
161            crate::common::types::SwapKind::GivenOut => {
162                // Check max trade size ratio for output
163                self.check_max_trade_size(amount_scaled_18, balance_out)?;
164
165                let amount_in_scaled_18 = compute_in_given_exact_out(
166                    balance_in,
167                    &weight_in,
168                    balance_out,
169                    &weight_out,
170                    amount_scaled_18,
171                )?;
172
173                // Check max trade size ratio for input
174                self.check_max_trade_size(&amount_in_scaled_18, balance_in)?;
175
176                Ok(amount_in_scaled_18)
177            }
178        }
179    }
180
181    fn compute_invariant(
182        &self,
183        balances_live_scaled_18: &[U256],
184        rounding: Rounding,
185    ) -> Result<U256, PoolError> {
186        match rounding {
187            Rounding::RoundDown => {
188                compute_invariant_down(&self.normalized_weights, balances_live_scaled_18)
189            }
190            Rounding::RoundUp => {
191                compute_invariant_up(&self.normalized_weights, balances_live_scaled_18)
192            }
193        }
194    }
195
196    fn compute_balance(
197        &self,
198        balances_live_scaled_18: &[U256],
199        token_in_index: usize,
200        invariant_ratio: &U256,
201    ) -> Result<U256, PoolError> {
202        if token_in_index >= balances_live_scaled_18.len()
203            || token_in_index >= self.normalized_weights.len()
204        {
205            return Err(PoolError::InvalidTokenIndex);
206        }
207
208        let current_balance = &balances_live_scaled_18[token_in_index];
209        let weight = &self.normalized_weights[token_in_index];
210
211        // Calculate the new balance based on the invariant ratio
212        compute_balance_out_given_invariant(current_balance, weight, invariant_ratio)
213    }
214
215    fn get_maximum_invariant_ratio(&self) -> U256 {
216        MAX_INVARIANT_RATIO
217    }
218
219    fn get_minimum_invariant_ratio(&self) -> U256 {
220        MIN_INVARIANT_RATIO
221    }
222}
223
224impl From<QuantAmmState> for QuantAmmPool {
225    fn from(quant_amm_state: QuantAmmState) -> Self {
226        Self::new(quant_amm_state).expect("Failed to create QuantAmmPool from state")
227    }
228}