Skip to main content

balancer_maths_rust/pools/weighted/
weighted_pool.rs

1//! Weighted pool implementation
2
3use crate::common::errors::PoolError;
4use crate::common::pool_base::PoolBase;
5use crate::common::types::{Rounding, SwapParams};
6use crate::pools::weighted::weighted_data::WeightedState;
7use crate::pools::weighted::weighted_math::{MAX_INVARIANT_RATIO, MIN_INVARIANT_RATIO, *};
8use alloy_primitives::U256;
9
10/// Weighted pool implementation
11pub struct WeightedPool {
12    /// Normalized weights (scaled 18)
13    normalized_weights: Vec<U256>,
14}
15
16impl WeightedPool {
17    /// Create a new weighted pool
18    pub fn new(weights: Vec<U256>) -> Result<Self, PoolError> {
19        if weights.is_empty() {
20            return Err(PoolError::InvalidSwapParameters);
21        }
22        Ok(Self {
23            normalized_weights: weights,
24        })
25    }
26
27    /// Get the normalized weights
28    pub fn normalized_weights(&self) -> &[U256] {
29        &self.normalized_weights
30    }
31}
32
33impl PoolBase for WeightedPool {
34    fn on_swap(&self, swap_params: &SwapParams) -> Result<U256, PoolError> {
35        let token_in_index = swap_params.token_in_index;
36        let token_out_index = swap_params.token_out_index;
37
38        if token_in_index >= self.normalized_weights.len()
39            || token_out_index >= self.normalized_weights.len()
40        {
41            return Err(PoolError::InvalidTokenIndex);
42        }
43
44        let balance_in = &swap_params.balances_live_scaled_18[token_in_index];
45        let weight_in = &self.normalized_weights[token_in_index];
46        let balance_out = &swap_params.balances_live_scaled_18[token_out_index];
47        let weight_out = &self.normalized_weights[token_out_index];
48        let amount_scaled_18 = &swap_params.amount_scaled_18;
49
50        match swap_params.swap_kind {
51            crate::common::types::SwapKind::GivenIn => compute_out_given_exact_in(
52                balance_in,
53                weight_in,
54                balance_out,
55                weight_out,
56                amount_scaled_18,
57            ),
58            crate::common::types::SwapKind::GivenOut => compute_in_given_exact_out(
59                balance_in,
60                weight_in,
61                balance_out,
62                weight_out,
63                amount_scaled_18,
64            ),
65        }
66    }
67
68    fn compute_invariant(
69        &self,
70        balances_live_scaled_18: &[U256],
71        rounding: Rounding,
72    ) -> Result<U256, PoolError> {
73        match rounding {
74            Rounding::RoundDown => {
75                compute_invariant_down(&self.normalized_weights, balances_live_scaled_18)
76            }
77            Rounding::RoundUp => {
78                compute_invariant_up(&self.normalized_weights, balances_live_scaled_18)
79            }
80        }
81    }
82
83    fn compute_balance(
84        &self,
85        balances_live_scaled_18: &[U256],
86        token_in_index: usize,
87        invariant_ratio: &U256,
88    ) -> Result<U256, PoolError> {
89        if token_in_index >= balances_live_scaled_18.len()
90            || token_in_index >= self.normalized_weights.len()
91        {
92            return Err(PoolError::InvalidTokenIndex);
93        }
94
95        let current_balance = &balances_live_scaled_18[token_in_index];
96        let weight = &self.normalized_weights[token_in_index];
97
98        // Calculate the new balance based on the invariant ratio
99        compute_balance_out_given_invariant(current_balance, weight, invariant_ratio)
100    }
101
102    fn get_maximum_invariant_ratio(&self) -> U256 {
103        MAX_INVARIANT_RATIO
104    }
105
106    fn get_minimum_invariant_ratio(&self) -> U256 {
107        MIN_INVARIANT_RATIO
108    }
109}
110
111impl From<WeightedState> for WeightedPool {
112    fn from(weighted_state: WeightedState) -> Self {
113        Self {
114            normalized_weights: weighted_state.weights,
115        }
116    }
117}