From 04f0a0b89a3c5b79037aba5aef0dffa63c232f3b Mon Sep 17 00:00:00 2001
From: LucasLvy <lucaslevy10@gmail.com>
Date: Fri, 20 Oct 2023 15:08:34 +0200
Subject: [PATCH] feat(os output): add DA

---
 src/io/mod.rs | 20 +++++++++++++++++++-
 src/lib.rs    | 14 ++++++++++++--
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/src/io/mod.rs b/src/io/mod.rs
index d4190345..b7bbbeff 100644
--- a/src/io/mod.rs
+++ b/src/io/mod.rs
@@ -121,8 +121,14 @@ pub struct StarknetOsOutput {
     pub messages_to_l1: Vec<Felt252>,
     /// List of messages from L1 handled in this block
     pub messages_to_l2: Vec<Felt252>,
+    /// List of the storage updates.
+    pub state_updates: Vec<Felt252>,
+    /// List of the newly declared contract classes.
+    pub contract_class_diff: Vec<Felt252>,
 }
+
 impl StarknetOsOutput {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         prev_state_root: Felt252,
         new_state_root: Felt252,
@@ -131,8 +137,20 @@ impl StarknetOsOutput {
         config_hash: Felt252,
         messages_to_l1: Vec<Felt252>,
         messages_to_l2: Vec<Felt252>,
+        state_updates: Vec<Felt252>,
+        contract_class_diff: Vec<Felt252>,
     ) -> Self {
-        Self { prev_state_root, new_state_root, block_number, block_hash, config_hash, messages_to_l1, messages_to_l2 }
+        Self {
+            prev_state_root,
+            new_state_root,
+            block_number,
+            block_hash,
+            config_hash,
+            messages_to_l1,
+            messages_to_l2,
+            state_updates,
+            contract_class_diff,
+        }
     }
 }
 
diff --git a/src/lib.rs b/src/lib.rs
index cd43fd88..a5bd22ba 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -128,8 +128,16 @@ impl SnOsRunner {
         let messages_to_l1 = os_output[1..messages_to_l1_size].to_vec();
 
         let os_output = &os_output[messages_to_l1_size + 1..];
-        let messages_to_l2 =
-            os_output[1..<usize>::from_be_bytes(os_output[0].to_be_bytes()[..8].try_into().unwrap())].to_vec();
+        let messages_to_l2_size = <usize>::from_be_bytes(os_output[0].to_be_bytes()[..8].try_into().unwrap());
+        let messages_to_l2 = os_output[1..messages_to_l2_size].to_vec();
+        let os_output = &os_output[messages_to_l2_size + 1..];
+
+        let state_updates_size = <usize>::from_be_bytes(os_output[0].to_be_bytes()[..8].try_into().unwrap());
+        let state_updates = os_output[1..state_updates_size].to_vec();
+        let os_output = &os_output[state_updates_size + 1..];
+
+        let contract_class_diff_size = <usize>::from_be_bytes(os_output[0].to_be_bytes()[..8].try_into().unwrap());
+        let contract_class_diff = os_output[1..contract_class_diff_size].to_vec();
         StarknetOsOutput::new(
             prev_state_root,
             new_state_root,
@@ -138,6 +146,8 @@ impl SnOsRunner {
             config_hash,
             messages_to_l1,
             messages_to_l2,
+            state_updates,
+            contract_class_diff,
         );
 
         vm.verify_auto_deductions().map_err(|e| SnOsError::Runner(e.into()))?;