sim/
lib.rs

1//! Simulations of the StableSwap invariant compared to Curve's reference implementation.
2
3use pyo3::prelude::*;
4use pyo3::types::PyTuple;
5use std::fs::File;
6use std::io::prelude::*;
7
8pub const MODEL_FEE_NUMERATOR: u64 = 10000000;
9pub const MODEL_FEE_DENOMINATOR: u64 = 10000000000;
10
11const DEFAULT_POOL_TOKENS: u64 = 0;
12const DEFAULT_TARGET_PRICE: u128 = 1000000000000000000;
13const FILE_NAME: &str = "simulation.py";
14const FILE_PATH: &str = "sim/simulation.py";
15const MODULE_NAME: &str = "simulation";
16
17pub struct Model {
18    pub py_src: String,
19    pub amp_factor: u64,
20    pub balances: Vec<u64>,
21    pub n_coins: u8,
22    pub target_prices: Vec<u128>,
23    pub pool_tokens: u64,
24}
25
26impl Model {
27    /// Constructs a new [`Model`].
28    pub fn new(amp_factor: u64, balances: Vec<u64>, n_coins: u8) -> Model {
29        let src_file = File::open(FILE_PATH);
30        let mut src_file = match src_file {
31            Ok(file) => file,
32            Err(error) => {
33                panic!("{:?}\n Please run `curl -L
34            https://raw.githubusercontent.com/curvefi/curve-contract/master/tests/simulation.py > sim/simulation.py`", error)
35            }
36        };
37        let mut src_content = String::new();
38        let _ = src_file.read_to_string(&mut src_content);
39
40        Self {
41            py_src: src_content,
42            amp_factor,
43            balances,
44            n_coins,
45            target_prices: vec![DEFAULT_TARGET_PRICE, DEFAULT_TARGET_PRICE],
46            pool_tokens: DEFAULT_POOL_TOKENS,
47        }
48    }
49
50    pub fn new_with_pool_tokens(
51        amp_factor: u64,
52        balances: Vec<u64>,
53        n_coins: u8,
54        pool_token_amount: u64,
55    ) -> Model {
56        let src_file = File::open(FILE_PATH);
57        let mut src_file = match src_file {
58            Ok(file) => file,
59            Err(error) => {
60                panic!("{:?}\n Please run `curl -L
61            https://raw.githubusercontent.com/curvefi/curve-contract/master/tests/simulation.py > sim/simulation.py`", error)
62            }
63        };
64        let mut src_content = String::new();
65        let _ = src_file.read_to_string(&mut src_content);
66
67        Self {
68            py_src: src_content,
69            amp_factor,
70            balances,
71            n_coins,
72            target_prices: vec![DEFAULT_TARGET_PRICE, DEFAULT_TARGET_PRICE],
73            pool_tokens: pool_token_amount,
74        }
75    }
76
77    pub fn sim_d(&self) -> u128 {
78        let gil = Python::acquire_gil();
79        return self
80            .call0(gil.python(), "D")
81            .unwrap()
82            .extract(gil.python())
83            .unwrap();
84    }
85
86    pub fn sim_dy(&self, i: u128, j: u128, dx: u128) -> u128 {
87        let gil = Python::acquire_gil();
88        return self
89            .call1(gil.python(), "dy", (i, j, dx))
90            .unwrap()
91            .extract(gil.python())
92            .unwrap();
93    }
94
95    pub fn sim_exchange(&self, i: u128, j: u128, dx: u128) -> u64 {
96        let gil = Python::acquire_gil();
97        return self
98            .call1(gil.python(), "exchange", (i, j, dx))
99            .unwrap()
100            .extract(gil.python())
101            .unwrap();
102    }
103
104    pub fn sim_xp(&self) -> Vec<u128> {
105        let gil = Python::acquire_gil();
106        return self
107            .call0(gil.python(), "xp")
108            .unwrap()
109            .extract(gil.python())
110            .unwrap();
111    }
112
113    pub fn sim_y(&self, i: u128, j: u128, x: u64) -> u128 {
114        let gil = Python::acquire_gil();
115        return self
116            .call1(gil.python(), "y", (i, j, x))
117            .unwrap()
118            .extract(gil.python())
119            .unwrap();
120    }
121
122    pub fn sim_y_d(&self, i: u128, d: u128) -> u128 {
123        let gil = Python::acquire_gil();
124        return self
125            .call1(gil.python(), "y_D", (i, d))
126            .unwrap()
127            .extract(gil.python())
128            .unwrap();
129    }
130
131    pub fn sim_remove_liquidity_imbalance(&self, amounts: Vec<u128>) -> u128 {
132        let gil = Python::acquire_gil();
133        return self
134            .call1(
135                gil.python(),
136                "remove_liquidity_imbalance",
137                PyTuple::new(gil.python(), amounts.to_vec()),
138            )
139            .unwrap()
140            .extract(gil.python())
141            .unwrap();
142    }
143
144    pub fn sim_calc_withdraw_one_coin(&self, token_amount: u64, i: u128) -> (u64, u64) {
145        let gil = Python::acquire_gil();
146        return self
147            .call1(gil.python(), "calc_withdraw_one_coin", (token_amount, i))
148            .unwrap()
149            .extract(gil.python())
150            .unwrap();
151    }
152
153    fn call0(&self, py: Python, method_name: &str) -> Result<PyObject, PyErr> {
154        let sim = PyModule::from_code(py, &self.py_src, FILE_NAME, MODULE_NAME).unwrap();
155        let model = sim
156            .getattr("Curve")?
157            .call1((
158                self.amp_factor,
159                self.balances.to_vec(),
160                self.n_coins,
161                self.target_prices.to_vec(),
162                self.pool_tokens,
163            ))
164            .unwrap()
165            .to_object(py);
166        let py_ret = model.as_ref(py).call_method0(method_name);
167        self.extract_py_ret(py, py_ret)
168    }
169
170    fn call1(
171        &self,
172        py: Python,
173        method_name: &str,
174        args: impl IntoPy<Py<PyTuple>>,
175    ) -> Result<PyObject, PyErr> {
176        let sim = PyModule::from_code(py, &self.py_src, FILE_NAME, MODULE_NAME).unwrap();
177        let model = sim
178            .getattr("Curve")?
179            .call1((
180                self.amp_factor,
181                self.balances.to_vec(),
182                self.n_coins,
183                self.target_prices.to_vec(),
184                self.pool_tokens,
185            ))
186            .unwrap()
187            .to_object(py);
188        let py_ret = model.as_ref(py).call_method1(method_name, args);
189        self.extract_py_ret(py, py_ret)
190    }
191
192    fn extract_py_ret(&self, py: Python, ret: PyResult<&PyAny>) -> Result<PyObject, PyErr> {
193        match ret {
194            Ok(v) => v.extract(),
195            Err(e) => {
196                e.print_and_set_sys_last_vars(py);
197                panic!("Python exeuction failed.")
198            }
199        }
200    }
201
202    pub fn print_src(&self) {
203        println!("{}", self.py_src);
204    }
205}