Recursive proof

crates/proof/Cargo.toml

@@ -9,6 +9,7 @@ repository.workspace = true
 
 [dependencies]
 risc0-zkvm.workspace = true
+serde.workspace = true
 serde_json.workspace = true
 thiserror = { workspace = true, optional = true }
 tracing = { workspace = true, optional = true }

crates/proof/src/guest.rs

@@ -1,20 +1,17 @@
+use crate::types::StateCommit;
 use risc0_zkvm::{guest, sha::Sha256};
 use void_types::Block;
 
-pub struct Output<D> {
-    pub block_hash: [u8; 32],
-    pub block_height: u64,
-    pub pre_digest: D,
-    pub post_digest: D,
-}
+pub mod recursive;
 
 pub fn proof<W, D, F>(state_transition_function: F)
 where
     W: From<Vec<u8>>,
     F: Fn(&Block, &mut W),
     D: for<'a> From<&'a W>,
-    Vec<u8>: From<D>,
+    Vec<u8>: for<'a> From<&'a D>,
 {
+    // Read block from guest env and deserialize
     let block_bytes_len: u64 = guest::env::read();
     let mut block_bytes = vec![0u8; block_bytes_len as usize];
     guest::env::read_slice(&mut block_bytes);
@@ -23,6 +20,7 @@ where
 
     let block = Block::from_bytes(&block_bytes).expect("Block failed to deserialize");
 
+    // Read witness from guest env and deserialize
     let witness_bytes_len: u64 = guest::env::read();
     let mut witness_bytes = vec![0u8; witness_bytes_len as usize];
     guest::env::read_slice(&mut witness_bytes);
@@ -35,52 +33,53 @@ where
 
     let post_digest = D::from(&witness);
 
-    let output = Output {
+    let output = StateCommit {
         block_hash: (*block_hash).into(),
-        // block_height: block.height,
-        block_height: 0,
+        parent_hash: block.parent_hash,
+        block_height: block.height,
         pre_digest,
         post_digest,
     };
 
-    let output = to_output_bytes(output);
+    let mut bytes = Vec::new();
+    to_output_bytes(&mut bytes, &output);
 
-    guest::env::commit_slice(&output);
+    guest::env::commit_slice(&bytes);
 }
 
-pub fn to_output_bytes<D>(output: Output<D>) -> Vec<u8>
+pub fn to_output_bytes<D>(bytes: &mut Vec<u8>, output: &StateCommit<D>)
 where
-    Vec<u8>: From<D>,
+    Vec<u8>: for<'a> From<&'a D>,
 {
-    let mut bytes = Vec::new();
     bytes.extend(&output.block_hash);
+    bytes.extend(&output.parent_hash);
     bytes.extend(&output.block_height.to_be_bytes());
-    let pre_bytes: Vec<u8> = output.pre_digest.into();
+    let pre_bytes: Vec<u8> = (&output.pre_digest).into();
     bytes.extend(&(pre_bytes.len() as u64).to_be_bytes());
     bytes.extend(&pre_bytes);
-    let post_bytes: Vec<u8> = output.post_digest.into();
+    let post_bytes: Vec<u8> = (&output.post_digest).into();
     bytes.extend(&(post_bytes.len() as u64).to_be_bytes());
     bytes.extend(&post_bytes);
-    bytes
 }
 
-pub fn from_output_bytes<D>(bytes: &[u8]) -> Option<Output<D>>
+pub fn from_output_bytes<D>(bytes: &[u8]) -> Option<StateCommit<D>>
 where
     D: for<'a> TryFrom<&'a [u8]>,
 {
-    if bytes.len() < 32 + 8 + 8 {
+    if bytes.len() < 32 + 32 + 8 + 8 {
         return None;
     }
     let block_hash: [u8; 32] = bytes[0..32].try_into().expect("Converting to sized array");
+    let parent_hash: [u8; 32] = bytes[32..64].try_into().expect("Converting to sized array");
     let block_height =
-        u64::from_be_bytes(bytes[32..40].try_into().expect("Converting to sized array"));
+        u64::from_be_bytes(bytes[64..72].try_into().expect("Converting to sized array"));
     let pre_len =
-        u64::from_be_bytes(bytes[40..48].try_into().expect("Converting to sized array")) as usize;
-    if bytes.len() < 48 + pre_len + 8 {
+        u64::from_be_bytes(bytes[72..80].try_into().expect("Converting to sized array")) as usize;
+    if bytes.len() < 80 + pre_len + 8 {
         return None;
     }
-    let pre_bytes = &bytes[48..48 + pre_len];
-    let post_len_start = 48 + pre_len;
+    let pre_bytes = &bytes[80..80 + pre_len];
+    let post_len_start = 80 + pre_len;
     let post_len = u64::from_be_bytes(
         bytes[post_len_start..post_len_start + 8]
             .try_into()
@@ -90,8 +89,9 @@ where
         return None;
     }
     let post_bytes = &bytes[post_len_start + 8..post_len_start + 8 + post_len];
-    Some(Output {
+    Some(StateCommit {
         block_hash,
+        parent_hash,
         block_height,
         // TODO: check that it is safe to lose the error
         pre_digest: D::try_from(pre_bytes).ok()?,

crates/proof/src/guest/recursive.rs

@@ -0,0 +1,162 @@
+use crate::types::{RecursiveCommit, RecursiveType, StateCommit};
+use risc0_zkvm::guest::env;
+use std::fmt::Debug;
+
+pub fn proof<D>(state_image_id: [u32; 8])
+where
+    D: for<'a> TryFrom<&'a [u8]> + Debug + Default + PartialEq,
+    Vec<u8>: for<'a> From<&'a D>,
+{
+    let state_commits_count = env::read();
+    let mut state_commits = Vec::with_capacity(state_commits_count);
+
+    for _ in 0..state_commits_count {
+        let state_commit_len = env::read();
+
+        let mut state_commit_bytes = vec![0u8; state_commit_len];
+        env::read_slice(&mut state_commit_bytes);
+
+        let state_commit: StateCommit<D> = super::from_output_bytes(&state_commit_bytes)
+            .expect("Failed to deserialize state commit");
+
+        state_commits.push(state_commit);
+    }
+
+    let recursive_type: RecursiveType = env::read();
+
+    match recursive_type {
+        RecursiveType::Init => {
+            let image_id: [u32; 8] = env::read();
+
+            let first_state_commit = state_commits
+                .first()
+                .expect("Must have at least one state commit");
+
+            assert_eq!(first_state_commit.parent_hash, [0u8; 32]);
+            assert_eq!(first_state_commit.pre_digest, D::default());
+            assert_eq!(first_state_commit.block_height, 0);
+
+            let mut bytes = Vec::new();
+            super::to_output_bytes(&mut bytes, first_state_commit);
+
+            env::verify(state_image_id, &bytes).expect("init state proof failed to verify");
+
+            verify_state_steps(state_image_id, &state_commits);
+
+            let recursive_commit = RecursiveCommit {
+                this_image_id: image_id,
+                state_commit: state_commits
+                    .pop()
+                    .expect("Must have at least one state commit"),
+            };
+
+            let mut bytes = Vec::new();
+            to_output_bytes(&mut bytes, &recursive_commit);
+            env::commit_slice(&bytes[..]);
+        }
+        RecursiveType::Step => {
+            let recursive_commit_len = env::read();
+
+            let mut recursive_commit_bytes = vec![0u8; recursive_commit_len];
+            env::read_slice(&mut recursive_commit_bytes);
+
+            let recursive_commit: RecursiveCommit<D> = from_output_bytes(&recursive_commit_bytes)
+                .expect("Failed to deserialize state commit");
+
+            let first_state_commit = state_commits
+                .first()
+                .expect("Must have at least one state commit");
+
+            verify_step(
+                state_image_id,
+                &recursive_commit.state_commit,
+                first_state_commit,
+            );
+            verify_state_steps(state_image_id, &state_commits);
+
+            let recursive_commit = RecursiveCommit {
+                this_image_id: recursive_commit.this_image_id,
+                state_commit: state_commits
+                    .pop()
+                    .expect("Must have at least one state commit"),
+            };
+
+            let mut bytes = Vec::new();
+            to_output_bytes(&mut bytes, &recursive_commit);
+            env::commit_slice(&bytes[..]);
+        }
+    }
+}
+
+// Verify a single step in the state commit chain.
+fn verify_step<D>(image_id: [u32; 8], parent: &StateCommit<D>, commit: &StateCommit<D>)
+where
+    D: Debug + PartialEq,
+    Vec<u8>: for<'a> From<&'a D>,
+{
+    assert_eq!(commit.parent_hash, parent.block_hash);
+    assert_eq!(commit.pre_digest, parent.post_digest);
+    assert_eq!(
+        commit.block_height,
+        parent.block_height.checked_add(1).unwrap()
+    );
+
+    let mut bytes = Vec::new();
+    super::to_output_bytes(&mut bytes, commit);
+
+    env::verify(image_id, &bytes).expect("step state proof failed to verify");
+}
+
+// Verify a state commit chain.
+fn verify_state_steps<D>(image_id: [u32; 8], state_commits: &[StateCommit<D>])
+where
+    D: Debug + PartialEq,
+    Vec<u8>: for<'a> From<&'a D>,
+{
+    for (parent, commit) in state_commits.iter().zip(state_commits.iter().skip(1)) {
+        verify_step(image_id, parent, commit);
+    }
+}
+
+pub fn to_output_bytes<D>(bytes: &mut Vec<u8>, commit: &RecursiveCommit<D>)
+where
+    Vec<u8>: for<'a> From<&'a D>,
+{
+    let mut image_id: [u8; 32] = [0u8; 32];
+    for (input, output) in commit
+        .this_image_id
+        .iter()
+        .map(|e| e.to_be_bytes())
+        .flatten()
+        .zip(image_id.iter_mut())
+    {
+        *output = input;
+    }
+    bytes.extend(image_id);
+
+    super::to_output_bytes(bytes, &commit.state_commit);
+}
+
+pub fn from_output_bytes<D>(bytes: &[u8]) -> Option<RecursiveCommit<D>>
+where
+    D: for<'a> TryFrom<&'a [u8]>,
+{
+    if bytes.len() < 32 {
+        return None;
+    }
+
+    let mut image_id: [u32; 8] = [0u32; 8];
+    for (input, output) in bytes[0..32]
+        .chunks_exact(std::mem::size_of::<u32>())
+        .zip(image_id.iter_mut())
+    {
+        *output = u32::from_be_bytes(input.try_into().expect("The sizes match"));
+    }
+
+    let state_commit = super::from_output_bytes(&bytes[32..])?;
+
+    Some(RecursiveCommit {
+        this_image_id: image_id,
+        state_commit,
+    })
+}

crates/proof/src/host.rs

@@ -3,6 +3,8 @@ use thiserror::Error;
 use tracing::{debug, info, instrument, warn};
 use void_types::Block;
 
+pub mod recursive;
+
 #[derive(Debug, Error)]
 pub enum ProverError {
     #[error("Failed to write to environment: {0}")]