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::{
268 ekubo_v3::test_cases::*, test_utils::try_decode_snapshot_with_defaults,
269 };
270
271 #[apply(all_cases)]
272 #[tokio::test]
273 async fn test_try_from_with_header(case: TestCase) {
274 let snapshot = ComponentWithState {
275 state: ResponseProtocolState {
276 attributes: case.state_attributes,
277 ..Default::default()
278 },
279 component: case.component,
280 component_tvl: None,
281 entrypoints: Vec::new(),
282 };
283
284 let result = try_decode_snapshot_with_defaults::<EkuboV3State>(snapshot)
285 .await
286 .expect("reconstructing state");
287
288 assert_eq!(result, case.state_before_transition);
289 }
290
291 #[apply(all_cases)]
292 #[tokio::test]
293 async fn test_try_from_invalid(case: TestCase) {
294 for missing_attribute in case.required_attributes {
295 let mut component = case.component.clone();
296 let mut attributes = case.state_attributes.clone();
297
298 component
299 .static_attributes
300 .remove(&missing_attribute);
301 attributes.remove(&missing_attribute);
302
303 let snapshot = ComponentWithState {
304 state: ResponseProtocolState {
305 attributes,
306 component_id: Default::default(),
307 balances: Default::default(),
308 },
309 component,
310 component_tvl: None,
311 entrypoints: Vec::new(),
312 };
313
314 EkuboV3State::try_from_with_header(
315 snapshot,
316 BlockHeader::default(),
317 &HashMap::default(),
318 &HashMap::default(),
319 &DecoderContext::new(),
320 )
321 .await
322 .unwrap_err();
323 }
324 }
325}