Feat/streaming prover

crypto/crypto/src/merkle_tree/merkle.rs

@@ -55,6 +55,11 @@ pub(crate) struct MmapNodeBacking {
 pub struct MerkleTree<B: IsMerkleTreeBackend> {
     pub root: B::Node,
     nodes: Vec<B::Node>,
+    /// Number of leaves in the tree (always a power of two). Stored explicitly
+    /// so the leaf count is still known after [`MerkleTree::drop_leaves`] frees
+    /// the leaf half (after which `nodes.len() == leaves_len - 1`, so the usual
+    /// `(node_count + 1) / 2` recovery no longer applies).
+    leaves_len: usize,
     #[cfg(feature = "disk-spill")]
     #[cfg_attr(feature = "serde", serde(skip))]
     mmap_backing: Option<MmapNodeBacking>,
@@ -79,13 +84,14 @@ where
 {
     fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
         use serde::ser::SerializeStruct;
-        let mut s = serializer.serialize_struct("MerkleTree", 2)?;
+        let mut s = serializer.serialize_struct("MerkleTree", 3)?;
         s.serialize_field("root", &self.root)?;
         if self.mmap_backing.is_some() {
             s.serialize_field("nodes", &MmapNodesSeq(self))?;
         } else {
             s.serialize_field("nodes", &self.nodes)?;
         }
+        s.serialize_field("leaves_len", &self.leaves_len)?;
         s.end()
     }
 }
@@ -154,6 +160,7 @@ where
         Some(MerkleTree {
             root: nodes[ROOT].clone(),
             nodes,
+            leaves_len,
             #[cfg(feature = "disk-spill")]
             mmap_backing: None,
         })
@@ -211,6 +218,93 @@ where
         self.create_proof(merkle_path)
     }
 
+    /// Number of leaves in the tree (always a power of two). Valid both before
+    /// and after [`drop_leaves`](Self::drop_leaves).
+    pub fn leaves_len(&self) -> usize {
+        self.leaves_len
+    }
+
+    /// Whether the leaf half of the node buffer has been freed by
+    /// [`drop_leaves`](Self::drop_leaves). When `true`, the leaf-level sibling
+    /// must be supplied by the caller to build an opening (see
+    /// [`get_proof_by_pos_with_leaf_sibling`](Self::get_proof_by_pos_with_leaf_sibling)).
+    pub fn leaves_dropped(&self) -> bool {
+        self.node_count() == self.leaves_len - 1
+    }
+
+    /// Free the leaf half of the node buffer, keeping only the inner nodes
+    /// (`nodes[0..leaves_len - 1]`, root at index 0). This roughly halves the
+    /// tree's memory footprint. The root and every inner node are retained, so
+    /// the only path node that must be regenerated at open time is the
+    /// leaf-level sibling — see
+    /// [`get_proof_by_pos_with_leaf_sibling`](Self::get_proof_by_pos_with_leaf_sibling).
+    ///
+    /// Idempotent: a no-op if the leaves were already dropped. A single-leaf
+    /// tree (`leaves_len == 1`) has no leaf half to drop and is left unchanged.
+    pub fn drop_leaves(&mut self) {
+        if self.leaves_len <= 1 {
+            return;
+        }
+        let inner_count = self.leaves_len - 1;
+        // `disk-spill` mmap backing is read-only and never populated together
+        // with leaf-dropping in the prover, so only the heap path is handled.
+        #[cfg(feature = "disk-spill")]
+        if self.mmap_backing.is_some() {
+            return;
+        }
+        if self.nodes.len() > inner_count {
+            self.nodes.truncate(inner_count);
+            self.nodes.shrink_to_fit();
+        }
+    }
+
+    /// Builds the same opening that [`get_proof_by_pos`](Self::get_proof_by_pos)
+    /// would, but sources the leaf-level sibling from `leaf_sibling` instead of
+    /// the (possibly dropped) leaf buffer. All higher levels read the retained
+    /// inner nodes. The resulting [`Proof`] is byte-identical to the full-tree
+    /// `get_proof_by_pos(pos)` provided `leaf_sibling` equals the hash the
+    /// builder stored for the leaf at `sibling_leaf_position(pos)`.
+    ///
+    /// Works whether or not the leaves have been dropped.
+    pub fn get_proof_by_pos_with_leaf_sibling(
+        &self,
+        pos: usize,
+        leaf_sibling: B::Node,
+    ) -> Option<Proof<B::Node>> {
+        let leaf_node = pos + (self.leaves_len - 1);
+        // Single-leaf tree: the leaf is the root, the path is empty.
+        if leaf_node == ROOT {
+            return Some(Proof {
+                merkle_path: Vec::new(),
+            });
+        }
+
+        let tree_depth = self.leaves_len.ilog2() as usize;
+        let mut merkle_path = Vec::with_capacity(tree_depth);
+
+        // Bottom level: the leaf-level sibling, supplied by the caller (its
+        // node lives in the dropped leaf half).
+        merkle_path.push(leaf_sibling);
+
+        // Higher levels: every sibling here is an inner node, still resident.
+        let mut node = parent_index(leaf_node);
+        while node != ROOT {
+            let sibling = self.node_get(sibling_index(node))?;
+            merkle_path.push(sibling.clone());
+            node = parent_index(node);
+        }
+
+        Some(Proof { merkle_path })
+    }
+
+    /// 0-based leaf position of the leaf-level sibling needed to open `pos`.
+    /// The caller regenerates that leaf's hash (e.g. by re-hashing the LDE row)
+    /// to pass into [`get_proof_by_pos_with_leaf_sibling`](Self::get_proof_by_pos_with_leaf_sibling).
+    pub fn sibling_leaf_position(&self, pos: usize) -> usize {
+        let leaf_node = pos + (self.leaves_len - 1);
+        sibling_index(leaf_node) - (self.leaves_len - 1)
+    }
+
     /// Creates a proof from a Merkle pasth
     fn create_proof(&self, merkle_path: Vec<B::Node>) -> Option<Proof<B::Node>> {
         Some(Proof { merkle_path })

crypto/crypto/src/tests/merkle_proof_tests.rs

@@ -333,3 +333,42 @@ fn batch_proof_verify_sparse_leaves_across_tree() {
         16
     ));
 }
+
+use crate::merkle_tree::traits::IsMerkleTreeBackend;
+use crate::merkle_tree::utils::complete_until_power_of_two;
+
+/// Leaf-drop opener must produce byte-identical proofs to the full-tree opener.
+/// This is the core correctness invariant for the streaming "leaf-drop" mode.
+#[test]
+fn leaf_dropped_opening_is_byte_identical_to_full_tree_opening() {
+    type B = TestBackend<Ecgfp5>;
+    let values: Vec<Ecgfp5FE> = (1..1000).map(Ecgfp5FE::new).collect();
+    let leaves = complete_until_power_of_two(<B as IsMerkleTreeBackend>::hash_leaves(&values));
+
+    let full = TestMerkleTreeEcgfp::build(&values).unwrap();
+    let mut dropped = TestMerkleTreeEcgfp::build(&values).unwrap();
+    dropped.drop_leaves();
+    assert!(dropped.leaves_dropped());
+    assert_eq!(dropped.leaves_len(), leaves.len());
+
+    for pos in [0usize, 1, 2, 7, 9349 % leaves.len(), leaves.len() - 1] {
+        let full_proof = full.get_proof_by_pos(pos).unwrap();
+        let sib_leaf = leaves[dropped.sibling_leaf_position(pos)];
+        let dropped_proof = dropped
+            .get_proof_by_pos_with_leaf_sibling(pos, sib_leaf)
+            .unwrap();
+        assert_eq!(
+            full_proof.merkle_path, dropped_proof.merkle_path,
+            "leaf-dropped proof differs from full-tree proof at pos {pos}"
+        );
+    }
+}
+
+#[test]
+fn drop_leaves_is_idempotent() {
+    let values: Vec<Ecgfp5FE> = (1..100).map(Ecgfp5FE::new).collect();
+    let mut tree = TestMerkleTreeEcgfp::build(&values).unwrap();
+    tree.drop_leaves();
+    tree.drop_leaves();
+    assert!(tree.leaves_dropped());
+}

crypto/stark/src/instruments.rs

@@ -86,6 +86,8 @@ thread_local! {
     static R2_SUB: RefCell<Option<(Duration, Duration, Duration)>> = const { RefCell::new(None) };
     /// Round 4 sub-timings: (fft, merkle, deep_comp, queries)
     static R4_SUB: RefCell<Option<(Duration, Duration, Duration, Duration)>> = const { RefCell::new(None) };
+    /// Round 3 OOD evaluation timing.
+    static R3_OOD: RefCell<Option<Duration>> = const { RefCell::new(None) };
     /// Assembled sub-ops from prove_rounds_2_to_4 (without reconstruct_round1 LDE time).
     static ROUND_SUB_OPS: RefCell<Option<TableSubOps>> = const { RefCell::new(None) };
 }
@@ -141,6 +143,9 @@ pub fn reset_all() {
     R4_SUB.with(|cell| {
         cell.borrow_mut().take();
     });
+    R3_OOD.with(|cell| {
+        cell.borrow_mut().take();
+    });
     ROUND_SUB_OPS.with(|cell| {
         cell.borrow_mut().take();
     });
@@ -154,6 +159,14 @@ pub fn take_r2_sub() -> Option<(Duration, Duration, Duration)> {
     R2_SUB.with(|cell| cell.borrow_mut().take())
 }
 
+pub fn store_r3_ood(d: Duration) {
+    R3_OOD.with(|cell| *cell.borrow_mut() = Some(d));
+}
+
+pub fn take_r3_ood() -> Option<Duration> {
+    R3_OOD.with(|cell| cell.borrow_mut().take())
+}
+
 pub fn store_r4_sub(fft: Duration, merkle: Duration, deep_comp: Duration, queries: Duration) {
     R4_SUB.with(|cell| *cell.borrow_mut() = Some((fft, merkle, deep_comp, queries)));
 }

crypto/stark/src/proof/stark.rs

@@ -30,6 +30,40 @@ pub struct DeepPolynomialOpening<F: IsSubFieldOf<E>, E: IsField> {
 
 pub type DeepPolynomialOpenings<F, E> = Vec<DeepPolynomialOpening<F, E>>;
 
+/// Per-(chunk, lde_size) batched FRI instance (Approach 1, batched FRI within a
+/// chunk).
+///
+/// One per height bucket inside a chunk: every bucket-mate's individual DEEP
+/// composition polynomial is linearly combined with successive powers of the
+/// bucket's `delta_fri` challenge (sampled from the chunk-shared `bucket_seed`),
+/// and a single FRI commit + grinding + query is run on the combined
+/// polynomial. The `members` list pins the canonical bucket-local order used to
+/// derive `delta_fri^i` on the verifier side; reordering the list rejects the
+/// proof.
+///
+/// `decommitments` length equals `air.options().fri_number_of_queries` (one
+/// decommitment per shared iota). `nonce` is `Some` when the AIR's grinding
+/// factor > 0 (`None` otherwise).
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[serde(bound = "")]
+pub struct ChunkBucketFri<E: IsField> {
+    /// LDE size shared by every bucket-mate. Equal to
+    /// `trace_length * blowup_factor` for each member.
+    pub lde_size: u32,
+    /// Chunk-local indices of the bucket-mates, in canonical (chunk-local
+    /// index ascending) order. Index `i` here corresponds to `delta_fri^i`
+    /// in the linear combination.
+    pub members: Vec<usize>,
+    /// `[pₖ]` for the committed FRI layers.
+    pub layer_roots: Vec<Commitment>,
+    /// `pₙ` — the final folded constant.
+    pub last_value: FieldElement<E>,
+    /// One FRI decommitment per shared iota.
+    pub decommitments: Vec<FriDecommitment<E>>,
+    /// Grinding nonce, when `grinding_factor > 0`.
+    pub nonce: Option<u64>,
+}
+
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 #[serde(bound = "PI: serde::Serialize + serde::de::DeserializeOwned")]
 pub struct StarkProof<F: IsSubFieldOf<E>, E: IsField, PI> {
@@ -50,17 +84,14 @@ pub struct StarkProof<F: IsSubFieldOf<E>, E: IsField, PI> {
     pub composition_poly_root: Commitment,
     // Hᵢ(z^N)
     pub composition_poly_parts_ood_evaluation: Vec<FieldElement<E>>,
-    // [pₖ]
-    pub fri_layers_merkle_roots: Vec<Commitment>,
-    // pₙ
-    pub fri_last_value: FieldElement<E>,
-    // Open(pₖ(Dₖ), −𝜐ₛ^(2ᵏ))
-    pub query_list: Vec<FriDecommitment<E>>,
     // Open(H₁(D_LDE, 𝜐ᵢ), Open(H₂(D_LDE, 𝜐ᵢ), Open(tⱼ(D_LDE), 𝜐ᵢ)
     // Open(H₁(D_LDE, -𝜐ᵢ), Open(H₂(D_LDE, -𝜐ᵢ), Open(tⱼ(D_LDE), -𝜐ᵢ)
+    //
+    // FRI for this table is no longer per-table: it is run once per
+    // (chunk, lde_size) bucket and lives in
+    // [`MultiProof::fri_chunk_buckets`]. These DEEP openings are evaluated
+    // at the bucket-shared query indices (iotas).
     pub deep_poly_openings: DeepPolynomialOpenings<F, E>,
-    // nonce obtained from grinding
-    pub nonce: Option<u64>,
     // Bus interaction public inputs for the accumulated column.
     // Contains the table contribution (L), used for:
     // 1. Circular constraint offset: L/N per row
@@ -77,4 +108,12 @@ pub struct StarkProof<F: IsSubFieldOf<E>, E: IsField, PI> {
 #[serde(bound = "PI: serde::Serialize + serde::de::DeserializeOwned")]
 pub struct MultiProof<F: IsSubFieldOf<E>, E: IsField, PI> {
     pub proofs: Vec<StarkProof<F, E, PI>>,
+    /// Per-(chunk, lde_size-bucket) batched FRI instances. Outer Vec is indexed
+    /// by chunk (chunks of `chunk_size` tables in proof order); inner Vec lists
+    /// buckets in canonical first-encounter (chunk-local-index ascending) order.
+    pub fri_chunk_buckets: Vec<Vec<ChunkBucketFri<E>>>,
+    /// Pinned chunk size (= the prover's `table_parallelism()` at proving time).
+    /// The verifier uses this to chunk the proof slice into the same per-chunk
+    /// grouping the prover used.
+    pub chunk_size: u32,
 }