Optimize verification data size calculation (#2055)

Author: MauroToscano

crates/batcher/src/lib.rs

@@ -2050,22 +2050,9 @@ impl Batcher {
         client_msg: &SubmitProofMessage,
         ws_conn_sink: &WsMessageSink,
     ) -> bool {
-        let verification_data = match cbor_serialize(&client_msg.verification_data) {
-            Ok(data) => data,
-            // This should never happened, the user sent all his data serialized
-            Err(_) => {
-                error!("Proof serialization error");
-                send_message(
-                    ws_conn_sink.clone(),
-                    SubmitProofResponseMessage::Error("Proof serialization error".to_string()),
-                )
-                .await;
-                self.metrics.user_error(&["proof_serialization_error", ""]);
-                return false;
-            }
-        };
+        let verification_data_size = client_msg.verification_data.cbor_size_upper_bound();
 
-        if verification_data.len() > self.max_proof_size {
+        if verification_data_size > self.max_proof_size {
             error!("Proof size exceeds the maximum allowed size.");
             send_message(
                 ws_conn_sink.clone(),

crates/batcher/src/types/batch_queue.rs

@@ -1,9 +1,6 @@
-use aligned_sdk::{
-    common::{
-        constants::CBOR_ARRAY_MAX_OVERHEAD,
-        types::{NoncedVerificationData, VerificationDataCommitment},
-    },
-    communication::serialization::cbor_serialize,
+use aligned_sdk::common::{
+    constants::CBOR_ARRAY_MAX_OVERHEAD,
+    types::{NoncedVerificationData, VerificationDataCommitment},
 };
 use ethers::types::{Address, Signature, U256};
 use priority_queue::PriorityQueue;
@@ -124,14 +121,9 @@ pub(crate) type BatchQueue = PriorityQueue<BatchQueueEntry, BatchQueueEntryPrior
 /// Calculates the size of the batch represented by the given batch queue.
 pub(crate) fn calculate_batch_size(batch_queue: &BatchQueue) -> Result<usize, BatcherError> {
     let folded_result = batch_queue.iter().try_fold(0, |acc, (entry, _)| {
-        if let Ok(verification_data_bytes) =
-            cbor_serialize(&entry.nonced_verification_data.verification_data)
-        {
-            let current_batch_size = acc + verification_data_bytes.len();
-            ControlFlow::Continue(current_batch_size)
-        } else {
-            ControlFlow::Break(())
-        }
+        let verification_data_size = entry.nonced_verification_data.cbor_size_upper_bound();
+        let current_batch_size = acc + verification_data_size;
+        ControlFlow::<usize, usize>::Continue(current_batch_size)
     });
 
     if let ControlFlow::Continue(batch_size) = folded_result {
@@ -178,12 +170,7 @@ pub(crate) fn try_build_batch(
             // * Subtract this entry size to the size of the batch size.
             // * Push the current entry to the resulting batch queue.
 
-            // It is safe to call `.unwrap()` here since any serialization error should have been caught
-            // when calculating the total size of the batch with the `calculate_batch_size` function
-            let verification_data_size =
-                cbor_serialize(&entry.nonced_verification_data.verification_data)
-                    .unwrap()
-                    .len();
+            let verification_data_size = entry.nonced_verification_data.cbor_size_upper_bound();
             batch_size -= verification_data_size;
 
             finalized_batch.pop();

crates/sdk/src/common/types.rs

@@ -97,6 +97,34 @@ impl NoncedVerificationData {
             payment_service_addr,
         }
     }
+
+    /// Returns an upper bound for the CBOR encoding size without performing serialization.
+    /// Sums the length of all Vec<u8> fields in the inner VerificationData and adds 1024 bytes as overhead. This covers the constant size data and extra headers
+    pub fn cbor_size_upper_bound(&self) -> usize {
+        const CBOR_OVERHEAD_BYTES: usize = 1024;
+        let mut total_size = 0;
+
+        // Add length of proof Vec<u8>
+        total_size += self.verification_data.proof.len();
+
+        // Add length of pub_input Option<Vec<u8>>
+        if let Some(ref pub_input) = self.verification_data.pub_input {
+            total_size += pub_input.len();
+        }
+
+        // Add length of verification_key Option<Vec<u8>>
+        if let Some(ref verification_key) = self.verification_data.verification_key {
+            total_size += verification_key.len();
+        }
+
+        // Add length of vm_program_code Option<Vec<u8>>
+        if let Some(ref vm_program_code) = self.verification_data.vm_program_code {
+            total_size += vm_program_code.len();
+        }
+
+        // Add overhead bytes for the full NoncedVerificationData structure
+        total_size + CBOR_OVERHEAD_BYTES
+    }
 }
 
 // Defines a price estimate type for the user.