balancer_maths_rust/pools/quantamm/
quantamm_pool.rs1use 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
15pub struct QuantAmmPool {
17 normalized_weights: Vec<U256>,
19 state: QuantAmmState,
21}
22
23impl QuantAmmPool {
24 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 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 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 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 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 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 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 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 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 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}