Small refactor

Author: ali-behjati

hermes/Cargo.lock

@@ -59,7 +59,6 @@ solana-sdk = "=1.13.3"
 solana-account-decoder = "=1.13.3"
 
 moka = { version = "0.11.0", features = ["future"] }
-derive_builder = "0.12.0"
 byteorder = "1.4.3"
 serde_qs = { version = "0.12.0", features = ["axum"] }
 

hermes/Cargo.toml

@@ -18,7 +18,10 @@ use {
             construct_message_states_proofs,
             store_wormhole_merkle_verified_message,
         },
-        storage::AccumulatorState,
+        storage::{
+            AccumulatorState,
+            CompletedAccumulatorState,
+        },
         types::{
             MessageState,
             ProofSet,
@@ -113,7 +116,7 @@ impl Store {
                     .unwrap_or(AccumulatorState {
                         slot,
                         accumulator_messages: None,
-                        wormhole_merkle_proof: None,
+                        wormhole_merkle_state: None,
                     });
                 accumulator_state.accumulator_messages = Some(accumulator_messages);
                 self.storage
@@ -128,21 +131,32 @@ impl Store {
             None => return Ok(()),
         };
 
-        let (accumulator_messages, wormhole_merkle_proof) =
-            match (state.accumulator_messages, state.wormhole_merkle_proof) {
-                (Some(accumulator_messages), Some(wormhole_merkle_proof)) => {
-                    (accumulator_messages, wormhole_merkle_proof)
-                }
-                _ => return Ok(()),
-            };
+        let completed_state = state.try_into();
+        let completed_state: CompletedAccumulatorState = match completed_state {
+            Ok(completed_state) => completed_state,
+            Err(_) => {
+                return Ok(());
+            }
+        };
+
+        // Once the accumulator reaches a complete state for a specific slot
+        // we can build the message states
+        self.build_message_states(completed_state).await?;
 
+        self.update_tx.send(()).await?;
+
+        Ok(())
+    }
+
+    async fn build_message_states(&self, completed_state: CompletedAccumulatorState) -> Result<()> {
         let wormhole_merkle_message_states_proofs =
-            construct_message_states_proofs(&accumulator_messages, &wormhole_merkle_proof)?;
+            construct_message_states_proofs(&completed_state)?;
 
         let current_time: UnixTimestamp =
             SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs() as _;
 
-        let message_states = accumulator_messages
+        let message_states = completed_state
+            .accumulator_messages
             .messages
             .iter()
             .enumerate()
@@ -158,7 +172,7 @@ impl Store {
                             .ok_or(anyhow!("Missing proof for message"))?
                             .clone(),
                     },
-                    accumulator_messages.slot,
+                    completed_state.slot,
                     current_time,
                 ))
             })
@@ -168,8 +182,6 @@ impl Store {
 
         self.storage.store_message_states(message_states).await?;
 
-        self.update_tx.send(()).await?;
-
         Ok(())
     }
 

hermes/src/store.rs

@@ -1,10 +1,10 @@
 use {
     crate::store::{
-        storage::AccumulatorState,
-        types::{
-            AccumulatorMessages,
-            MessageState,
+        storage::{
+            AccumulatorState,
+            CompletedAccumulatorState,
         },
+        types::MessageState,
         Store,
     },
     anyhow::{
@@ -32,6 +32,12 @@ use {
     },
 };
 
+#[derive(Clone, PartialEq, Debug)]
+pub struct WormholeMerkleState {
+    pub root: WormholeMerkleRoot,
+    pub vaa:  Vec<u8>,
+}
+
 #[derive(Clone, PartialEq, Debug)]
 pub struct WormholeMerkleMessageProof {
     pub vaa:   Vec<u8>,
@@ -40,20 +46,23 @@ pub struct WormholeMerkleMessageProof {
 
 pub async fn store_wormhole_merkle_verified_message(
     store: &Store,
-    proof: WormholeMerkleRoot,
+    root: WormholeMerkleRoot,
     vaa_bytes: Vec<u8>,
 ) -> Result<()> {
     let mut accumulator_state = store
         .storage
-        .fetch_accumulator_state(proof.slot)
+        .fetch_accumulator_state(root.slot)
         .await?
         .unwrap_or(AccumulatorState {
-            slot:                  proof.slot,
+            slot:                  root.slot,
             accumulator_messages:  None,
-            wormhole_merkle_proof: None,
+            wormhole_merkle_state: None,
         });
 
-    accumulator_state.wormhole_merkle_proof = Some((proof, vaa_bytes));
+    accumulator_state.wormhole_merkle_state = Some(WormholeMerkleState {
+        root,
+        vaa: vaa_bytes,
+    });
     store
         .storage
         .store_accumulator_state(accumulator_state)
@@ -62,9 +71,11 @@ pub async fn store_wormhole_merkle_verified_message(
 }
 
 pub fn construct_message_states_proofs(
-    accumulator_messages: &AccumulatorMessages,
-    wormhole_merkle_proof: &(WormholeMerkleRoot, Vec<u8>),
+    completed_accumulator_state: &CompletedAccumulatorState,
 ) -> Result<Vec<WormholeMerkleMessageProof>> {
+    let accumulator_messages = &completed_accumulator_state.accumulator_messages;
+    let wormhole_merkle_state = &completed_accumulator_state.wormhole_merkle_state;
+
     // Check whether the state is valid
     let merkle_acc = match MerkleAccumulator::<Keccak160>::from_set(
         accumulator_messages.messages.iter().map(|m| m.as_ref()),
@@ -73,9 +84,7 @@ pub fn construct_message_states_proofs(
         None => return Ok(vec![]), // It only happens when the message set is empty
     };
 
-    let (proof, vaa) = &wormhole_merkle_proof;
-
-    if merkle_acc.root != proof.root {
+    if merkle_acc.root != wormhole_merkle_state.root.root {
         return Err(anyhow!("Invalid merkle root"));
     }
 
@@ -84,7 +93,7 @@ pub fn construct_message_states_proofs(
         .iter()
         .map(|m| {
             Ok(WormholeMerkleMessageProof {
-                vaa:   vaa.clone(),
+                vaa:   wormhole_merkle_state.vaa.clone(),
                 proof: merkle_acc
                     .prove(m.as_ref())
                     .ok_or(anyhow!("Failed to prove message"))?,

hermes/src/store/proof/wormhole_merkle.rs

@@ -1,13 +1,19 @@
 use {
-    super::types::{
-        AccumulatorMessages,
-        MessageIdentifier,
-        MessageState,
-        MessageType,
-        RequestTime,
-        Slot,
+    super::{
+        proof::wormhole_merkle::WormholeMerkleState,
+        types::{
+            AccumulatorMessages,
+            MessageIdentifier,
+            MessageState,
+            MessageType,
+            RequestTime,
+            Slot,
+        },
+    },
+    anyhow::{
+        anyhow,
+        Result,
     },
-    anyhow::Result,
     async_trait::async_trait,
     pyth_sdk::PriceIdentifier,
     pythnet_sdk::wire::v1::WormholeMerkleRoot,
@@ -19,7 +25,32 @@ pub mod local_storage;
 pub struct AccumulatorState {
     pub slot:                  Slot,
     pub accumulator_messages:  Option<AccumulatorMessages>,
-    pub wormhole_merkle_proof: Option<(WormholeMerkleRoot, Vec<u8>)>,
+    pub wormhole_merkle_state: Option<WormholeMerkleState>,
+}
+
+#[derive(Clone, PartialEq, Debug)]
+pub struct CompletedAccumulatorState {
+    pub slot:                  Slot,
+    pub accumulator_messages:  AccumulatorMessages,
+    pub wormhole_merkle_state: WormholeMerkleState,
+}
+
+impl TryFrom<AccumulatorState> for CompletedAccumulatorState {
+    type Error = anyhow::Error;
+
+    fn try_from(state: AccumulatorState) -> Result<Self> {
+        let accumulator_messages = state
+            .accumulator_messages
+            .ok_or_else(|| anyhow!("missing accumulator messages"))?;
+        let wormhole_merkle_state = state
+            .wormhole_merkle_state
+            .ok_or_else(|| anyhow!("missing wormhole merkle state"))?;
+        Ok(Self {
+            slot: state.slot,
+            accumulator_messages,
+            wormhole_merkle_state,
+        })
+    }
 }
 
 #[derive(Clone, Copy)]

hermes/src/store/storage.rs

@@ -52,12 +52,6 @@ impl MessageExt for Message {
 
 pub type RawMessage = Vec<u8>;
 
-#[derive(Clone, PartialEq, Debug)]
-pub struct WormholeMerkleState {
-    pub digest_proof: Vec<u8>,
-    pub tree:         Option<Vec<Vec<u8>>>,
-}
-
 #[derive(Clone, PartialEq, Eq, Debug, Hash)]
 pub struct MessageIdentifier {
     pub price_id: PriceIdentifier,