Skip to main content

balancer_maths_rust/vault/
mod.rs

1//! Vault operations for Balancer pools
2
3pub mod add_liquidity;
4pub mod base_pool_math;
5pub mod remove_liquidity;
6pub mod swap;
7
8use crate::common::errors::PoolError;
9use crate::common::pool_base::PoolBase;
10use crate::common::types::*;
11use crate::hooks::types::HookState;
12use crate::hooks::{
13    AkronHook, DefaultHook, DirectionalFeeHook, ExitFeeHook, HookBase, StableSurgeHook,
14};
15use crate::pools::buffer::erc4626_buffer_wrap_or_unwrap;
16use crate::vault::add_liquidity::add_liquidity;
17use crate::vault::remove_liquidity::remove_liquidity;
18use crate::vault::swap::swap;
19use alloy_primitives::U256;
20
21/// Main vault interface for pool operations
22pub struct Vault;
23
24impl Vault {
25    /// Create a new vault instance
26    pub fn new() -> Self {
27        Vault
28    }
29
30    /// Get hook instance based on hook type
31    fn get_hook(
32        &self,
33        hook_type: &Option<String>,
34        hook_state: Option<&HookState>,
35    ) -> Box<dyn HookBase> {
36        match hook_type {
37            Some(hook_type) => match hook_type.as_str() {
38                "Akron" => {
39                    if let Some(HookState::Akron(_)) = hook_state {
40                        Box::new(AkronHook::new())
41                    } else {
42                        Box::new(DefaultHook::new())
43                    }
44                }
45                "DirectionalFee" => {
46                    if let Some(HookState::DirectionalFee(_)) = hook_state {
47                        Box::new(DirectionalFeeHook::new())
48                    } else {
49                        Box::new(DefaultHook::new())
50                    }
51                }
52                "StableSurge" => {
53                    if let Some(HookState::StableSurge(_)) = hook_state {
54                        Box::new(StableSurgeHook::new())
55                    } else {
56                        Box::new(DefaultHook::new())
57                    }
58                }
59                "ExitFee" => {
60                    if let Some(HookState::ExitFee(_)) = hook_state {
61                        Box::new(ExitFeeHook::new())
62                    } else {
63                        Box::new(DefaultHook::new())
64                    }
65                }
66                _ => Box::new(DefaultHook::new()),
67            },
68            None => Box::new(DefaultHook::new()),
69        }
70    }
71
72    /// Perform a swap operation
73    pub fn swap(
74        &self,
75        swap_input: &SwapInput,
76        pool_state_or_buffer: &PoolStateOrBuffer,
77        hook_state: Option<&HookState>,
78    ) -> Result<U256, PoolError> {
79        match pool_state_or_buffer {
80            PoolStateOrBuffer::Pool(pool_state) => {
81                let base_state = pool_state.base();
82
83                // Create pool instance
84                let pool: Box<dyn PoolBase> = match pool_state.as_ref() {
85                    PoolState::Weighted(weighted_state) => {
86                        // Use the weights from the WeightedState directly
87                        Box::new(crate::pools::weighted::WeightedPool::from(
88                            weighted_state.clone(),
89                        ))
90                    }
91                    PoolState::Stable(stable_state) => Box::new(
92                        crate::pools::stable::StablePool::new(stable_state.mutable.clone()),
93                    ),
94                    PoolState::GyroECLP(gyro_eclp_state) => Box::new(
95                        crate::pools::gyro::GyroECLPPool::new(gyro_eclp_state.immutable.clone()),
96                    ),
97                    PoolState::QuantAmm(quant_amm_state) => Box::new(
98                        crate::pools::quantamm::QuantAmmPool::from(quant_amm_state.clone()),
99                    ),
100                    PoolState::LiquidityBootstrapping(liquidity_bootstrapping_state) => Box::new(
101                        crate::pools::liquidity_bootstrapping::LiquidityBootstrappingPool::from(
102                            liquidity_bootstrapping_state.clone(),
103                        ),
104                    ),
105                    PoolState::FixedPriceLBP(fixed_price_lbp_state) => {
106                        Box::new(crate::pools::fixed_price_lbp::FixedPriceLBPPool::from(
107                            fixed_price_lbp_state.clone(),
108                        ))
109                    }
110                    PoolState::ReClamm(re_clamm_state) => Box::new(
111                        crate::pools::reclamm::ReClammPool::new(re_clamm_state.clone()),
112                    ),
113                    PoolState::ReClammV2(re_clamm_v2_state) => Box::new(
114                        crate::pools::reclammv2::ReClammV2Pool::new(re_clamm_v2_state.clone()),
115                    ),
116                    _ => return Err(PoolError::UnsupportedPoolType(base_state.pool_type.clone())),
117                };
118
119                // Get hook instance
120                let hook: Box<dyn HookBase> = self.get_hook(&base_state.hook_type, hook_state);
121
122                // Execute swap
123                swap(
124                    swap_input,
125                    pool_state,
126                    pool.as_ref(),
127                    hook.as_ref(),
128                    hook_state,
129                )
130            }
131            PoolStateOrBuffer::Buffer(buffer_state) => Ok(erc4626_buffer_wrap_or_unwrap(
132                swap_input,
133                buffer_state.as_ref(),
134            )?),
135        }
136    }
137
138    /// Add liquidity to a pool
139    pub fn add_liquidity(
140        &self,
141        add_liquidity_input: &AddLiquidityInput,
142        pool_state: &PoolState,
143        hook_state: Option<&HookState>,
144    ) -> Result<AddLiquidityResult, PoolError> {
145        let base_state = pool_state.base();
146
147        // Create pool instance
148        let pool: Box<dyn PoolBase> = match pool_state {
149            PoolState::Weighted(weighted_state) => {
150                // Use the weights from the WeightedState directly
151                Box::new(crate::pools::weighted::WeightedPool::from(
152                    weighted_state.clone(),
153                ))
154            }
155            PoolState::Stable(stable_state) => Box::new(crate::pools::stable::StablePool::new(
156                stable_state.mutable.clone(),
157            )),
158            PoolState::GyroECLP(gyro_eclp_state) => Box::new(
159                crate::pools::gyro::GyroECLPPool::new(gyro_eclp_state.immutable.clone()),
160            ),
161            PoolState::QuantAmm(quant_amm_state) => Box::new(
162                crate::pools::quantamm::QuantAmmPool::from(quant_amm_state.clone()),
163            ),
164            PoolState::LiquidityBootstrapping(liquidity_bootstrapping_state) => Box::new(
165                crate::pools::liquidity_bootstrapping::LiquidityBootstrappingPool::from(
166                    liquidity_bootstrapping_state.clone(),
167                ),
168            ),
169            PoolState::FixedPriceLBP(fixed_price_lbp_state) => {
170                Box::new(crate::pools::fixed_price_lbp::FixedPriceLBPPool::from(
171                    fixed_price_lbp_state.clone(),
172                ))
173            }
174            PoolState::ReClamm(re_clamm_state) => Box::new(
175                crate::pools::reclamm::ReClammPool::new(re_clamm_state.clone()),
176            ),
177            PoolState::ReClammV2(re_clamm_v2_state) => Box::new(
178                crate::pools::reclammv2::ReClammV2Pool::new(re_clamm_v2_state.clone()),
179            ),
180            _ => return Err(PoolError::UnsupportedPoolType(base_state.pool_type.clone())),
181        };
182
183        // Create hook instance
184        let hook: Box<dyn HookBase> = self.get_hook(&base_state.hook_type, hook_state);
185
186        add_liquidity(
187            add_liquidity_input,
188            pool_state,
189            pool.as_ref(),
190            hook.as_ref(),
191            hook_state,
192        )
193    }
194
195    /// Remove liquidity from a pool
196    pub fn remove_liquidity(
197        &self,
198        remove_liquidity_input: &RemoveLiquidityInput,
199        pool_state: &PoolState,
200        hook_state: Option<&HookState>,
201    ) -> Result<RemoveLiquidityResult, PoolError> {
202        let base_state = pool_state.base();
203
204        // Create pool instance
205        let pool: Box<dyn PoolBase> = match pool_state {
206            PoolState::Weighted(weighted_state) => {
207                // Use the weights from the WeightedState directly
208                Box::new(crate::pools::weighted::WeightedPool::from(
209                    weighted_state.clone(),
210                ))
211            }
212            PoolState::Stable(stable_state) => Box::new(crate::pools::stable::StablePool::new(
213                stable_state.mutable.clone(),
214            )),
215            PoolState::GyroECLP(gyro_eclp_state) => Box::new(
216                crate::pools::gyro::GyroECLPPool::new(gyro_eclp_state.immutable.clone()),
217            ),
218            PoolState::QuantAmm(quant_amm_state) => Box::new(
219                crate::pools::quantamm::QuantAmmPool::from(quant_amm_state.clone()),
220            ),
221            PoolState::LiquidityBootstrapping(liquidity_bootstrapping_state) => Box::new(
222                crate::pools::liquidity_bootstrapping::LiquidityBootstrappingPool::from(
223                    liquidity_bootstrapping_state.clone(),
224                ),
225            ),
226            PoolState::FixedPriceLBP(fixed_price_lbp_state) => {
227                Box::new(crate::pools::fixed_price_lbp::FixedPriceLBPPool::from(
228                    fixed_price_lbp_state.clone(),
229                ))
230            }
231            PoolState::ReClamm(re_clamm_state) => Box::new(
232                crate::pools::reclamm::ReClammPool::new(re_clamm_state.clone()),
233            ),
234            PoolState::ReClammV2(re_clamm_v2_state) => Box::new(
235                crate::pools::reclammv2::ReClammV2Pool::new(re_clamm_v2_state.clone()),
236            ),
237            _ => return Err(PoolError::UnsupportedPoolType(base_state.pool_type.clone())),
238        };
239
240        // Create hook instance
241        let hook: Box<dyn HookBase> = self.get_hook(&base_state.hook_type, hook_state);
242
243        remove_liquidity(
244            remove_liquidity_input,
245            pool_state,
246            pool.as_ref(),
247            hook.as_ref(),
248            hook_state,
249        )
250    }
251}
252
253impl Default for Vault {
254    fn default() -> Self {
255        Vault::new()
256    }
257}