1use 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 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}