triton_vm/
proof_stream.rs

1use arbitrary::Arbitrary;
2use twenty_first::prelude::*;
3
4use crate::error::ProofStreamError;
5use crate::proof::Proof;
6use crate::proof_item::ProofItem;
7
8#[derive(Debug, Default, Clone, Eq, PartialEq, BFieldCodec, Arbitrary)]
9pub struct ProofStream {
10    pub items: Vec<ProofItem>,
11
12    #[bfield_codec(ignore)]
13    pub items_index: usize,
14
15    #[bfield_codec(ignore)]
16    pub sponge: Tip5,
17}
18
19impl ProofStream {
20    pub fn new() -> Self {
21        ProofStream {
22            items: vec![],
23            items_index: 0,
24            sponge: Tip5::init(),
25        }
26    }
27
28    /// The number of field elements required to encode the proof.
29    pub fn transcript_length(&self) -> usize {
30        let Proof(b_field_elements) = self.into();
31        b_field_elements.len()
32    }
33
34    /// Alters the Fiat-Shamir's sponge state with the encoding of the given item.
35    /// Does _not_ record the given item in the proof stream.
36    /// This is useful for items that are not sent to the verifier, _e.g._, the
37    /// [`Claim`](crate::proof::Claim).
38    ///
39    /// See also [`Self::enqueue()`] and [`Self::dequeue()`].
40    pub fn alter_fiat_shamir_state_with(&mut self, item: &impl BFieldCodec) {
41        self.sponge.pad_and_absorb_all(&item.encode())
42    }
43
44    /// Send a proof item as prover to verifier.
45    /// Some items do not need to be included in the Fiat-Shamir heuristic, _i.e._, they do not
46    /// need to modify the sponge state. For those items, namely those that evaluate to `false`
47    /// according to [`ProofItem::include_in_fiat_shamir_heuristic`], the sponge state is not
48    /// modified.
49    /// For example:
50    /// - Merkle authentication structure do not need to be hashed if the root of the tree
51    ///     in question was hashed previously.
52    /// - If the proof stream is not used to sample any more randomness, _i.e._, after the last
53    ///     round of interaction, no further items need to be hashed.
54    pub fn enqueue(&mut self, item: ProofItem) {
55        if item.include_in_fiat_shamir_heuristic() {
56            self.alter_fiat_shamir_state_with(&item);
57        }
58        self.items.push(item);
59    }
60
61    /// Receive a proof item from prover as verifier.
62    /// See [`ProofStream::enqueue`] for more details.
63    pub fn dequeue(&mut self) -> Result<ProofItem, ProofStreamError> {
64        let Some(item) = self.items.get(self.items_index) else {
65            return Err(ProofStreamError::EmptyQueue);
66        };
67        let item = item.to_owned();
68        if item.include_in_fiat_shamir_heuristic() {
69            self.alter_fiat_shamir_state_with(&item);
70        }
71        self.items_index += 1;
72        Ok(item)
73    }
74
75    /// Given an `upper_bound` that is a power of 2, produce `num_indices` uniform random numbers
76    /// in the interval `[0; upper_bound)`.
77    ///
78    /// - `upper_bound`: The (non-inclusive) upper bound. Must be a power of two.
79    /// - `num_indices`: The number of indices to sample
80    pub fn sample_indices(&mut self, upper_bound: usize, num_indices: usize) -> Vec<usize> {
81        assert!(upper_bound.is_power_of_two());
82        assert!(upper_bound <= BFieldElement::MAX as usize);
83        self.sponge
84            .sample_indices(upper_bound as u32, num_indices)
85            .into_iter()
86            .map(|i| i as usize)
87            .collect()
88    }
89
90    /// A thin wrapper around [`Tip5::sample_scalars`].
91    pub fn sample_scalars(&mut self, num_scalars: usize) -> Vec<XFieldElement> {
92        self.sponge.sample_scalars(num_scalars)
93    }
94}
95
96impl TryFrom<&Proof> for ProofStream {
97    type Error = ProofStreamError;
98
99    fn try_from(proof: &Proof) -> Result<Self, ProofStreamError> {
100        let proof_stream = *ProofStream::decode(&proof.0)?;
101        Ok(proof_stream)
102    }
103}
104
105impl From<&ProofStream> for Proof {
106    fn from(proof_stream: &ProofStream) -> Self {
107        Proof(proof_stream.encode())
108    }
109}
110
111impl From<ProofStream> for Proof {
112    fn from(proof_stream: ProofStream) -> Self {
113        (&proof_stream).into()
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use std::collections::VecDeque;
120
121    use assert2::assert;
122    use assert2::let_assert;
123    use itertools::Itertools;
124    use proptest::collection::vec;
125    use proptest_arbitrary_interop::arb;
126    use test_strategy::proptest;
127    use twenty_first::math::other::random_elements;
128
129    use crate::proof_item::FriResponse;
130    use crate::proof_item::ProofItem;
131    use crate::shared_tests::LeavedMerkleTreeTestData;
132    use crate::table::AuxiliaryRow;
133    use crate::table::MainRow;
134    use crate::table::QuotientSegments;
135
136    use super::*;
137
138    #[proptest]
139    fn serialize_proof_with_fiat_shamir(
140        #[strategy(vec(arb(), 2..100))] main_rows: Vec<MainRow<BFieldElement>>,
141        #[strategy(vec(arb(), 2..100))] aux_rows: Vec<AuxiliaryRow>,
142        #[strategy(arb())] ood_main_row: Box<MainRow<XFieldElement>>,
143        #[strategy(arb())] ood_aux_row: Box<AuxiliaryRow>,
144        #[strategy(arb())] quot_elements: Vec<QuotientSegments>,
145        leaved_merkle_tree: LeavedMerkleTreeTestData,
146    ) {
147        let auth_structure = leaved_merkle_tree.auth_structure.clone();
148        let root = leaved_merkle_tree.root();
149        let fri_codeword = leaved_merkle_tree.leaves().to_owned();
150        let fri_response = leaved_merkle_tree.into_fri_response();
151
152        let mut sponge_states = VecDeque::new();
153        let mut proof_stream = ProofStream::new();
154
155        sponge_states.push_back(proof_stream.sponge.state);
156        proof_stream.enqueue(ProofItem::AuthenticationStructure(auth_structure.clone()));
157        sponge_states.push_back(proof_stream.sponge.state);
158        proof_stream.enqueue(ProofItem::MasterMainTableRows(main_rows.clone()));
159        sponge_states.push_back(proof_stream.sponge.state);
160        proof_stream.enqueue(ProofItem::MasterAuxTableRows(aux_rows.clone()));
161        sponge_states.push_back(proof_stream.sponge.state);
162        proof_stream.enqueue(ProofItem::OutOfDomainMainRow(ood_main_row.clone()));
163        sponge_states.push_back(proof_stream.sponge.state);
164        proof_stream.enqueue(ProofItem::OutOfDomainAuxRow(ood_aux_row.clone()));
165        sponge_states.push_back(proof_stream.sponge.state);
166        proof_stream.enqueue(ProofItem::MerkleRoot(root));
167        sponge_states.push_back(proof_stream.sponge.state);
168        proof_stream.enqueue(ProofItem::QuotientSegmentsElements(quot_elements.clone()));
169        sponge_states.push_back(proof_stream.sponge.state);
170        proof_stream.enqueue(ProofItem::FriCodeword(fri_codeword.clone()));
171        sponge_states.push_back(proof_stream.sponge.state);
172        proof_stream.enqueue(ProofItem::FriResponse(fri_response.clone()));
173        sponge_states.push_back(proof_stream.sponge.state);
174
175        let proof = proof_stream.into();
176        let mut proof_stream = ProofStream::try_from(&proof).unwrap();
177
178        assert!(sponge_states.pop_front() == Some(proof_stream.sponge.state));
179        let_assert!(Ok(proof_item) = proof_stream.dequeue());
180        let_assert!(ProofItem::AuthenticationStructure(auth_structure_) = proof_item);
181        assert!(auth_structure == auth_structure_);
182
183        assert!(sponge_states.pop_front() == Some(proof_stream.sponge.state));
184        let_assert!(Ok(ProofItem::MasterMainTableRows(main_rows_)) = proof_stream.dequeue());
185        assert!(main_rows == main_rows_);
186
187        assert!(sponge_states.pop_front() == Some(proof_stream.sponge.state));
188        let_assert!(Ok(ProofItem::MasterAuxTableRows(aux_rows_)) = proof_stream.dequeue());
189        assert!(aux_rows == aux_rows_);
190
191        assert!(sponge_states.pop_front() == Some(proof_stream.sponge.state));
192        let_assert!(Ok(ProofItem::OutOfDomainMainRow(ood_main_row_)) = proof_stream.dequeue());
193        assert!(ood_main_row == ood_main_row_);
194
195        assert!(sponge_states.pop_front() == Some(proof_stream.sponge.state));
196        let_assert!(Ok(ProofItem::OutOfDomainAuxRow(ood_aux_row_)) = proof_stream.dequeue());
197        assert!(ood_aux_row == ood_aux_row_);
198
199        assert!(sponge_states.pop_front() == Some(proof_stream.sponge.state));
200        let_assert!(Ok(ProofItem::MerkleRoot(root_)) = proof_stream.dequeue());
201        assert!(root == root_);
202
203        assert!(sponge_states.pop_front() == Some(proof_stream.sponge.state));
204        let_assert!(Ok(proof_item) = proof_stream.dequeue());
205        let_assert!(ProofItem::QuotientSegmentsElements(quot_elements_) = proof_item);
206        assert!(quot_elements == quot_elements_);
207
208        assert!(sponge_states.pop_front() == Some(proof_stream.sponge.state));
209        let_assert!(Ok(ProofItem::FriCodeword(fri_codeword_)) = proof_stream.dequeue());
210        assert!(fri_codeword == fri_codeword_);
211
212        assert!(sponge_states.pop_front() == Some(proof_stream.sponge.state));
213        let_assert!(Ok(ProofItem::FriResponse(fri_response_)) = proof_stream.dequeue());
214        assert!(fri_response == fri_response_);
215
216        assert!(sponge_states.pop_front() == Some(proof_stream.sponge.state));
217        assert!(0 == sponge_states.len());
218    }
219
220    #[test]
221    fn enqueue_dequeue_verify_partial_authentication_structure() {
222        let tree_height = 8;
223        let num_leaves = 1 << tree_height;
224        let leaf_values: Vec<XFieldElement> = random_elements(num_leaves);
225        let leaf_digests = leaf_values.iter().map(|&xfe| xfe.into()).collect_vec();
226        let merkle_tree = MerkleTree::par_new(&leaf_digests).unwrap();
227        let indices_to_check = vec![5, 173, 175, 167, 228, 140, 252, 149, 232, 182, 5, 5, 182];
228        let auth_structure = merkle_tree
229            .authentication_structure(&indices_to_check)
230            .unwrap();
231        let revealed_leaves = indices_to_check
232            .iter()
233            .map(|&idx| leaf_values[idx])
234            .collect_vec();
235        let fri_response = FriResponse {
236            auth_structure,
237            revealed_leaves,
238        };
239
240        let mut proof_stream = ProofStream::new();
241        proof_stream.enqueue(ProofItem::FriResponse(fri_response));
242
243        // TODO: Also check that deserializing from Proof works here.
244
245        let proof_item = proof_stream.dequeue().unwrap();
246        let maybe_same_fri_response = proof_item.try_into_fri_response().unwrap();
247        let FriResponse {
248            auth_structure,
249            revealed_leaves,
250        } = maybe_same_fri_response;
251        let maybe_same_leaf_digests = revealed_leaves.iter().map(|&xfe| xfe.into()).collect_vec();
252        let indexed_leafs = indices_to_check
253            .into_iter()
254            .zip_eq(maybe_same_leaf_digests)
255            .collect();
256
257        let inclusion_proof = MerkleTreeInclusionProof {
258            tree_height,
259            indexed_leafs,
260            authentication_structure: auth_structure,
261        };
262        assert!(inclusion_proof.verify(merkle_tree.root()));
263    }
264
265    #[test]
266    fn dequeuing_from_empty_stream_fails() {
267        let mut proof_stream = ProofStream::new();
268        let_assert!(Err(ProofStreamError::EmptyQueue) = proof_stream.dequeue());
269    }
270
271    #[test]
272    fn dequeuing_more_items_than_have_been_enqueued_fails() {
273        let mut proof_stream = ProofStream::new();
274        proof_stream.enqueue(ProofItem::FriCodeword(vec![]));
275        proof_stream.enqueue(ProofItem::Log2PaddedHeight(7));
276
277        let_assert!(Ok(_) = proof_stream.dequeue());
278        let_assert!(Ok(_) = proof_stream.dequeue());
279        let_assert!(Err(ProofStreamError::EmptyQueue) = proof_stream.dequeue());
280    }
281
282    #[test]
283    fn encoded_length_of_prove_stream_is_not_known_at_compile_time() {
284        assert!(ProofStream::static_length().is_none());
285    }
286}
triton_vm/proof_stream.rs

triton_vm/
proof_stream.rs
