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

triton_vm/
proof_stream.rs
