consta_pool/
lib.rs

1use std::cmp::Ordering;
2
3type Result<T> = std::result::Result<T, PoolError>;
4
5#[derive(Debug, Clone)]
6pub struct LiquidityPool {
7    initial_token_reserve: u64,
8    native_reserve: u64,
9    token_reserve: u64,
10    constant_product: u128,
11}
12
13#[derive(Debug, thiserror::Error, PartialEq)]
14pub enum PoolError {
15    #[error("Slippage too high")]
16    SlippageExceeded,
17    #[error("Invalid funds in the pool")]
18    InsufficientPoolFunds,
19    #[error("Invalid amount")]
20    InvalidAmount,
21    #[error("Overflow")]
22    Overflow,
23}
24
25impl LiquidityPool {
26    pub fn new(native_reserve: u64, token_reserve: u64) -> Result<Self> {
27        if native_reserve == 0 || token_reserve == 0 {
28            return Err(PoolError::InvalidAmount);
29        }
30        let constant_product = native_reserve as u128 * token_reserve as u128;
31        Ok(Self {
32            initial_token_reserve: token_reserve,
33            native_reserve,
34            token_reserve,
35            constant_product,
36        })
37    }
38
39    pub fn get_native_reserve(&self) -> u64 {
40        self.native_reserve
41    }
42    pub fn get_token_reserve(&self) -> u64 {
43        self.token_reserve
44    }
45
46    pub fn get_constant_product(&self) -> u128 {
47        self.constant_product
48    }
49
50    /// Returns the current market price of tokens in terms of native currency.
51    pub fn market_price(&self) -> f64 {
52        self.native_reserve as f64 / self.initial_token_reserve as f64
53    }
54
55    /// Buys `token_amount` tokens from the pool, checking if the native currency spent does not exceed `max_native`.
56    pub fn buy(&mut self, token_amount: u64, max_native: Option<u64>) -> Result<u64> {
57        if token_amount == 0 {
58            return Err(PoolError::InvalidAmount);
59        }
60        let new_token_reserve = self
61            .token_reserve
62            .checked_sub(token_amount)
63            .ok_or(PoolError::InsufficientPoolFunds)?;
64        let new_native_reserve = self
65            .constant_product
66            .checked_div(new_token_reserve as u128)
67            .ok_or(PoolError::Overflow)? as u64;
68        let native_sold = new_native_reserve - self.native_reserve;
69        if let Some(max_native) = max_native {
70            if native_sold > max_native {
71                return Err(PoolError::SlippageExceeded);
72            }
73        }
74        self.native_reserve = new_native_reserve;
75        self.token_reserve = new_token_reserve;
76        Ok(native_sold)
77    }
78
79    /// Sells `token_amount` tokens to the pool, checking if the native currency received is at least `min_native`.
80    pub fn sell(&mut self, token_amount: u64, min_native: Option<u64>) -> Result<u64> {
81        if token_amount == 0 {
82            return Err(PoolError::InvalidAmount);
83        }
84        if token_amount > self.token_reserve {
85            return Err(PoolError::InsufficientPoolFunds);
86        }
87        let new_token_reserve = self
88            .token_reserve
89            .checked_add(token_amount)
90            .ok_or(PoolError::Overflow)?;
91        let new_native_reserve = self
92            .constant_product
93            .checked_div(new_token_reserve as u128)
94            .ok_or(PoolError::Overflow)? as u64;
95        let native_bought = self.native_reserve - new_native_reserve;
96        if let Some(min_native) = min_native {
97            if native_bought < min_native {
98                return Err(PoolError::SlippageExceeded);
99            }
100        }
101        self.native_reserve = new_native_reserve;
102        self.token_reserve = new_token_reserve;
103        Ok(native_bought)
104    }
105
106    /// Simulates buying `token_amount` tokens and calculates the native currency that would be spent.
107    pub fn simulate_buy(&self, token_amount: u64, min_native: Option<u64>) -> Result<u64> {
108        if token_amount == 0 {
109            return Err(PoolError::InvalidAmount);
110        }
111        let new_token_reserve = self
112            .token_reserve
113            .checked_sub(token_amount)
114            .ok_or(PoolError::InsufficientPoolFunds)?;
115        let new_native_reserve = self
116            .constant_product
117            .checked_div(new_token_reserve as u128)
118            .ok_or(PoolError::Overflow)? as u64;
119        let native_sold = new_native_reserve - self.native_reserve;
120        if let Some(min_native) = min_native {
121            if native_sold < min_native {
122                return Err(PoolError::SlippageExceeded);
123            }
124        }
125        Ok(native_sold)
126    }
127
128    /// Simulates selling `token_amount` tokens and calculates the native currency that would be received.
129    pub fn simulate_sell(&self, token_amount: u64, max_native: Option<u64>) -> Result<u64> {
130        if token_amount == 0 {
131            return Err(PoolError::InvalidAmount);
132        }
133        if token_amount > self.token_reserve {
134            return Err(PoolError::InsufficientPoolFunds);
135        }
136        let new_token_reserve = self
137            .token_reserve
138            .checked_add(token_amount)
139            .ok_or(PoolError::Overflow)?;
140        let new_native_reserve = self
141            .constant_product
142            .checked_div(new_token_reserve as u128)
143            .ok_or(PoolError::Overflow)? as u64;
144        let native_sold = self.native_reserve - new_native_reserve;
145        if let Some(max_native) = max_native {
146            if native_sold > max_native {
147                return Err(PoolError::SlippageExceeded);
148            }
149        }
150        Ok(native_sold)
151    }
152
153    /// Calculates the amount of tokens that would be received for spending a specific amount of native currency.
154    pub fn calculate_tokens_received(&self, native_amount: u64) -> Result<u64> {
155        if native_amount == 0 {
156            return Err(PoolError::InvalidAmount);
157        }
158        let new_native_reserve = self.native_reserve + native_amount;
159        let new_token_reserve = self
160            .constant_product
161            .checked_div(new_native_reserve as u128)
162            .ok_or(PoolError::Overflow)? as u64;
163        Ok(self.token_reserve - new_token_reserve)
164    }
165
166    /// Buys tokens using a specified amount of native currency.
167    pub fn buy_tokens_with_native(&mut self, native_amount: u64) -> Result<u64> {
168        if native_amount == 0 {
169            return Err(PoolError::InvalidAmount);
170        }
171        let token_amount = self.calculate_tokens_received(native_amount)?;
172        self.buy(token_amount, None)?;
173        Ok(token_amount)
174    }
175
176    pub fn calculate_price_impact(&self, token_amount: u64) -> f64 {
177        let initial_price = self.market_price();
178        let new_token_reserve = self.token_reserve - token_amount;
179        let new_native_reserve = self.constant_product / new_token_reserve as u128;
180        let new_price = new_native_reserve as f64 / new_token_reserve as f64;
181        (new_price - initial_price) / initial_price
182    }
183
184    /// Calculates the number of additional tokens required to reach a desired native currency amount.
185    pub fn calculate_additional_tokens_for_desired_native(
186        &mut self,
187        sell_tokens: u64,
188        desired_native: u64,
189    ) -> Result<u64> {
190        if sell_tokens == 0 || desired_native == 0 {
191            return Err(PoolError::InvalidAmount);
192        }
193
194        let mut low = 0u64;
195        let mut high = self.token_reserve;
196        let mut best_guess = high;
197
198        // Binary search to find the best amount of tokens to buy
199        while low <= high {
200            let mid = low + (high - low) / 2;
201
202            // Clone the pool to simulate the impact of buying and selling without affecting the actual pool
203            let mut temp_pool = self.clone();
204
205            // Attempt to buy `mid` tokens
206            temp_pool.buy(mid, None)?;
207
208            // Simulate selling `sell_tokens` tokens to see how much native currency would be received
209            let native_received = temp_pool.simulate_sell(sell_tokens, None)?;
210
211            match native_received.cmp(&desired_native) {
212                Ordering::Equal => return Ok(mid),
213                Ordering::Less => {
214                    // Not enough native currency received, increase the number of tokens to buy
215                    low = mid.checked_add(1).ok_or(PoolError::Overflow)?;
216                }
217                Ordering::Greater => {
218                    // Too much native currency received, decrease the number of tokens to buy
219                    best_guess = mid;
220                    high = mid.checked_sub(1).ok_or(PoolError::Overflow)?;
221                }
222            }
223            // Exit the loop when the search range is narrowed sufficiently
224            if high - low <= 1 {
225                break;
226            }
227        }
228
229        Ok(best_guess)
230    }
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    impl LiquidityPool {
238        /// Checks the integrity of the pool by comparing the constant product after operations.
239        fn check_pool_integrity(&self) -> Result<f64> {
240            if self.native_reserve == 0 || self.token_reserve == 0 {
241                return Err(PoolError::InvalidAmount);
242            }
243            let k = self.native_reserve as u128 * self.token_reserve as u128;
244            Ok((k as f64 - self.constant_product as f64) / self.constant_product as f64)
245        }
246    }
247
248    impl Default for LiquidityPool {
249        fn default() -> Self {
250            let native_reserve = 10u64.pow(9);
251            let token_reserve = 1_000_000_000 * 10u64.pow(6);
252            Self::new(native_reserve, token_reserve).unwrap()
253        }
254    }
255
256    #[test]
257    fn test_buy() {
258        let mut pool = LiquidityPool::default();
259        let native_reserve = pool.get_native_reserve();
260        let token_reserve = pool.get_token_reserve();
261        let token_amount = 1_000_000 * 10u64.pow(6);
262        let native = pool.buy(token_amount, None).unwrap();
263        println!(
264            "Get {} tokens by paying {:.6} NATIVE",
265            token_amount / 10u64.pow(6),
266            native as f64 / 10u64.pow(9) as f64
267        );
268        assert_eq!(pool.get_native_reserve(), native_reserve + native);
269        assert_eq!(pool.get_token_reserve(), token_reserve - token_amount);
270        println!("Integrity: {}", pool.check_pool_integrity().unwrap());
271        assert!(pool.check_pool_integrity().unwrap().abs() < 0.000000001);
272    }
273
274    #[test]
275    fn test_sell() {
276        let mut pool = LiquidityPool::default();
277        let native_reserve = pool.get_native_reserve();
278        let token_reserve = pool.get_token_reserve();
279        let token_amount = 1_000_000 * 10u64.pow(6);
280        let native = pool.sell(token_amount, None).unwrap();
281        println!(
282            "Sell {} tokens and get {:.6} NATIVE",
283            token_amount / 10u64.pow(6),
284            native as f64 / 10u64.pow(9) as f64
285        );
286        assert_eq!(pool.get_native_reserve(), native_reserve - native);
287        assert_eq!(pool.get_token_reserve(), token_reserve + token_amount);
288        println!("Integrity: {}", pool.check_pool_integrity().unwrap());
289        assert!(pool.check_pool_integrity().unwrap().abs() < 0.000000001);
290    }
291
292    #[test]
293    fn test_simulate_buy() {
294        let pool = LiquidityPool::default();
295        let token_amount = 1_000_000 * 10u64.pow(6);
296        let native = pool.simulate_buy(token_amount, None).unwrap();
297        let new_k = (pool.get_native_reserve() + native) as u128
298            * (pool.get_token_reserve() - token_amount) as u128;
299        let difference = (new_k as f64 - pool.get_constant_product() as f64)
300            / pool.get_constant_product() as f64;
301        assert!(difference.abs() < 0.000000001);
302    }
303
304    #[test]
305    fn test_simulate_sell() {
306        let pool = LiquidityPool::default();
307        let token_amount = 1_000_000 * 10u64.pow(6);
308        let native = pool.simulate_sell(token_amount, None).unwrap();
309        println!(
310            "Simulate sell {} tokens and get {:.6} NATIVE",
311            token_amount / 10u64.pow(6),
312            native as f64 / 10u64.pow(9) as f64
313        );
314        let new_k = (pool.get_native_reserve() - native) as u128
315            * (pool.get_token_reserve() + token_amount) as u128;
316        let difference = (new_k as f64 - pool.get_constant_product() as f64)
317            / pool.get_constant_product() as f64;
318        println!("Integrity: {}", difference);
319        assert!(difference.abs() < 0.000000001);
320    }
321
322    #[test]
323    fn test_calculate_tokens_received() {
324        let mut pool = LiquidityPool::default();
325        let native_amount = 10u64.pow(9);
326        let token_amount = pool.calculate_tokens_received(native_amount).unwrap();
327        println!(
328            "Should get {:.6} tokens by paying {:.6} NATIVE",
329            token_amount as f64 / 10u64.pow(6) as f64,
330            native_amount as f64 / 10u64.pow(9) as f64
331        );
332        let native = pool.buy(token_amount, None).unwrap();
333        assert_eq!(native, native_amount);
334    }
335
336    #[test]
337    fn test_buy_tokens_with_native() {
338        let mut pool = LiquidityPool::default();
339        let native_amount = 10u64.pow(9);
340        let initial_native_reserve = pool.get_native_reserve();
341        let token_amount = pool.buy_tokens_with_native(native_amount).unwrap();
342        println!(
343            "Get {} tokens by paying {:.6} NATIVE",
344            token_amount / 10u64.pow(6),
345            native_amount as f64 / 10u64.pow(9) as f64
346        );
347        assert_eq!(
348            pool.get_native_reserve(),
349            initial_native_reserve + native_amount
350        );
351    }
352
353    #[test]
354    fn test_buy_invalid_slippage() {
355        let mut pool = LiquidityPool::default();
356        let token_amount = 1_000_000 * 10u64.pow(6);
357        let native_cost = pool.simulate_buy(token_amount, None).unwrap();
358        let result = pool.buy(token_amount, Some(native_cost - 1));
359        assert_eq!(result, Err(PoolError::SlippageExceeded));
360    }
361
362    #[test]
363    fn test_sell_invalid_slippage() {
364        let mut pool = LiquidityPool::default();
365        let token_amount = 1_000_000 * 10u64.pow(6);
366        let native_gain = pool.simulate_sell(token_amount, None).unwrap();
367        let result = pool.sell(token_amount, Some(native_gain + 1));
368        assert_eq!(result, Err(PoolError::SlippageExceeded));
369    }
370
371    #[test]
372    fn test_calculate_missing_tokens() {
373        let mut pool = LiquidityPool::default();
374        let tokens_to_buy = 50_000_000 * 10u64.pow(6);
375        let native_spent = pool.buy(tokens_to_buy, None).unwrap();
376
377        let additional_native_needed = native_spent + 10u64.pow(9);
378        let missing_tokens = pool
379            .calculate_additional_tokens_for_desired_native(tokens_to_buy, additional_native_needed)
380            .unwrap();
381
382        pool.buy(missing_tokens, None).unwrap();
383
384        let native_received = pool.sell(tokens_to_buy, None).unwrap();
385
386        assert!(native_received >= additional_native_needed);
387        assert!(native_received <= additional_native_needed + 1);
388    }
389
390    #[test]
391    fn test_many_operations() {
392        let mut pool = LiquidityPool::default();
393        let token_buy_amouts = [1_000_000, 10_000_000, 3_000_000, 3_000_000, 1_000_000];
394        let token_sell_amouts = [1_000_000, 1_000_000, 2_000_000, 1_000_000, 5_000_000];
395        for i in 0..5 {
396            let token_amount = token_buy_amouts[i] * 10u64.pow(6);
397            let _ = pool.buy(token_amount, None).unwrap();
398            let token_amount = token_sell_amouts[i] * 10u64.pow(6);
399            let _ = pool.sell(token_amount, None).unwrap();
400        }
401
402        assert!(pool.check_pool_integrity().unwrap().abs() < 0.000000001);
403    }
404}
consta_pool/lib.rs

consta_pool/
lib.rs
