1use std::collections::HashMap;
2
3use lunarbase_pmm_math::U256;
4use tycho_client::feed::{synchronizer::ComponentWithState, BlockHeader};
5use tycho_common::{models::token::Token, Bytes};
6
7use super::state::{Address, LunarBaseState};
8use crate::protocol::{
9 errors::InvalidSnapshotError,
10 models::{DecoderContext, TryFromWithBlock},
11};
12
13mod attrs {
14 pub const ANCHOR_PRICE_X96: &str = "anchor_price_x96";
15 pub const FEE_ASK_X24: &str = "fee_ask_x24";
16 pub const FEE_BID_X24: &str = "fee_bid_x24";
17 pub const LATEST_UPDATE_BLOCK: &str = "latest_update_block";
18 pub const RESERVE_X: &str = "reserve_x";
19 pub const RESERVE_Y: &str = "reserve_y";
20 pub const MAX_PUNISHMENT_X24: &str = "max_punishment_x24";
21 pub const BLOCK_DELAY: &str = "block_delay";
22 pub const PAUSED: &str = "paused";
23}
24
25impl TryFromWithBlock<ComponentWithState, BlockHeader> for LunarBaseState {
26 type Error = InvalidSnapshotError;
27
28 async fn try_from_with_header(
29 snapshot: ComponentWithState,
30 block: BlockHeader,
31 _account_balances: &HashMap<Bytes, HashMap<Bytes, Bytes>>,
32 _all_tokens: &HashMap<Bytes, Token>,
33 _decoder_context: &DecoderContext,
34 ) -> Result<Self, Self::Error> {
35 let mut state = decode_lunarbase_snapshot(&snapshot)?;
36 state.head_block = block.number;
37 Ok(state)
38 }
39}
40
41#[cfg(test)]
42pub fn encode_state(state: &LunarBaseState) -> HashMap<String, Bytes> {
43 HashMap::from([
44 (
45 attrs::ANCHOR_PRICE_X96.to_owned(),
46 Bytes::from(
47 state
48 .anchor_price_x96
49 .to_be_bytes::<32>()[12..]
50 .to_vec(),
51 ),
52 ),
53 (attrs::FEE_ASK_X24.to_owned(), Bytes::from(state.fee_ask_x24)),
54 (attrs::FEE_BID_X24.to_owned(), Bytes::from(state.fee_bid_x24)),
55 (attrs::LATEST_UPDATE_BLOCK.to_owned(), Bytes::from(state.latest_update_block)),
56 (attrs::RESERVE_X.to_owned(), Bytes::from(state.reserve_x)),
57 (attrs::RESERVE_Y.to_owned(), Bytes::from(state.reserve_y)),
58 (attrs::MAX_PUNISHMENT_X24.to_owned(), Bytes::from(state.max_punishment_x24)),
59 (attrs::BLOCK_DELAY.to_owned(), Bytes::from(state.block_delay)),
60 (attrs::PAUSED.to_owned(), Bytes::from([u8::from(state.paused)])),
61 ])
62}
63
64pub fn apply_delta(
65 state: &mut LunarBaseState,
66 updated_attributes: HashMap<String, Bytes>,
67) -> Result<(), InvalidSnapshotError> {
68 let mut next = state.clone();
69 for (name, value) in updated_attributes {
70 match name.as_str() {
71 attrs::ANCHOR_PRICE_X96 => {
72 next.anchor_price_x96 = decode_uint(&name, &value, 160, 20)?;
73 }
74 attrs::FEE_ASK_X24 => next.fee_ask_x24 = decode_uint(&name, &value, 24, 4)?.to(),
75 attrs::FEE_BID_X24 => next.fee_bid_x24 = decode_uint(&name, &value, 24, 4)?.to(),
76 attrs::LATEST_UPDATE_BLOCK => {
77 next.latest_update_block = decode_uint(&name, &value, 48, 8)?.to();
78 }
79 attrs::RESERVE_X => next.reserve_x = decode_uint(&name, &value, 112, 16)?.to(),
80 attrs::RESERVE_Y => next.reserve_y = decode_uint(&name, &value, 112, 16)?.to(),
81 attrs::MAX_PUNISHMENT_X24 => {
82 next.max_punishment_x24 = decode_uint(&name, &value, 24, 4)?.to();
83 }
84 attrs::BLOCK_DELAY => next.block_delay = decode_uint(&name, &value, 48, 8)?.to(),
85 attrs::PAUSED => next.paused = decode_bool(attrs::PAUSED, &value)?,
86 "block_number" => next.head_block = decode_uint(&name, &value, 64, 8)?.to(),
87 _ => {}
88 }
89 }
90 *state = next;
91 Ok(())
92}
93
94pub fn decode_lunarbase_snapshot(
95 snapshot: &ComponentWithState,
96) -> Result<LunarBaseState, InvalidSnapshotError> {
97 let attrs = &snapshot.state.attributes;
98
99 Ok(LunarBaseState {
100 pool: component_pool(snapshot)?,
101 token_x: component_token(snapshot, 0)?,
102 token_y: component_token(snapshot, 1)?,
103 anchor_price_x96: decode_uint(
104 attrs::ANCHOR_PRICE_X96,
105 required_attr(attrs, attrs::ANCHOR_PRICE_X96)?,
106 160,
107 20,
108 )?,
109 fee_ask_x24: decode_uint(
110 attrs::FEE_ASK_X24,
111 required_attr(attrs, attrs::FEE_ASK_X24)?,
112 24,
113 4,
114 )?
115 .to(),
116 fee_bid_x24: decode_uint(
117 attrs::FEE_BID_X24,
118 required_attr(attrs, attrs::FEE_BID_X24)?,
119 24,
120 4,
121 )?
122 .to(),
123 latest_update_block: decode_uint(
124 attrs::LATEST_UPDATE_BLOCK,
125 required_attr(attrs, attrs::LATEST_UPDATE_BLOCK)?,
126 48,
127 8,
128 )?
129 .to(),
130 reserve_x: decode_uint(attrs::RESERVE_X, required_attr(attrs, attrs::RESERVE_X)?, 112, 16)?
131 .to(),
132 reserve_y: decode_uint(attrs::RESERVE_Y, required_attr(attrs, attrs::RESERVE_Y)?, 112, 16)?
133 .to(),
134 max_punishment_x24: decode_uint(
135 attrs::MAX_PUNISHMENT_X24,
136 required_attr(attrs, attrs::MAX_PUNISHMENT_X24)?,
137 24,
138 4,
139 )?
140 .to(),
141 block_delay: decode_uint(
142 attrs::BLOCK_DELAY,
143 required_attr(attrs, attrs::BLOCK_DELAY)?,
144 48,
145 8,
146 )?
147 .to(),
148 paused: decode_bool(attrs::PAUSED, required_attr(attrs, attrs::PAUSED)?)?,
149 head_block: 0,
150 })
151}
152
153fn component_pool(snapshot: &ComponentWithState) -> Result<Address, InvalidSnapshotError> {
154 address_from_component_id(&snapshot.component.id)
155}
156
157fn component_token(
158 snapshot: &ComponentWithState,
159 idx: usize,
160) -> Result<Address, InvalidSnapshotError> {
161 snapshot
162 .component
163 .tokens
164 .get(idx)
165 .map(|token| token.as_ref())
166 .ok_or_else(|| InvalidSnapshotError::ValueError(format!("missing token index {idx}")))
167 .and_then(address_from_bytes)
168}
169
170fn required_attr<'a>(
171 attrs: &'a HashMap<String, Bytes>,
172 name: &'static str,
173) -> Result<&'a Bytes, InvalidSnapshotError> {
174 attrs
175 .get(name)
176 .ok_or_else(|| InvalidSnapshotError::MissingAttribute(name.to_owned()))
177}
178
179fn decode_uint(
182 name: &str,
183 value: &Bytes,
184 bits: usize,
185 max_bytes: usize,
186) -> Result<U256, InvalidSnapshotError> {
187 if value.is_empty() || value.len() > max_bytes {
188 return Err(InvalidSnapshotError::ValueError(format!(
189 "attribute {name} must contain 1..={max_bytes} bytes, got {}",
190 value.len()
191 )));
192 }
193 let decoded = U256::from_be_slice(value.as_ref());
194 if decoded.bit_len() > bits {
195 return Err(InvalidSnapshotError::ValueError(format!(
196 "attribute {name} exceeds uint{bits}"
197 )));
198 }
199 Ok(decoded)
200}
201
202fn decode_bool(name: &'static str, value: &Bytes) -> Result<bool, InvalidSnapshotError> {
203 if value.len() != 1 {
204 return Err(invalid_length(name, 1, value.len()));
205 }
206 match value[0] {
207 0 => Ok(false),
208 1 => Ok(true),
209 other => Err(InvalidSnapshotError::ValueError(format!(
210 "attribute {name} must be 0 or 1, got {other}"
211 ))),
212 }
213}
214
215fn address_from_bytes(value: &[u8]) -> Result<Address, InvalidSnapshotError> {
216 value.try_into().map_err(|_| {
217 InvalidSnapshotError::ValueError(format!("expected 20-byte address, got {}", value.len()))
218 })
219}
220
221fn address_from_component_id(value: &str) -> Result<Address, InvalidSnapshotError> {
222 let value = value
223 .strip_prefix("0x")
224 .unwrap_or(value);
225 if value.len() != 40 {
226 return Err(InvalidSnapshotError::ValueError(format!(
227 "expected 20-byte hex address component id, got {value}"
228 )));
229 }
230
231 let mut out = [0u8; 20];
232 for (idx, byte) in out.iter_mut().enumerate() {
233 let start = idx * 2;
234 *byte = u8::from_str_radix(&value[start..start + 2], 16).map_err(|err| {
235 InvalidSnapshotError::ValueError(format!("invalid LunarBase component id hex: {err}"))
236 })?;
237 }
238 Ok(out)
239}
240
241fn invalid_length(name: &'static str, expected: usize, actual: usize) -> InvalidSnapshotError {
242 InvalidSnapshotError::ValueError(format!(
243 "attribute {name} has invalid length: expected {expected}, got {actual}"
244 ))
245}
246
247#[cfg(test)]
248mod tests {
249 use super::*;
250
251 fn addr(byte: u8) -> [u8; 20] {
252 [byte; 20]
253 }
254
255 fn state() -> LunarBaseState {
256 LunarBaseState {
257 pool: addr(9),
258 token_x: addr(1),
259 token_y: addr(2),
260 anchor_price_x96: U256::from(1u128 << 96),
261 fee_ask_x24: 10,
262 fee_bid_x24: 11,
263 latest_update_block: 100,
264 reserve_x: 1_000_000,
265 reserve_y: 2_000_000,
266 max_punishment_x24: 4096,
267 block_delay: 2,
268 paused: false,
269 head_block: 100,
270 }
271 }
272
273 #[test]
274 fn encodes_full_state_attributes() {
275 let attrs = encode_state(&state());
276
277 assert_eq!(
278 U256::from_be_slice(attrs[attrs::ANCHOR_PRICE_X96].as_ref()),
279 U256::from(1u128 << 96)
280 );
281 assert_eq!(u32::from(attrs[attrs::FEE_ASK_X24].clone()), 10);
282 assert_eq!(u64::from(attrs[attrs::LATEST_UPDATE_BLOCK].clone()), 100);
283 assert!(!decode_bool(attrs::PAUSED, &attrs[attrs::PAUSED]).unwrap());
284 }
285
286 #[test]
287 fn applies_partial_state_updated_delta() {
288 let mut state = state();
289 let updated = HashMap::from([
290 (attrs::ANCHOR_PRICE_X96.to_owned(), Bytes::from(2u128 << 96)),
291 (attrs::FEE_ASK_X24.to_owned(), Bytes::from(20u32)),
292 (attrs::FEE_BID_X24.to_owned(), Bytes::from(21u32)),
293 (attrs::LATEST_UPDATE_BLOCK.to_owned(), Bytes::from(101u64)),
294 ]);
295
296 apply_delta(&mut state, updated).unwrap();
297
298 assert_eq!(state.anchor_price_x96, U256::from(2u128 << 96));
299 assert_eq!(state.fee_ask_x24, 20);
300 assert_eq!(state.fee_bid_x24, 21);
301 assert_eq!(state.latest_update_block, 101);
302 assert_eq!(state.reserve_x, 1_000_000);
303 }
304
305 #[test]
306 fn accepts_full_uint160_anchor_and_legacy_sixteen_byte_encoding() {
307 let mut state = state();
308 let large = U256::from(1u64) << 159usize;
309 apply_delta(
310 &mut state,
311 HashMap::from([(
312 attrs::ANCHOR_PRICE_X96.to_owned(),
313 Bytes::from(large.to_be_bytes::<32>()[12..].to_vec()),
314 )]),
315 )
316 .unwrap();
317 assert_eq!(state.anchor_price_x96, large);
318
319 apply_delta(
320 &mut state,
321 HashMap::from([(attrs::ANCHOR_PRICE_X96.to_owned(), Bytes::from(1u128 << 96))]),
322 )
323 .unwrap();
324 assert_eq!(state.anchor_price_x96, U256::from(1u128 << 96));
325 }
326
327 #[test]
328 fn punishment_delta_updates_fees_without_refreshing_operator_block() {
329 let mut state = state();
330 apply_delta(
331 &mut state,
332 HashMap::from([
333 (attrs::FEE_ASK_X24.to_owned(), Bytes::from(123u32)),
334 (attrs::FEE_BID_X24.to_owned(), Bytes::from(456u32)),
335 (attrs::MAX_PUNISHMENT_X24.to_owned(), Bytes::from(16_778u32)),
336 ("block_number".to_owned(), Bytes::from(101u64)),
337 ]),
338 )
339 .unwrap();
340 assert_eq!((state.fee_ask_x24, state.fee_bid_x24), (123, 456));
341 assert_eq!(state.max_punishment_x24, 16_778);
342 assert_eq!(state.latest_update_block, 100);
343 assert_eq!(state.head_block, 101);
344 }
345
346 #[test]
347 fn rejects_malformed_and_out_of_range_attributes_without_partial_changes() {
348 let cases = [
349 (attrs::ANCHOR_PRICE_X96, Bytes::from(Vec::<u8>::new())),
350 (attrs::ANCHOR_PRICE_X96, Bytes::from(vec![0u8; 21])),
351 (attrs::FEE_ASK_X24, Bytes::from(1u32 << 24)),
352 (attrs::FEE_BID_X24, Bytes::from(vec![0u8; 5])),
353 (attrs::RESERVE_X, Bytes::from(1u128 << 112)),
354 (attrs::RESERVE_Y, Bytes::from(vec![0u8; 17])),
355 (attrs::MAX_PUNISHMENT_X24, Bytes::from(1u32 << 24)),
356 (attrs::LATEST_UPDATE_BLOCK, Bytes::from(1u64 << 48)),
357 (attrs::BLOCK_DELAY, Bytes::from(1u64 << 48)),
358 (attrs::PAUSED, Bytes::from([2u8])),
359 ("block_number", Bytes::from(vec![0u8; 9])),
360 ];
361 for (name, value) in cases {
362 let mut state = state();
363 let before = state.clone();
364 let mut delta = HashMap::from([(attrs::RESERVE_Y.to_owned(), Bytes::from(123u128))]);
365 delta.insert(name.to_owned(), value);
366 assert!(apply_delta(&mut state, delta).is_err(), "accepted {name}");
367 assert_eq!(state, before, "partially applied {name}");
368 }
369 }
370}