Merge pull request #333 from propeller-heads/hooks/tnl/ENG-5030-pass-block

fix: Update block in UniswapV4State and EVMPoolState

Author: tamaralipows

src/evm/decoder.rs

@@ -10,7 +10,7 @@ use alloy::primitives::{Address, U256};
 use thiserror::Error;
 use tokio::sync::{RwLock, RwLockReadGuard};
 use tracing::{debug, error, info, warn};
-use tycho_client::feed::{synchronizer::ComponentWithState, FeedMessage, HeaderLike};
+use tycho_client::feed::{synchronizer::ComponentWithState, BlockHeader, FeedMessage, HeaderLike};
 use tycho_common::{
     dto::{ChangeType, ProtocolStateDelta},
     models::{token::Token, Chain},
@@ -229,6 +229,11 @@ where
             .header
             .clone();
 
+        let block_number_or_timestamp = header
+            .clone()
+            .block_number_or_timestamp();
+        let current_block = header.clone().block();
+
         for (protocol, protocol_msg) in msg.state_msgs.iter() {
             // Add any new tokens
             if let Some(deltas) = protocol_msg.deltas.as_ref() {
@@ -698,9 +703,11 @@ where
 
                 // update states with protocol state deltas (attribute changes etc.)
                 for (id, update) in deltas.state_updates {
+                    let update_with_block =
+                        Self::add_block_info_to_delta(update, current_block.clone());
                     match Self::apply_update(
                         &id,
-                        update,
+                        update_with_block,
                         &mut updated_states,
                         &state_guard,
                         &all_balances,
@@ -733,9 +740,13 @@ where
 
                 // update remaining pools linked to updated contracts/updated balances
                 for pool in pools_to_update {
+                    let default_delta_with_block = Self::add_block_info_to_delta(
+                        ProtocolStateDelta::default(),
+                        current_block.clone(),
+                    );
                     match Self::apply_update(
                         &pool,
-                        ProtocolStateDelta::default(),
+                        default_delta_with_block,
                         &mut updated_states,
                         &state_guard,
                         &all_balances,
@@ -792,11 +803,31 @@ where
         }
 
         // Send the tick with all updated states
-        Ok(Update::new(header.block_number_or_timestamp(), updated_states, new_pairs)
+        Ok(Update::new(block_number_or_timestamp, updated_states, new_pairs)
             .set_removed_pairs(removed_pairs)
             .set_sync_states(msg.sync_states.clone()))
     }
 
+    /// Add block information (number and timestamp) to a ProtocolStateDelta
+    fn add_block_info_to_delta(
+        mut delta: ProtocolStateDelta,
+        block_header_opt: Option<BlockHeader>,
+    ) -> ProtocolStateDelta {
+        if let Some(header) = block_header_opt {
+            // Add block_number and block_timestamp attributes to ensure pool states
+            // receive current block information during delta_transition
+            delta.updated_attributes.insert(
+                "block_number".to_string(),
+                Bytes::from(header.number.to_be_bytes().to_vec()),
+            );
+            delta.updated_attributes.insert(
+                "block_timestamp".to_string(),
+                Bytes::from(header.timestamp.to_be_bytes().to_vec()),
+            );
+        }
+        delta
+    }
+
     fn apply_update(
         id: &String,
         update: ProtocolStateDelta,

src/evm/protocol/uniswap_v4/state.rs

@@ -712,6 +712,39 @@ impl ProtocolSim for UniswapV4State {
             }
         }
 
+        // Update block information if provided in the delta attributes
+        let block_number_bytes = delta
+            .updated_attributes
+            .get("block_number")
+            .ok_or_else(|| {
+                SimulationError::FatalError("block_number not found in updated attributes".into())
+            })?;
+
+        self.block.number = u64::from_be_bytes(
+            block_number_bytes[block_number_bytes.len() - 8..]
+                .try_into()
+                .map_err(|_| {
+                    TransitionError::DecodeError("Invalid block_number bytes".to_string())
+                })?,
+        );
+
+        let block_timestamp_bytes = delta
+            .updated_attributes
+            .get("block_timestamp")
+            .ok_or_else(|| {
+                SimulationError::FatalError(
+                    "block_timestamp not found in updated attributes".into(),
+                )
+            })?;
+
+        self.block.timestamp = u64::from_be_bytes(
+            block_timestamp_bytes[block_timestamp_bytes.len() - 8..]
+                .try_into()
+                .map_err(|_| {
+                    TransitionError::DecodeError("Invalid block_timestamp bytes".to_string())
+                })?,
+        );
+
         // Apply attribute changes
         if let Some(liquidity) = delta
             .updated_attributes
@@ -862,6 +895,46 @@ mod tests {
         )
     }
 
+    #[test]
+    fn test_delta_transition_missing_block_update() {
+        let block = BlockHeader {
+            number: 1000,
+            hash: Bytes::from_str(
+                "0x28d41d40f2ac275a4f5f621a636b9016b527d11d37d610a45ac3a821346ebf8c",
+            )
+            .expect("Invalid block hash"),
+            parent_hash: Bytes::from(vec![0; 32]),
+            revert: false,
+            timestamp: 1758201863,
+        };
+        let mut pool = UniswapV4State::new(
+            1000,
+            U256::from_str("1000").unwrap(),
+            UniswapV4Fees { zero_for_one: 100, one_for_zero: 90, lp_fee: 700 },
+            100,
+            60,
+            vec![TickInfo::new(120, 10000), TickInfo::new(180, -10000)],
+            block,
+        );
+
+        assert_eq!(pool.block.number, 1000);
+        assert_eq!(pool.block.timestamp, 1758201863);
+
+        let attributes: HashMap<String, Bytes> =
+            [("block_number".to_string(), Bytes::from(2000_u64.to_be_bytes().to_vec()))]
+                .into_iter()
+                .collect();
+
+        let delta = ProtocolStateDelta {
+            component_id: "State1".to_owned(),
+            updated_attributes: attributes,
+            deleted_attributes: HashSet::new(),
+        };
+
+        let result = pool.delta_transition(delta, &HashMap::new(), &Balances::default());
+        assert!(result.is_err())
+    }
+
     #[test]
     fn test_delta_transition() {
         let block = BlockHeader {
@@ -893,6 +966,8 @@ mod tests {
             ("fee".to_string(), Bytes::from(100_u32.to_be_bytes().to_vec())),
             ("ticks/-120/net_liquidity".to_string(), Bytes::from(10200_u64.to_be_bytes().to_vec())),
             ("ticks/120/net_liquidity".to_string(), Bytes::from(9800_u64.to_be_bytes().to_vec())),
+            ("block_number".to_string(), Bytes::from(2000_u64.to_be_bytes().to_vec())),
+            ("block_timestamp".to_string(), Bytes::from(1758201935_u64.to_be_bytes().to_vec())),
         ]
         .into_iter()
         .collect();
@@ -926,6 +1001,9 @@ mod tests {
                 .net_liquidity,
             9800
         );
+        // Verify block was updated
+        assert_eq!(pool.block.number, 2000);
+        assert_eq!(pool.block.timestamp, 1758201935);
     }
 
     #[tokio::test]

src/evm/protocol/vm/state.rs

@@ -656,6 +656,39 @@ where
         tokens: &HashMap<Bytes, Token>,
         balances: &Balances,
     ) -> Result<(), TransitionError<String>> {
+        // Update block information if provided in the delta attributes
+        let block_number_bytes = delta
+            .updated_attributes
+            .get("block_number")
+            .ok_or_else(|| {
+                SimulationError::FatalError("block_number not found in updated attributes".into())
+            })?;
+
+        self.block.number = u64::from_be_bytes(
+            block_number_bytes[block_number_bytes.len() - 8..]
+                .try_into()
+                .map_err(|_| {
+                    TransitionError::DecodeError("Invalid block_number bytes".to_string())
+                })?,
+        );
+
+        let block_timestamp_bytes = delta
+            .updated_attributes
+            .get("block_timestamp")
+            .ok_or_else(|| {
+                SimulationError::FatalError(
+                    "block_timestamp not found in updated attributes".into(),
+                )
+            })?;
+
+        self.block.timestamp = u64::from_be_bytes(
+            block_timestamp_bytes[block_timestamp_bytes.len() - 8..]
+                .try_into()
+                .map_err(|_| {
+                    TransitionError::DecodeError("Invalid block_timestamp bytes".to_string())
+                })?,
+        );
+
         if self.manual_updates {
             // Directly check for "update_marker" in `updated_attributes`
             if let Some(marker) = delta
@@ -1230,4 +1263,40 @@ mod tests {
             "New token balance should be unchanged"
         );
     }
+
+    #[tokio::test]
+    async fn test_delta_transition_block_update() {
+        let mut pool_state = setup_pool_state().await;
+
+        let initial_number = pool_state.block.number;
+        let initial_timestamp = pool_state.block.timestamp;
+
+        let new_block_number = initial_number + 1;
+        let new_timestamp = initial_timestamp + 100;
+
+        let attributes: HashMap<String, Bytes> = [
+            ("block_number".to_string(), Bytes::from(new_block_number.to_be_bytes().to_vec())),
+            ("block_timestamp".to_string(), Bytes::from(new_timestamp.to_be_bytes().to_vec())),
+        ]
+        .into_iter()
+        .collect();
+
+        let delta = ProtocolStateDelta {
+            component_id: "test_pool".to_owned(),
+            updated_attributes: attributes,
+            deleted_attributes: HashSet::new(),
+        };
+
+        let mut tokens = HashMap::new();
+        tokens.insert(dai().address, dai());
+        tokens.insert(bal().address, bal());
+        let balances = Balances::default();
+
+        pool_state
+            .delta_transition(delta, &tokens, &balances)
+            .unwrap();
+
+        assert_eq!(pool_state.block.number, new_block_number);
+        assert_eq!(pool_state.block.timestamp, new_timestamp);
+    }
 }