feat: Update get_block_inputs for new BlockInputs

crates/block-producer/Cargo.toml

@@ -21,6 +21,7 @@ tracing-forest = ["miden-node-utils/tracing-forest"]
 async-trait           = { version = "0.1" }
 futures               = { version = "0.3" }
 itertools             = { workspace = true }
+miden-block-prover    = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "pgackst-batch-expiration" }
 miden-lib             = { workspace = true }
 miden-node-proto      = { workspace = true }
 miden-node-utils      = { workspace = true }

crates/block-producer/src/block_builder/mod.rs

@@ -1,29 +1,27 @@
-use std::{
-    collections::BTreeSet,
-    ops::{Add, Range},
-};
+use std::{collections::BTreeSet, ops::Range};
 
 use futures::FutureExt;
+use miden_block_prover::LocalBlockProver;
 use miden_node_utils::tracing::OpenTelemetrySpanExt;
 use miden_objects::{
     account::AccountId,
     batch::ProvenBatch,
-    block::{BlockNumber, ProvenBlock},
+    block::{BlockInputs, BlockNumber, ProposedBlock, ProvenBlock},
     note::{NoteHeader, NoteId, Nullifier},
     transaction::{InputNoteCommitment, OutputNote},
+    MIN_PROOF_SECURITY_LEVEL,
 };
 use rand::Rng;
 use tokio::time::Duration;
 use tracing::{instrument, Span};
 
 use crate::{
-    block::BlockInputs, errors::BuildBlockError, mempool::SharedMempool, store::StoreClient,
-    COMPONENT, SERVER_BLOCK_FREQUENCY,
+    errors::BuildBlockError, mempool::SharedMempool, store::StoreClient, COMPONENT,
+    SERVER_BLOCK_FREQUENCY,
 };
 
-pub(crate) mod prover;
-
-use self::prover::{block_witness::BlockWitness, BlockProver};
+// pub(crate) mod prover;
+// use self::prover::BlockProver;
 
 // BLOCK BUILDER
 // =================================================================================================
@@ -39,7 +37,7 @@ pub struct BlockBuilder {
     pub failure_rate: f64,
 
     pub store: StoreClient,
-    pub block_kernel: BlockProver,
+    // pub block_prover: LocalBlockProver,
 }
 
 impl BlockBuilder {
@@ -49,7 +47,7 @@ impl BlockBuilder {
             // Note: The range cannot be empty.
             simulated_proof_time: Duration::ZERO..Duration::from_millis(1),
             failure_rate: 0.0,
-            block_kernel: BlockProver::new(),
+            // block_prover: LocalBlockProver::new(MIN_PROOF_SECURITY_LEVEL),
             store,
         }
     }
@@ -125,52 +123,70 @@ impl BlockBuilder {
         selected_block: SelectedBlock,
     ) -> Result<BlockSummaryAndInputs, BuildBlockError> {
         let SelectedBlock { block_number: _, batches } = selected_block;
-        let summary = BlockSummary::summarize_batches(&batches);
+        // let summary = BlockSummary::summarize_batches(&batches);
+
+        // For a given set of batches, we need to get the following block inputs from the store:
+        // - note inclusion proofs for unauthenticated notes (not required to be complete due to the
+        //   possibility of note erasure)
+        // - a chain MMR with:
+        //   - all blocks referenced by batches
+        //   - all blocks referenced by note inclusion proofs
+        // - account witnesses for all accounts updated in the block
+        // - nullifier witnesses for all nullifiers created in the block
+        //   - since we don't yet know which nullifiers the block will actually create, we will
+        //     supply a superset of all possible nullifiers the block might create, which are the
+        //     nullifiers of all proven batches. However, if we knew that a certain note will be
+        //     erased, we would not have to supply a nullifier witness for it.
+
+        let batch_iter = batches.iter();
+
+        let unauthenticated_notes_iter = batch_iter.clone().flat_map(|batch| {
+            batch.input_notes().iter().filter_map(|note| note.header().map(NoteHeader::id))
+        });
+        let block_references_iter = batch_iter.clone().map(ProvenBatch::reference_block_num);
+        let account_ids = batch_iter.clone().flat_map(ProvenBatch::updated_accounts);
+        let created_nullifiers = batch_iter.flat_map(ProvenBatch::produced_nullifiers);
 
         let inputs = self
             .store
             .get_block_inputs(
-                summary.updated_accounts.iter().copied(),
-                summary.nullifiers.iter(),
-                summary.dangling_notes.iter(),
+                account_ids,
+                created_nullifiers,
+                unauthenticated_notes_iter,
+                block_references_iter,
             )
             .await
             .map_err(BuildBlockError::GetBlockInputsFailed)?;
 
-        let missing_notes: Vec<_> = summary
-            .dangling_notes
-            .difference(&inputs.found_unauthenticated_notes.note_ids())
-            .copied()
-            .collect();
-        if !missing_notes.is_empty() {
-            return Err(BuildBlockError::UnauthenticatedNotesNotFound(missing_notes));
-        }
-
-        Ok(BlockSummaryAndInputs { batches, summary, inputs })
+        // No longer needed, we have note erasure at block level.
+        // let missing_notes: Vec<_> = summary
+        //     .dangling_notes
+        //     .difference(&inputs.found_unauthenticated_notes.note_ids())
+        //     .copied()
+        //     .collect();
+        // if !missing_notes.is_empty() {
+        //     return Err(BuildBlockError::UnauthenticatedNotesNotFound(missing_notes));
+        // }
+
+        Ok(BlockSummaryAndInputs { batches, inputs })
     }
 
     #[instrument(target = COMPONENT, name = "block_builder.prove_block", skip_all, err)]
     async fn prove_block(
         &self,
         preimage: BlockSummaryAndInputs,
     ) -> Result<ProvenBlockWrapper, BuildBlockError> {
-        let BlockSummaryAndInputs { batches, summary, inputs } = preimage;
+        let BlockSummaryAndInputs { batches, inputs } = preimage;
 
-        let (block_header_witness, updated_accounts) = BlockWitness::new(inputs, &batches)?;
+        let proposed_block = ProposedBlock::new(inputs, batches).expect("TODO: error");
 
-        let new_block_header = self.block_kernel.prove(block_header_witness)?;
-
-        // TODO: Update. Temporarily left in an incorrect state.
-        let block = ProvenBlock::new_unchecked(
-            new_block_header,
-            updated_accounts,
-            vec![],
-            summary.nullifiers,
-        );
+        let proven_block = LocalBlockProver::new(MIN_PROOF_SECURITY_LEVEL)
+            .prove(proposed_block)
+            .expect("TODO: error");
 
         self.simulate_proving().await;
 
-        Ok(ProvenBlockWrapper { block })
+        Ok(ProvenBlockWrapper { block: proven_block })
     }
 
     #[instrument(target = COMPONENT, name = "block_builder.commit_block", skip_all, err)]
@@ -273,7 +289,6 @@ struct SelectedBlock {
 }
 struct BlockSummaryAndInputs {
     batches: Vec<ProvenBatch>,
-    summary: BlockSummary,
     inputs: BlockInputs,
 }
 
@@ -297,21 +312,17 @@ impl BlockSummaryAndInputs {
         // SAFETY: We do not expect to have more than u32::MAX of any count per block.
         span.set_attribute(
             "block.updated_accounts.count",
-            i64::try_from(self.summary.updated_accounts.len())
+            i64::try_from(self.inputs.account_witnesses().len())
                 .expect("less than u32::MAX account updates"),
         );
-        span.set_attribute(
-            "block.output_notes.count",
-            i64::try_from(self.summary.output_notes.iter().fold(0, |acc, x| acc.add(x.len())))
-                .expect("less than u32::MAX output notes"),
-        );
         span.set_attribute(
             "block.nullifiers.count",
-            i64::try_from(self.summary.nullifiers.len()).expect("less than u32::MAX nullifiers"),
+            i64::try_from(self.inputs.nullifier_witnesses().len())
+                .expect("less than u32::MAX nullifiers"),
         );
         span.set_attribute(
-            "block.dangling_notes.count",
-            i64::try_from(self.summary.dangling_notes.len())
+            "block.unauthenticated_notes.count",
+            i64::try_from(self.inputs.unauthenticated_note_proofs().len())
                 .expect("less than u32::MAX dangling notes"),
         );
     }
@@ -328,6 +339,13 @@ impl ProvenBlockWrapper {
 
         span.set_attribute("block.protocol.version", i64::from(header.version()));
 
+        // Question: Should this be here?
+        span.set_attribute(
+            "block.output_notes.count",
+            i64::try_from(self.block.output_notes().count())
+                .expect("less than u32::MAX output notes"),
+        );
+
         span.set_attribute("block.commitments.kernel", header.kernel_root());
         span.set_attribute("block.commitments.nullifier", header.nullifier_root());
         span.set_attribute("block.commitments.account", header.account_root());

crates/block-producer/src/lib.rs

@@ -10,7 +10,6 @@ mod errors;
 mod mempool;
 mod store;
 
-pub mod block;
 pub mod config;
 pub mod server;