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 num_bigint::BigInt;
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<BigInt, 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::ReClamm(re_clamm_state) => Box::new(
106                        crate::pools::reclamm::ReClammPool::new(re_clamm_state.clone()),
107                    ),
108                    PoolState::ReClammV2(re_clamm_v2_state) => Box::new(
109                        crate::pools::reclammv2::ReClammV2Pool::new(re_clamm_v2_state.clone()),
110                    ),
111                    _ => return Err(PoolError::UnsupportedPoolType(base_state.pool_type.clone())),
112                };
113
114                // Get hook instance
115                let hook: Box<dyn HookBase> = self.get_hook(&base_state.hook_type, hook_state);
116
117                // Execute swap
118                swap(
119                    swap_input,
120                    pool_state,
121                    pool.as_ref(),
122                    hook.as_ref(),
123                    hook_state,
124                )
125            }
126            PoolStateOrBuffer::Buffer(buffer_state) => Ok(erc4626_buffer_wrap_or_unwrap(
127                swap_input,
128                buffer_state.as_ref(),
129            )?),
130        }
131    }
132
133    /// Add liquidity to a pool
134    pub fn add_liquidity(
135        &self,
136        add_liquidity_input: &AddLiquidityInput,
137        pool_state: &PoolState,
138        hook_state: Option<&HookState>,
139    ) -> Result<AddLiquidityResult, PoolError> {
140        let base_state = pool_state.base();
141
142        // Create pool instance
143        let pool: Box<dyn PoolBase> = match pool_state {
144            PoolState::Weighted(weighted_state) => {
145                // Use the weights from the WeightedState directly
146                Box::new(crate::pools::weighted::WeightedPool::from(
147                    weighted_state.clone(),
148                ))
149            }
150            PoolState::Stable(stable_state) => Box::new(crate::pools::stable::StablePool::new(
151                stable_state.mutable.clone(),
152            )),
153            PoolState::GyroECLP(gyro_eclp_state) => Box::new(
154                crate::pools::gyro::GyroECLPPool::new(gyro_eclp_state.immutable.clone()),
155            ),
156            PoolState::QuantAmm(quant_amm_state) => Box::new(
157                crate::pools::quantamm::QuantAmmPool::from(quant_amm_state.clone()),
158            ),
159            PoolState::LiquidityBootstrapping(liquidity_bootstrapping_state) => Box::new(
160                crate::pools::liquidity_bootstrapping::LiquidityBootstrappingPool::from(
161                    liquidity_bootstrapping_state.clone(),
162                ),
163            ),
164            PoolState::ReClamm(re_clamm_state) => Box::new(
165                crate::pools::reclamm::ReClammPool::new(re_clamm_state.clone()),
166            ),
167            PoolState::ReClammV2(re_clamm_v2_state) => Box::new(
168                crate::pools::reclammv2::ReClammV2Pool::new(re_clamm_v2_state.clone()),
169            ),
170            _ => return Err(PoolError::UnsupportedPoolType(base_state.pool_type.clone())),
171        };
172
173        // Create hook instance
174        let hook: Box<dyn HookBase> = self.get_hook(&base_state.hook_type, hook_state);
175
176        add_liquidity(
177            add_liquidity_input,
178            pool_state,
179            pool.as_ref(),
180            hook.as_ref(),
181            hook_state,
182        )
183    }
184
185    /// Remove liquidity from a pool
186    pub fn remove_liquidity(
187        &self,
188        remove_liquidity_input: &RemoveLiquidityInput,
189        pool_state: &PoolState,
190        hook_state: Option<&HookState>,
191    ) -> Result<RemoveLiquidityResult, PoolError> {
192        let base_state = pool_state.base();
193
194        // Create pool instance
195        let pool: Box<dyn PoolBase> = match pool_state {
196            PoolState::Weighted(weighted_state) => {
197                // Use the weights from the WeightedState directly
198                Box::new(crate::pools::weighted::WeightedPool::from(
199                    weighted_state.clone(),
200                ))
201            }
202            PoolState::Stable(stable_state) => Box::new(crate::pools::stable::StablePool::new(
203                stable_state.mutable.clone(),
204            )),
205            PoolState::GyroECLP(gyro_eclp_state) => Box::new(
206                crate::pools::gyro::GyroECLPPool::new(gyro_eclp_state.immutable.clone()),
207            ),
208            PoolState::QuantAmm(quant_amm_state) => Box::new(
209                crate::pools::quantamm::QuantAmmPool::from(quant_amm_state.clone()),
210            ),
211            PoolState::LiquidityBootstrapping(liquidity_bootstrapping_state) => Box::new(
212                crate::pools::liquidity_bootstrapping::LiquidityBootstrappingPool::from(
213                    liquidity_bootstrapping_state.clone(),
214                ),
215            ),
216            PoolState::ReClamm(re_clamm_state) => Box::new(
217                crate::pools::reclamm::ReClammPool::new(re_clamm_state.clone()),
218            ),
219            PoolState::ReClammV2(re_clamm_v2_state) => Box::new(
220                crate::pools::reclammv2::ReClammV2Pool::new(re_clamm_v2_state.clone()),
221            ),
222            _ => return Err(PoolError::UnsupportedPoolType(base_state.pool_type.clone())),
223        };
224
225        // Create hook instance
226        let hook: Box<dyn HookBase> = self.get_hook(&base_state.hook_type, hook_state);
227
228        remove_liquidity(
229            remove_liquidity_input,
230            pool_state,
231            pool.as_ref(),
232            hook.as_ref(),
233            hook_state,
234        )
235    }
236}
237
238impl Default for Vault {
239    fn default() -> Self {
240        Vault::new()
241    }
242}