tasm_lib/verifier/fri/
test_helpers.rs

1use itertools::Itertools;
2use triton_vm::challenges::Challenges;
3use triton_vm::prelude::*;
4use triton_vm::proof_stream::ProofStream;
5use triton_vm::table::NUM_QUOTIENT_SEGMENTS;
6use triton_vm::table::master_table::MasterAuxTable;
7use triton_vm::table::master_table::MasterMainTable;
8use twenty_first::prelude::*;
9
10use crate::prelude::Digest;
11
12pub struct StarkProofExtraction {
13    pub fri_proof_stream: ProofStream,
14    pub main_tree_authentication_paths: Vec<Vec<Digest>>,
15    pub aux_tree_authentication_paths: Vec<Vec<Digest>>,
16    pub quot_tree_authentication_paths: Vec<Vec<Digest>>,
17}
18
19/// Extracts a proof stream that will work for FRI verification from a proof stream that works for
20/// the whole STARK verification.
21pub fn extract_fri_proof(
22    proof_stream: &ProofStream,
23    claim: &Claim,
24    stark: &Stark,
25) -> StarkProofExtraction {
26    let mut proof_stream = proof_stream.to_owned();
27    let log2_padded_height = proof_stream
28        .dequeue()
29        .unwrap()
30        .try_into_log2_padded_height()
31        .unwrap();
32    proof_stream.alter_fiat_shamir_state_with(claim);
33
34    // Main-table Merkle root
35    let main_table_root = proof_stream
36        .dequeue()
37        .unwrap()
38        .try_into_merkle_root()
39        .unwrap();
40
41    // Aux challenge weights
42    proof_stream.sample_scalars(Challenges::SAMPLE_COUNT);
43
44    // Aux-table Merkle root
45    let aux_mt_root = proof_stream
46        .dequeue()
47        .unwrap()
48        .try_into_merkle_root()
49        .unwrap();
50
51    // Quotient codeword weights
52    proof_stream.sample_scalars(MasterAuxTable::NUM_CONSTRAINTS);
53
54    // Quotient codeword Merkle root
55    let quotient_root = proof_stream
56        .dequeue()
57        .unwrap()
58        .try_into_merkle_root()
59        .unwrap();
60
61    // Out-of-domain point current row
62    proof_stream.sample_scalars(1);
63
64    // Five out-of-domain values
65    proof_stream
66        .dequeue()
67        .unwrap()
68        .try_into_out_of_domain_main_row()
69        .unwrap();
70    proof_stream
71        .dequeue()
72        .unwrap()
73        .try_into_out_of_domain_aux_row()
74        .unwrap();
75    proof_stream
76        .dequeue()
77        .unwrap()
78        .try_into_out_of_domain_main_row()
79        .unwrap();
80    proof_stream
81        .dequeue()
82        .unwrap()
83        .try_into_out_of_domain_aux_row()
84        .unwrap();
85    proof_stream
86        .dequeue()
87        .unwrap()
88        .try_into_out_of_domain_quot_segments()
89        .unwrap();
90
91    // `beqd_weights`
92    const NUM_DEEP_CODEWORD_COMPONENTS: usize = 3;
93    proof_stream.sample_scalars(
94        MasterMainTable::NUM_COLUMNS
95            + MasterAuxTable::NUM_COLUMNS
96            + NUM_QUOTIENT_SEGMENTS
97            + NUM_DEEP_CODEWORD_COMPONENTS,
98    );
99
100    let padded_height = 1 << log2_padded_height;
101    let fri = stark.fri(padded_height).unwrap();
102    let fri_proof_stream = proof_stream.clone();
103    let fri_verify_result = fri.verify(&mut proof_stream).unwrap();
104    let indices = fri_verify_result.iter().map(|(i, _)| *i).collect_vec();
105    let tree_height = fri.domain.length.ilog2();
106
107    // main
108    let main_table_rows = proof_stream
109        .dequeue()
110        .unwrap()
111        .try_into_master_main_table_rows()
112        .unwrap();
113    let main_authentication_structure = proof_stream
114        .dequeue()
115        .unwrap()
116        .try_into_authentication_structure()
117        .unwrap();
118    let main_tree_authentication_paths = extract_paths(
119        main_table_root,
120        &indices,
121        &main_table_rows,
122        &main_authentication_structure,
123        tree_height,
124    );
125
126    // auxiliary
127    let aux_table_rows = proof_stream
128        .dequeue()
129        .unwrap()
130        .try_into_master_aux_table_rows()
131        .unwrap();
132    let aux_authentication_structure = proof_stream
133        .dequeue()
134        .unwrap()
135        .try_into_authentication_structure()
136        .unwrap();
137    let aux_tree_authentication_paths = extract_paths(
138        aux_mt_root,
139        &indices,
140        &aux_table_rows,
141        &aux_authentication_structure,
142        tree_height,
143    );
144
145    // quotient
146    let quot_table_rows = proof_stream
147        .dequeue()
148        .unwrap()
149        .try_into_quot_segments_elements()
150        .unwrap();
151    let quot_authentication_structure = proof_stream
152        .dequeue()
153        .unwrap()
154        .try_into_authentication_structure()
155        .unwrap();
156    let quot_tree_authentication_paths = extract_paths(
157        quotient_root,
158        &indices,
159        &quot_table_rows,
160        &quot_authentication_structure,
161        tree_height,
162    );
163
164    StarkProofExtraction {
165        fri_proof_stream,
166        main_tree_authentication_paths,
167        aux_tree_authentication_paths,
168        quot_tree_authentication_paths,
169    }
170}
171
172fn extract_paths<const N: usize, T: BFieldCodec>(
173    root: Digest,
174    indices: &[usize],
175    rows: &[[T; N]],
176    authentication_structure: &[Digest],
177    tree_height: u32,
178) -> Vec<Vec<Digest>> {
179    let leafs = rows.iter().map(Tip5::hash).collect_vec();
180    let inclusion_proof = MerkleTreeInclusionProof {
181        tree_height,
182        indexed_leafs: indices.iter().cloned().zip(leafs).collect_vec(),
183        authentication_structure: authentication_structure.to_vec(),
184    };
185    assert!(inclusion_proof.clone().verify(root));
186    inclusion_proof.into_authentication_paths().unwrap()
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    #[test]
194    fn extract_fri_proof_works() {
195        let simple_program = triton_program!(halt);
196        let public_input = [];
197        let non_determinism = NonDeterminism::default();
198        let (stark, claim, proof) =
199            triton_vm::prove_program(simple_program, public_input.into(), non_determinism).unwrap();
200        let padded_height = proof.padded_height().unwrap();
201        let fri = stark.fri(padded_height).unwrap();
202
203        let proof_stream = ProofStream::try_from(&proof).unwrap();
204        let mut fri_proof_stream =
205            extract_fri_proof(&proof_stream, &claim, &stark).fri_proof_stream;
206        assert!(
207            fri.verify(&mut fri_proof_stream).is_ok(),
208            "Extracted proof must verify"
209        );
210    }
211}