From fc22673786c446ad9eb0b62821d9d7b6b38360ba Mon Sep 17 00:00:00 2001 From: Roman Proskuryakoff Date: Tue, 9 Jul 2024 18:37:48 +0400 Subject: [PATCH] Implement borsh in guest code --- .../adapters/mock-zkvm/src/lib.rs | 67 +++++++------- .../adapters/risc0-bonsai/src/host.rs | 89 +++++++++++-------- .../adapters/risc0/src/guest/native.rs | 54 +++++------ .../adapters/risc0/src/guest/zkvm.rs | 19 +++- .../sovereign-sdk/adapters/risc0/src/host.rs | 40 +++++---- .../rollup-interface/src/state_machine/da.rs | 2 +- .../src/state_machine/zk/mod.rs | 2 +- 7 files changed, 141 insertions(+), 132 deletions(-) diff --git a/crates/sovereign-sdk/adapters/mock-zkvm/src/lib.rs b/crates/sovereign-sdk/adapters/mock-zkvm/src/lib.rs index 6cd917d4e..c930f5832 100644 --- a/crates/sovereign-sdk/adapters/mock-zkvm/src/lib.rs +++ b/crates/sovereign-sdk/adapters/mock-zkvm/src/lib.rs @@ -8,7 +8,8 @@ use std::sync::{Arc, Condvar, Mutex}; use anyhow::ensure; use borsh::{BorshDeserialize, BorshSerialize}; use serde::{Deserialize, Serialize}; -use sov_rollup_interface::zk::{Matches, ValidityCondition}; +use sov_rollup_interface::da::BlockHeaderTrait; +use sov_rollup_interface::zk::{Matches, StateTransitionData, ValidityCondition}; /// A mock commitment to a particular zkVM program. #[derive(Debug, Clone, PartialEq, Eq, BorshDeserialize, BorshSerialize, Serialize, Deserialize)] @@ -152,15 +153,14 @@ impl sov_rollup_interface::zk::ZkvmHost type Guest = MockZkGuest; fn add_hint(&mut self, item: T) { - unimplemented!() - // let hint = bincode::serialize(&item).unwrap(); - // let proof_info = ProofInfo { - // hint, - // validity_condition: self.validity_condition, - // }; - - // let data = bincode::serialize(&proof_info).unwrap(); - // self.committed_data.push_back(data) + let hint = borsh::to_vec(&item).unwrap(); + let proof_info = ProofInfo { + hint, + validity_condition: self.validity_condition, + }; + + let data = borsh::to_vec(&proof_info).unwrap(); + self.committed_data.push_back(data) } fn simulate_with_hints(&mut self) -> Self::Guest { @@ -173,32 +173,29 @@ impl sov_rollup_interface::zk::ZkvmHost Ok(sov_rollup_interface::zk::Proof::PublicInput(data)) } - fn extract_output< - Da: sov_rollup_interface::da::DaSpec, - Root: Serialize + serde::de::DeserializeOwned, - >( + fn extract_output( proof: &sov_rollup_interface::zk::Proof, ) -> Result, Self::Error> { - unimplemented!() - // match proof { - // sov_rollup_interface::zk::Proof::PublicInput(pub_input) => { - // let data: ProofInfo = bincode::deserialize(pub_input)?; - // let st: StateTransitionData = bincode::deserialize(&data.hint)?; - - // Ok(sov_rollup_interface::zk::StateTransition { - // initial_state_root: st.initial_state_root, - // final_state_root: st.final_state_root, - // validity_condition: data.validity_condition, - // state_diff: Default::default(), - // da_slot_hash: st.da_block_header_of_commitments.hash(), - // sequencer_public_key: vec![], - // sequencer_da_public_key: vec![], - // }) - // } - // sov_rollup_interface::zk::Proof::Full(_) => { - // panic!("Mock DA doesn't generate real proofs") - // } - // } + match proof { + sov_rollup_interface::zk::Proof::PublicInput(pub_input) => { + let data: ProofInfo = bincode::deserialize(pub_input)?; + let st: StateTransitionData = + BorshDeserialize::deserialize(&mut &*data.hint)?; + + Ok(sov_rollup_interface::zk::StateTransition { + initial_state_root: st.initial_state_root, + final_state_root: st.final_state_root, + validity_condition: data.validity_condition, + state_diff: Default::default(), + da_slot_hash: st.da_block_header_of_commitments.hash(), + sequencer_public_key: vec![], + sequencer_da_public_key: vec![], + }) + } + sov_rollup_interface::zk::Proof::Full(_) => { + panic!("Mock DA doesn't generate real proofs") + } + } } } @@ -238,7 +235,7 @@ impl sov_rollup_interface::zk::ZkvmGuest for MockZkGuest { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, BorshDeserialize, BorshSerialize, Serialize, Deserialize)] struct ProofInfo { hint: Vec, validity_condition: ValidityCond, diff --git a/crates/sovereign-sdk/adapters/risc0-bonsai/src/host.rs b/crates/sovereign-sdk/adapters/risc0-bonsai/src/host.rs index ad11841eb..ac260dfab 100644 --- a/crates/sovereign-sdk/adapters/risc0-bonsai/src/host.rs +++ b/crates/sovereign-sdk/adapters/risc0-bonsai/src/host.rs @@ -6,7 +6,7 @@ use std::time::Duration; use anyhow::anyhow; use bonsai_sdk::alpha as bonsai_sdk; -use borsh::BorshSerialize; +use borsh::{BorshDeserialize, BorshSerialize}; use risc0_zkvm::sha::Digest; use risc0_zkvm::{ compute_image_id, ExecutorEnvBuilder, ExecutorImpl, Groth16Receipt, InnerReceipt, Journal, @@ -142,7 +142,7 @@ impl BonsaiClient { let _ = notify.send(res); } BonsaiRequest::Download { url, notify } => { - debug!(%url, "Bonsai:upload_input"); + debug!(%url, "Bonsai:download"); let res = client.download(&url); let res = unwrap_bonsai_response!(res, 'client, 'queue); let _ = notify.send(res); @@ -326,18 +326,29 @@ impl<'a> Risc0BonsaiHost<'a> { } fn add_hint_bonsai(&mut self, item: T) { - unimplemented!() - // // For running in "prove" mode. - - // // Prepare input data and upload it. - // let client = self.client.as_ref().unwrap(); - - // let input_data = to_vec(&item).unwrap(); - // let input_data = bytemuck::cast_slice(&input_data).to_vec(); - // // handle error - // let input_id = client.upload_input(input_data); - // tracing::info!("Uploaded input with id: {}", input_id); - // self.last_input_id = Some(input_id); + // For running in "prove" mode. + + // Prepare input data and upload it. + let client = self.client.as_ref().unwrap(); + + let mut input_data = vec![]; + let mut buf = borsh::to_vec(&item).unwrap(); + // append [0..] alignment + let rem = buf.len() % 4; + if rem > 0 { + buf.extend(vec![0; 4 - rem]); + } + let buf_u32: &[u32] = bytemuck::cast_slice(&buf); + // write len(u64) in LE + let len = buf_u32.len() as u64; + input_data.extend(len.to_le_bytes()); + // write buf + input_data.extend(buf); + + // handle error + let input_id = client.upload_input(input_data); + tracing::info!("Uploaded input with id: {}", input_id); + self.last_input_id = Some(input_id); } } @@ -345,25 +356,26 @@ impl<'a> ZkvmHost for Risc0BonsaiHost<'a> { type Guest = Risc0Guest; fn add_hint(&mut self, item: T) { - unimplemented!() - // // For running in "execute" mode. - - // // We use the in-memory size of `item` as an indication of how much - // // space to reserve. This is in no way guaranteed to be exact, but - // // usually the in-memory size and serialized data size are quite close. - // // - // // Note: this is just an optimization to avoid frequent reallocations, - // // it's not actually required. - // self.env - // .reserve(std::mem::size_of::() / std::mem::size_of::()); - - // let mut serializer = risc0_zkvm::serde::Serializer::new(&mut self.env); - // item.serialize(&mut serializer) - // .expect("Risc0 hint serialization is infallible"); - - // if self.client.is_some() { - // self.add_hint_bonsai(item) - // } + // For running in "execute" mode. + + let mut buf = borsh::to_vec(&item).expect("Risc0 hint serialization is infallible"); + // append [0..] alignment to cast &[u8] to &[u32] + let rem = buf.len() % 4; + if rem > 0 { + buf.extend(vec![0; 4 - rem]); + } + let buf: &[u32] = bytemuck::cast_slice(&buf); + // write len(u64) in LE + let len = buf.len() as u64; + let len_buf = &len.to_le_bytes()[..]; + let len_buf: &[u32] = bytemuck::cast_slice(len_buf); + self.env.extend_from_slice(len_buf); + // write buf + self.env.extend_from_slice(buf); + + if self.client.is_some() { + self.add_hint_bonsai(item) + } } /// Guest simulation (execute mode) is run inside the Risc0 VM locally @@ -496,19 +508,20 @@ impl<'a> ZkvmHost for Risc0BonsaiHost<'a> { } } - fn extract_output( + fn extract_output( proof: &Proof, ) -> Result, Self::Error> { - match proof { + let journal = match proof { Proof::PublicInput(journal) => { let journal: Journal = bincode::deserialize(journal)?; - Ok(journal.decode()?) + journal } Proof::Full(data) => { let receipt: Receipt = bincode::deserialize(data)?; - Ok(receipt.journal.decode()?) + receipt.journal } - } + }; + Ok(BorshDeserialize::try_from_slice(&journal.bytes)?) } } diff --git a/crates/sovereign-sdk/adapters/risc0/src/guest/native.rs b/crates/sovereign-sdk/adapters/risc0/src/guest/native.rs index d965e4aff..822f1e503 100644 --- a/crates/sovereign-sdk/adapters/risc0/src/guest/native.rs +++ b/crates/sovereign-sdk/adapters/risc0/src/guest/native.rs @@ -1,7 +1,6 @@ //! This module implements the `ZkvmGuest` trait for the RISC0 VM. use borsh::{BorshDeserialize, BorshSerialize}; -use risc0_zkvm::serde::WordRead; use sov_rollup_interface::zk::ZkvmGuest; #[derive(Default)] @@ -19,37 +18,12 @@ impl Hints { } } -impl WordRead for Hints { - fn read_words(&mut self, words: &mut [u32]) -> risc0_zkvm::serde::Result<()> { - if let Some(slice) = self.values.get(self.position..self.position + words.len()) { - words.copy_from_slice(slice); - self.position += words.len(); - Ok(()) - } else { - Err(risc0_zkvm::serde::Error::DeserializeUnexpectedEnd) - } - } - - fn read_padded_bytes(&mut self, bytes: &mut [u8]) -> risc0_zkvm::serde::Result<()> { - use risc0_zkvm::align_up; - use risc0_zkvm_platform::WORD_SIZE; - - let remaining_bytes: &[u8] = bytemuck::cast_slice(&self.values[self.position..]); - if bytes.len() > remaining_bytes.len() { - return Err(risc0_zkvm::serde::Error::DeserializeUnexpectedEnd); - } - bytes.copy_from_slice(&remaining_bytes[..bytes.len()]); - self.position += align_up(bytes.len(), WORD_SIZE) / WORD_SIZE; - Ok(()) - } -} - /// A guest for the RISC0 VM. Implements the `ZkvmGuest` trait /// using interior mutability to test the functionality. #[derive(Default)] pub struct Risc0Guest { hints: std::sync::Mutex, - commits: std::sync::Mutex>, + // commits: std::sync::Mutex>, } impl Risc0Guest { @@ -62,21 +36,33 @@ impl Risc0Guest { pub fn with_hints(hints: Vec) -> Self { Self { hints: std::sync::Mutex::new(Hints::with_hints(hints)), - commits: Default::default(), + // commits: Default::default(), } } } impl ZkvmGuest for Risc0Guest { fn read_from_host(&self) -> T { - unimplemented!("read_from_host") - // let mut hints = self.hints.lock().unwrap(); - // let mut hints = hints.deref_mut(); - // T::deserialize(&mut Deserializer::new(&mut hints)).unwrap() + let mut hints = self.hints.lock().unwrap(); + let hints = &mut *hints; + let pos = &mut hints.position; + let env = &hints.values; + // read len(u64) in LE + let len_buf = &env[*pos..*pos + 2]; + let len_bytes = bytemuck::cast_slice(len_buf); + let len_bytes: [u8; 8] = len_bytes.try_into().expect("Exactly 4 bytes"); + let len = u64::from_le_bytes(len_bytes) as usize; + *pos += 2; + // read buf + let buf = &env[*pos..*pos + len]; + let buf: &[u8] = bytemuck::cast_slice(buf); + *pos += len; + // deserialize + BorshDeserialize::deserialize(&mut &*buf).unwrap() } - fn commit(&self, item: &T) { - unimplemented!("commit") + fn commit(&self, _item: &T) { + unimplemented!("commitment never used in a test code") // self.commits.lock().unwrap().extend_from_slice( // &risc0_zkvm::serde::to_vec(item).expect("Serialization to vec is infallible"), // ); diff --git a/crates/sovereign-sdk/adapters/risc0/src/guest/zkvm.rs b/crates/sovereign-sdk/adapters/risc0/src/guest/zkvm.rs index eab6a9ed6..cd718d200 100644 --- a/crates/sovereign-sdk/adapters/risc0/src/guest/zkvm.rs +++ b/crates/sovereign-sdk/adapters/risc0/src/guest/zkvm.rs @@ -1,6 +1,7 @@ //! This module implements the `ZkvmGuest` trait for the RISC0 VM. use borsh::{BorshDeserialize, BorshSerialize}; use risc0_zkvm::guest::env; +use risc0_zkvm::guest::env::Write; use sov_rollup_interface::zk::ZkvmGuest; /// A guest for the RISC0 VM. Implements the `ZkvmGuest` trait @@ -17,12 +18,22 @@ impl Risc0Guest { impl ZkvmGuest for Risc0Guest { fn read_from_host(&self) -> T { - unimplemented!() - // FIXME: env::read() + // read len(u64) in LE + let mut len_buf = [0u8; 8]; + env::read_slice(&mut len_buf); + let len = u64::from_le_bytes(len_buf); + // read buf + let mut buf: Vec = vec![0; len as usize]; + env::read_slice(&mut buf); + let slice: &[u8] = bytemuck::cast_slice(&buf); + // deserialize + BorshDeserialize::deserialize(&mut &*slice).expect("Failed to deserialize input from host") } fn commit(&self, item: &T) { - unimplemented!() - // FIXME: env::commit(item); + // use risc0_zkvm::guest::env::Write as _; + let buf = borsh::to_vec(item).expect("Serialization to vec is infallible"); + let mut journal = env::journal(); + journal.write_slice(&buf); } } diff --git a/crates/sovereign-sdk/adapters/risc0/src/host.rs b/crates/sovereign-sdk/adapters/risc0/src/host.rs index 4e668b6b2..2a0edbee4 100644 --- a/crates/sovereign-sdk/adapters/risc0/src/host.rs +++ b/crates/sovereign-sdk/adapters/risc0/src/host.rs @@ -1,6 +1,6 @@ //! This module implements the [`ZkvmHost`] trait for the RISC0 VM. -use borsh::BorshSerialize; +use borsh::{BorshDeserialize, BorshSerialize}; use risc0_zkvm::{ExecutorEnvBuilder, ExecutorImpl, InnerReceipt, Journal, Receipt, Session}; use serde::de::DeserializeOwned; use serde::Serialize; @@ -71,19 +71,20 @@ impl<'a> ZkvmHost for Risc0Host<'a> { type Guest = Risc0Guest; fn add_hint(&mut self, item: T) { - unimplemented!() - // // We use the in-memory size of `item` as an indication of how much - // // space to reserve. This is in no way guaranteed to be exact, but - // // usually the in-memory size and serialized data size are quite close. - // // - // // Note: this is just an optimization to avoid frequent reallocations, - // // it's not actually required. - // self.env - // .reserve(std::mem::size_of::() / std::mem::size_of::()); - - // let mut serializer = risc0_zkvm::serde::Serializer::new(&mut self.env); - // item.serialize(&mut serializer) - // .expect("Risc0 hint serialization is infallible"); + let mut buf = borsh::to_vec(&item).expect("Risc0 hint serialization is infallible"); + // append [0..] alignment to cast &[u8] to &[u32] + let rem = buf.len() % 4; + if rem > 0 { + buf.extend(vec![0; 4 - rem]); + } + let buf: &[u32] = bytemuck::cast_slice(&buf); + // write len(u64) in LE + let len = buf.len() as u64; + let len_buf = &len.to_le_bytes()[..]; + let len_buf: &[u32] = bytemuck::cast_slice(len_buf); + self.env.extend_from_slice(len_buf); + // write buf + self.env.extend_from_slice(buf); } fn simulate_with_hints(&mut self) -> Self::Guest { @@ -102,19 +103,20 @@ impl<'a> ZkvmHost for Risc0Host<'a> { } } - fn extract_output( + fn extract_output( proof: &Proof, ) -> Result, Self::Error> { - match proof { + let journal = match proof { Proof::PublicInput(journal) => { let journal: Journal = bincode::deserialize(journal)?; - Ok(journal.decode()?) + journal } Proof::Full(data) => { let receipt: Receipt = bincode::deserialize(data)?; - Ok(receipt.journal.decode()?) + receipt.journal } - } + }; + Ok(BorshDeserialize::deserialize(&mut journal.bytes.as_ref())?) } } diff --git a/crates/sovereign-sdk/rollup-interface/src/state_machine/da.rs b/crates/sovereign-sdk/rollup-interface/src/state_machine/da.rs index 7f1c60f01..ec9c0bd24 100644 --- a/crates/sovereign-sdk/rollup-interface/src/state_machine/da.rs +++ b/crates/sovereign-sdk/rollup-interface/src/state_machine/da.rs @@ -216,7 +216,7 @@ pub trait BlobReaderTrait: /// Trait with collection of trait bounds for a block hash. pub trait BlockHashTrait: // so it is compatible with StorageManager implementation? - BorshSerialize + Serialize + DeserializeOwned + PartialEq + Debug + Send + Sync + Clone + Eq + Into<[u8; 32]> + core::hash::Hash + BorshDeserialize + BorshSerialize + Serialize + DeserializeOwned + PartialEq + Debug + Send + Sync + Clone + Eq + Into<[u8; 32]> + core::hash::Hash { } diff --git a/crates/sovereign-sdk/rollup-interface/src/state_machine/zk/mod.rs b/crates/sovereign-sdk/rollup-interface/src/state_machine/zk/mod.rs index f682f6505..b9180be97 100644 --- a/crates/sovereign-sdk/rollup-interface/src/state_machine/zk/mod.rs +++ b/crates/sovereign-sdk/rollup-interface/src/state_machine/zk/mod.rs @@ -51,7 +51,7 @@ pub trait ZkvmHost: Zkvm + Clone { fn run(&mut self, with_proof: bool) -> Result; /// Extracts public input form the proof. - fn extract_output( + fn extract_output( proof: &Proof, ) -> Result, Self::Error>; }