Skip to content

Commit

Permalink
chore: add recursive proof using rocksDB
Browse files Browse the repository at this point in the history
  • Loading branch information
Prabhat1308 committed Dec 19, 2024
1 parent 25eb68e commit 3e00fa4
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 3 deletions.
9 changes: 9 additions & 0 deletions core/src/prover.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,21 @@ pub fn run<Z: ZKVMEnv>() {
let touched_states: StateUpdate = Z::read_input::<StateUpdate>().unwrap();
let header: AvailHeader = Z::read_input::<AvailHeader>().unwrap();
let mut header_store: HeaderStore = Z::read_input::<HeaderStore>().unwrap();
let img_id: [u32; 8] = Z::read_input::<[u32; 8]>().unwrap();

let zkvm_state_machine = ZKVMStateMachine::<Z>::new();
let zkvm_result = zkvm_state_machine
.execute_batch(&header, &mut header_store, &txs, touched_states)
.expect("Should not have panicked.");

if header_store.inner.clone().len() != 0 {
Z::verify(img_id, &header_store.inner.first().unwrap()).unwrap();
// = 0 only for the case of genesis
if header_store.inner.first().unwrap().avail_header_hash != header.parent_hash {
panic!("match not found");
}
}

// let after_stf = env::cycle_count();
// eprintln!("after STF {}", after_stf);

Expand Down
27 changes: 24 additions & 3 deletions nexus/host/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use nexus_core::{
types::{
AvailHeader, HeaderStore, NexusBlock, NexusBlockWithPointers, NexusHeader,
Proof as NexusProof, Transaction, TransactionResult, TransactionStatus,
TransactionWithStatus, TransactionZKVM, TxParams, H256,
TransactionWithStatus, TransactionZKVM, TxParams, H256, StatementDigest,
},
zkvm::{
traits::{ZKVMEnv, ZKVMProof, ZKVMProver},
Expand Down Expand Up @@ -117,6 +117,7 @@ async fn execute_batch<
header: &AvailHeader,
header_store: &mut HeaderStore,
prover_mode: ProverMode,
prev_proof: Option<P>,
) -> Result<(P, NexusHeader, HashMap<H256, bool>, Option<TreeUpdateBatch>), Error>
where
<P as TryFrom<NexusProof>>::Error: std::fmt::Debug,
Expand Down Expand Up @@ -154,11 +155,17 @@ where
.collect();

let zkvm_txs = zkvm_txs?;
let img_id = StatementDigest(NEXUS_RUNTIME_ID);

zkvm_prover.add_input(&zkvm_txs).unwrap();
zkvm_prover.add_input(&state_update).unwrap();
zkvm_prover.add_input(&header).unwrap();
zkvm_prover.add_input(&header_store).unwrap();
zkvm_prover.add_input(&img_id).unwrap();
match prev_proof {
Some(i) => {zkvm_prover.add_proof_for_recursion(i).unwrap();}
_ => ()
}
let mut proof = zkvm_prover.prove()?;

let result: NexusHeader = proof.public_inputs()?;
Expand All @@ -181,6 +188,17 @@ pub async fn execution_engine_handle(
) -> Result<(), anyhow::Error> {
const MAX_HEADERS: usize = 5;
let mut header_array: Vec<Header> = Vec::new();

// Initialize prev_proof from RocksDB
let mut prev_proof: Option<Proof> = {
let db_lock = node_db.lock().await;
match db_lock.get::<Proof>(b"prev_proof") {
Ok(Some(proof)) => Some(proof),
Ok(None) => None,
Err(_) => return Err(anyhow!("Failed to retrieve prev_proof from RocksDB")),
}
};

loop {
if *shutdown_rx.borrow() {
println!("Shutdown signal received. Stopping execution engine...");
Expand Down Expand Up @@ -247,11 +265,14 @@ pub async fn execution_engine_handle(
&AvailHeader::from(&header),
&mut old_headers,
prover_mode.clone(),
prev_proof.clone()
)
.await
{
Ok((_, result, tx_result, tree_update_batch)) => {
//The execute_batch method on state machine would have updated the version in the storage.
Ok((proof, result, tx_result, tree_update_batch)) => {
prev_proof = Some(proof.clone());
let db_lock = node_db.lock().await;
db_lock.put(b"prev_proof", &proof)?;
let updated_version = state.lock().await.get_version(false)?;

save_batch_information(
Expand Down

0 comments on commit 3e00fa4

Please sign in to comment.