1use std::collections::HashMap;
2
3use alloy::primitives::aliases::B32;
4use ekubo_sdk::{
5 chain::evm::{EvmBasePoolKey, EvmOraclePoolKey, EvmPoolTypeConfig, EvmTwammPoolKey},
6 quoting::{
7 pools::{
8 full_range::{FullRangePoolKey, FullRangePoolState, FullRangePoolTypeConfig},
9 mev_capture::MevCapturePoolKey,
10 oracle::OraclePoolState,
11 stableswap::{StableswapPoolKey, StableswapPoolState},
12 twamm::TwammPoolState,
13 },
14 types::PoolConfig,
15 },
16 U256,
17};
18use itertools::Itertools;
19use revm::primitives::Address;
20use tycho_client::feed::{synchronizer::ComponentWithState, BlockHeader};
21use tycho_common::{models::token::Token, Bytes};
22
23use super::{
24 attributes::{sale_rate_deltas_from_attributes, ticks_from_attributes},
25 pool::{base::BasePool, full_range::FullRangePool, oracle::OraclePool, twamm::TwammPool},
26 state::EkuboV3State,
27};
28use crate::{
29 evm::protocol::ekubo_v3::pool::{mev_capture::MevCapturePool, stableswap::StableswapPool},
30 protocol::{
31 errors::InvalidSnapshotError,
32 models::{DecoderContext, TryFromWithBlock},
33 },
34};
35
36enum EkuboExtension {
37 Base,
38 Oracle,
39 Twamm,
40 MevCapture,
41}
42
43impl TryFrom<Bytes> for EkuboExtension {
44 type Error = InvalidSnapshotError;
45
46 fn try_from(value: Bytes) -> Result<Self, Self::Error> {
47 match i32::from(value) {
49 0 => Err(InvalidSnapshotError::ValueError("Unknown Ekubo extension".to_string())),
50 1 => Ok(Self::Base),
51 2 => Ok(Self::Oracle),
52 3 => Ok(Self::Twamm),
53 4 => Ok(Self::MevCapture),
54 discriminant => Err(InvalidSnapshotError::ValueError(format!(
55 "Unknown Ekubo extension discriminant {discriminant}"
56 ))),
57 }
58 }
59}
60
61impl TryFromWithBlock<ComponentWithState, BlockHeader> for EkuboV3State {
62 type Error = InvalidSnapshotError;
63
64 async fn try_from_with_header(
65 snapshot: ComponentWithState,
66 _block: BlockHeader,
67 _account_balances: &HashMap<Bytes, HashMap<Bytes, Bytes>>,
68 _all_tokens: &HashMap<Bytes, Token>,
69 _decoder_context: &DecoderContext,
70 ) -> Result<Self, Self::Error> {
71 let static_attrs = snapshot.component.static_attributes;
72 let state_attrs = snapshot.state.attributes;
73
74 let extension_id = attribute(&static_attrs, "extension_id")?
75 .clone()
76 .try_into()?;
77
78 let (token0, token1) = (
79 parse_address(attribute(&static_attrs, "token0")?, "token0")?,
80 parse_address(attribute(&static_attrs, "token1")?, "token1")?,
81 );
82
83 let fee = u64::from_be_bytes(
84 attribute(&static_attrs, "fee")?
85 .as_ref()
86 .try_into()
87 .map_err(|err| {
88 InvalidSnapshotError::ValueError(format!("fee length mismatch: {err:?}"))
89 })?,
90 );
91
92 let pool_type_config = EvmPoolTypeConfig::try_from(
93 B32::try_from(attribute(&static_attrs, "pool_type_config")?.as_ref()).map_err(
94 |err| {
95 InvalidSnapshotError::ValueError(format!(
96 "pool_type_config length mismatch: {err:?}"
97 ))
98 },
99 )?,
100 )
101 .map_err(|err| {
102 InvalidSnapshotError::ValueError(format!("parsing pool_type_config: {err}"))
103 })?;
104
105 let extension = parse_address(attribute(&static_attrs, "extension")?, "extension")?;
106
107 let liquidity = attribute(&state_attrs, "liquidity")?
108 .clone()
109 .into();
110
111 let sqrt_ratio = U256::try_from_be_slice(&attribute(&state_attrs, "sqrt_ratio")?[..])
112 .ok_or_else(|| InvalidSnapshotError::ValueError("invalid pool price".to_string()))?;
113
114 Ok(match extension_id {
115 EkuboExtension::Base => match pool_type_config {
116 EvmPoolTypeConfig::FullRange(pool_type_config) => {
117 Self::FullRange(FullRangePool::new(
118 FullRangePoolKey {
119 token0,
120 token1,
121 config: PoolConfig { extension, fee, pool_type_config },
122 },
123 FullRangePoolState { sqrt_ratio, liquidity },
124 )?)
125 }
126 EvmPoolTypeConfig::Stableswap(pool_type_config) => {
127 Self::Stableswap(StableswapPool::new(
128 StableswapPoolKey {
129 token0,
130 token1,
131 config: PoolConfig { extension, fee, pool_type_config },
132 },
133 StableswapPoolState { sqrt_ratio, liquidity },
134 )?)
135 }
136 EvmPoolTypeConfig::Concentrated(pool_type_config) => {
137 let tick = attribute(&state_attrs, "tick")?
138 .clone()
139 .into();
140
141 let mut ticks = ticks_from_attributes(state_attrs)
142 .map_err(InvalidSnapshotError::ValueError)?;
143
144 ticks.sort_unstable_by_key(|tick| tick.index);
145
146 Self::Base(BasePool::new(
147 EvmBasePoolKey {
148 token0,
149 token1,
150 config: PoolConfig { extension, fee, pool_type_config },
151 },
152 ticks,
153 sqrt_ratio,
154 liquidity,
155 tick,
156 )?)
157 }
158 },
159 EkuboExtension::Oracle => Self::Oracle(OraclePool::new(
160 EvmOraclePoolKey {
161 token0,
162 token1,
163 config: PoolConfig {
164 extension,
165 fee,
166 pool_type_config: FullRangePoolTypeConfig,
167 },
168 },
169 OraclePoolState {
170 full_range_pool_state: FullRangePoolState { sqrt_ratio, liquidity },
171 last_snapshot_time: 0, },
174 )?),
175 EkuboExtension::Twamm => {
176 let (token0_sale_rate, token1_sale_rate) = (
177 attribute(&state_attrs, "token0_sale_rate")?
178 .clone()
179 .into(),
180 attribute(&state_attrs, "token1_sale_rate")?
181 .clone()
182 .into(),
183 );
184
185 let last_execution_time = attribute(&state_attrs, "last_execution_time")?
186 .clone()
187 .into();
188
189 let mut virtual_order_deltas =
190 sale_rate_deltas_from_attributes(state_attrs, last_execution_time)
191 .map_err(InvalidSnapshotError::ValueError)?
192 .collect_vec();
193
194 virtual_order_deltas.sort_unstable_by_key(|delta| delta.time);
195
196 Self::Twamm(TwammPool::new(
197 EvmTwammPoolKey {
198 token0,
199 token1,
200 config: PoolConfig {
201 extension,
202 fee,
203 pool_type_config: FullRangePoolTypeConfig,
204 },
205 },
206 TwammPoolState {
207 full_range_pool_state: FullRangePoolState { sqrt_ratio, liquidity },
208 token0_sale_rate,
209 token1_sale_rate,
210 last_execution_time,
211 },
212 virtual_order_deltas,
213 )?)
214 }
215 EkuboExtension::MevCapture => {
216 let tick = attribute(&state_attrs, "tick")?
217 .clone()
218 .into();
219
220 let mut ticks =
221 ticks_from_attributes(state_attrs).map_err(InvalidSnapshotError::ValueError)?;
222
223 ticks.sort_unstable_by_key(|tick| tick.index);
224
225 let EvmPoolTypeConfig::Concentrated(pool_type_config) = pool_type_config else {
226 return Err(InvalidSnapshotError::ValueError(
227 "expected concentrated pool config for MEV-capture pool".to_string(),
228 ));
229 };
230
231 Self::MevCapture(MevCapturePool::new(
232 MevCapturePoolKey {
233 token0,
234 token1,
235 config: PoolConfig { extension, fee, pool_type_config },
236 },
237 ticks,
238 sqrt_ratio,
239 liquidity,
240 tick,
241 )?)
242 }
243 })
244 }
245}
246
247fn attribute<'a>(
248 map: &'a HashMap<String, Bytes>,
249 key: &str,
250) -> Result<&'a Bytes, InvalidSnapshotError> {
251 map.get(key)
252 .ok_or_else(|| InvalidSnapshotError::MissingAttribute(key.to_string()))
253}
254
255fn parse_address(bytes: &Bytes, attr_name: &str) -> Result<Address, InvalidSnapshotError> {
256 Address::try_from(&bytes[..])
257 .map_err(|err| InvalidSnapshotError::ValueError(format!("parsing {attr_name}: {err}")))
258}
259
260#[cfg(test)]
261mod tests {
262 use rstest::*;
263 use rstest_reuse::apply;
264 use tycho_common::dto::ResponseProtocolState;
265
266 use super::*;
267 use crate::evm::protocol::ekubo_v3::test_cases::*;
268
269 #[apply(all_cases)]
270 #[tokio::test]
271 async fn test_try_from_with_header(case: TestCase) {
272 let snapshot = ComponentWithState {
273 state: ResponseProtocolState {
274 attributes: case.state_attributes,
275 ..Default::default()
276 },
277 component: case.component,
278 component_tvl: None,
279 entrypoints: Vec::new(),
280 };
281
282 let result = EkuboV3State::try_from_with_header(
283 snapshot,
284 BlockHeader::default(),
285 &HashMap::new(),
286 &HashMap::new(),
287 &DecoderContext::new(),
288 )
289 .await
290 .expect("reconstructing state");
291
292 assert_eq!(result, case.state_before_transition);
293 }
294
295 #[apply(all_cases)]
296 #[tokio::test]
297 async fn test_try_from_invalid(case: TestCase) {
298 for missing_attribute in case.required_attributes {
299 let mut component = case.component.clone();
300 let mut attributes = case.state_attributes.clone();
301
302 component
303 .static_attributes
304 .remove(&missing_attribute);
305 attributes.remove(&missing_attribute);
306
307 let snapshot = ComponentWithState {
308 state: ResponseProtocolState {
309 attributes,
310 component_id: Default::default(),
311 balances: Default::default(),
312 },
313 component,
314 component_tvl: None,
315 entrypoints: Vec::new(),
316 };
317
318 EkuboV3State::try_from_with_header(
319 snapshot,
320 BlockHeader::default(),
321 &HashMap::default(),
322 &HashMap::default(),
323 &DecoderContext::new(),
324 )
325 .await
326 .unwrap_err();
327 }
328 }
329}