1use arbitrary::Arbitrary;
2use strum::Display;
3use strum::EnumCount;
4use strum::EnumDiscriminants;
5use strum::EnumIter;
6use twenty_first::prelude::*;
7
8use crate::error::ProofStreamError;
9use crate::error::ProofStreamError::UnexpectedItem;
10use crate::fri::AuthenticationStructure;
11use crate::table::AuxiliaryRow;
12use crate::table::MainRow;
13use crate::table::QuotientSegments;
14
15#[derive(Debug, Clone, Eq, PartialEq, Hash, BFieldCodec, Arbitrary)]
19pub struct FriResponse {
20 pub auth_structure: AuthenticationStructure,
22
23 pub revealed_leaves: Vec<XFieldElement>,
25}
26
27macro_rules! proof_items {
28 ($($variant:ident($payload:ty) => $in_fiat_shamir_heuristic:literal, $try_into_fn:ident,)+) => {
29 #[derive(
30 Debug,
31 Display,
32 Clone,
33 Eq,
34 PartialEq,
35 Hash,
36 EnumCount,
37 EnumDiscriminants,
38 BFieldCodec,
39 Arbitrary,
40 )]
41 #[strum_discriminants(name(ProofItemVariant))]
42 #[strum_discriminants(derive(Display, EnumIter, BFieldCodec, Arbitrary))]
44 pub enum ProofItem {
45 $( $variant($payload), )+
46 }
47
48 impl ProofItem {
49 pub const fn include_in_fiat_shamir_heuristic(&self) -> bool {
62 match self {
63 $( Self::$variant(_) => $in_fiat_shamir_heuristic, )+
64 }
65 }
66
67 $(
68 pub fn $try_into_fn(self) -> Result<$payload, ProofStreamError> {
69 match self {
70 Self::$variant(payload) => Ok(payload),
71 _ => Err(UnexpectedItem {
72 expected: ProofItemVariant::$variant,
73 got: self,
74 }),
75 }
76 }
77 )+
78 }
79
80 impl ProofItemVariant {
81 pub fn payload_static_length(self) -> Option<usize> {
82 match self {
83 $( Self::$variant => <$payload>::static_length(), )+
84 }
85 }
86
87 pub const fn include_in_fiat_shamir_heuristic(self) -> bool {
89 match self {
90 $( Self::$variant => $in_fiat_shamir_heuristic, )+
91 }
92 }
93
94 pub const fn payload_type(self) -> &'static str {
96 match self {
97 $( Self::$variant => stringify!($payload), )+
98 }
99 }
100 }
101 };
102}
103
104proof_items!(
105 MerkleRoot(Digest) => true, try_into_merkle_root,
106 OutOfDomainMainRow(Box<MainRow<XFieldElement>>) => true, try_into_out_of_domain_main_row,
107 OutOfDomainAuxRow(Box<AuxiliaryRow>) => true, try_into_out_of_domain_aux_row,
108 OutOfDomainQuotientSegments(QuotientSegments) => true, try_into_out_of_domain_quot_segments,
109
110 AuthenticationStructure(AuthenticationStructure) => false, try_into_authentication_structure,
112 MasterMainTableRows(Vec<MainRow<BFieldElement>>) => false, try_into_master_main_table_rows,
113 MasterAuxTableRows(Vec<AuxiliaryRow>) => false, try_into_master_aux_table_rows,
114 Log2PaddedHeight(u32) => false, try_into_log2_padded_height,
115 QuotientSegmentsElements(Vec<QuotientSegments>) => false, try_into_quot_segments_elements,
116 FriCodeword(Vec<XFieldElement>) => false, try_into_fri_codeword,
117 FriPolynomial(Polynomial<'static, XFieldElement>) => false, try_into_fri_polynomial,
118 FriResponse(FriResponse) => false, try_into_fri_response,
119);
120
121#[cfg(test)]
122#[cfg_attr(coverage_nightly, coverage(off))]
123pub(crate) mod tests {
124 use std::collections::HashSet;
125
126 use assert2::assert;
127 use assert2::let_assert;
128 use proptest::prelude::*;
129 use strum::IntoEnumIterator;
130 use test_strategy::proptest;
131
132 use crate::proof::Proof;
133 use crate::proof_stream::ProofStream;
134 use crate::shared_tests::LeavedMerkleTreeTestData;
135
136 use super::*;
137
138 #[proptest]
139 fn serialize_fri_response_in_isolation(leaved_merkle_tree: LeavedMerkleTreeTestData) {
140 let fri_response = leaved_merkle_tree.into_fri_response();
141 let encoding = fri_response.encode();
142 let_assert!(Ok(decoding) = FriResponse::decode(&encoding));
143 prop_assert_eq!(fri_response, *decoding);
144 }
145
146 #[proptest]
147 fn serialize_fri_response_in_proof_stream(leaved_merkle_tree: LeavedMerkleTreeTestData) {
148 let fri_response = leaved_merkle_tree.into_fri_response();
149 let mut proof_stream = ProofStream::new();
150 proof_stream.enqueue(ProofItem::FriResponse(fri_response.clone()));
151 let proof: Proof = proof_stream.into();
152
153 let_assert!(Ok(mut proof_stream) = ProofStream::try_from(&proof));
154 let_assert!(Ok(proof_item) = proof_stream.dequeue());
155 let_assert!(Ok(fri_response_) = proof_item.try_into_fri_response());
156 prop_assert_eq!(fri_response, fri_response_);
157 }
158
159 #[proptest]
160 fn serialize_authentication_structure_in_isolation(
161 leaved_merkle_tree: LeavedMerkleTreeTestData,
162 ) {
163 let auth_structure = leaved_merkle_tree.auth_structure;
164 let encoding = auth_structure.encode();
165 let_assert!(Ok(decoding) = AuthenticationStructure::decode(&encoding));
166 prop_assert_eq!(auth_structure, *decoding);
167 }
168
169 #[proptest]
170 fn serialize_authentication_structure_in_proof_stream(
171 leaved_merkle_tree: LeavedMerkleTreeTestData,
172 ) {
173 let auth_structure = leaved_merkle_tree.auth_structure;
174 let mut proof_stream = ProofStream::new();
175 proof_stream.enqueue(ProofItem::AuthenticationStructure(auth_structure.clone()));
176 let proof: Proof = proof_stream.into();
177
178 let_assert!(Ok(mut proof_stream) = ProofStream::try_from(&proof));
179 let_assert!(Ok(proof_item) = proof_stream.dequeue());
180 let_assert!(Ok(auth_structure_) = proof_item.try_into_authentication_structure());
181 prop_assert_eq!(auth_structure, auth_structure_);
182 }
183
184 #[test]
185 fn interpreting_a_merkle_root_as_anything_else_gives_appropriate_error() {
186 let fake_root = Digest::default();
187 let item = ProofItem::MerkleRoot(fake_root);
188 assert!(let Err(UnexpectedItem{..}) = item.clone().try_into_authentication_structure());
189 assert!(let Err(UnexpectedItem{..}) = item.clone().try_into_fri_response());
190 assert!(let Err(UnexpectedItem{..}) = item.clone().try_into_master_main_table_rows());
191 assert!(let Err(UnexpectedItem{..}) = item.clone().try_into_master_aux_table_rows());
192 assert!(let Err(UnexpectedItem{..}) = item.clone().try_into_out_of_domain_main_row());
193 assert!(let Err(UnexpectedItem{..}) = item.clone().try_into_out_of_domain_aux_row());
194 assert!(let Err(UnexpectedItem{..}) = item.clone().try_into_out_of_domain_quot_segments());
195 assert!(let Err(UnexpectedItem{..}) = item.clone().try_into_log2_padded_height());
196 assert!(let Err(UnexpectedItem{..}) = item.clone().try_into_quot_segments_elements());
197 assert!(let Err(UnexpectedItem{..}) = item.clone().try_into_fri_codeword());
198 assert!(let Err(UnexpectedItem{..}) = item.try_into_fri_polynomial());
199 }
200
201 #[test]
202 fn proof_item_payload_static_length_is_as_expected() {
203 assert!(let Some(_) = ProofItemVariant::MerkleRoot.payload_static_length());
204 assert!(let Some(_) = ProofItemVariant::Log2PaddedHeight.payload_static_length());
205 assert_eq!(None, ProofItemVariant::FriCodeword.payload_static_length());
206 assert_eq!(None, ProofItemVariant::FriResponse.payload_static_length());
207 }
208
209 #[test]
210 fn can_loop_over_proof_item_variants() {
211 let all_discriminants: HashSet<_> = ProofItemVariant::iter()
212 .map(|variant| variant.bfield_codec_discriminant())
213 .collect();
214 assert_eq!(ProofItem::COUNT, all_discriminants.len());
215 }
216
217 #[test]
218 fn proof_item_and_its_variant_have_same_bfield_codec_discriminant() {
219 assert_eq!(
220 ProofItem::MerkleRoot(Digest::default()).bfield_codec_discriminant(),
221 ProofItemVariant::MerkleRoot.bfield_codec_discriminant()
222 );
223 assert_eq!(
224 ProofItem::Log2PaddedHeight(0).bfield_codec_discriminant(),
225 ProofItemVariant::Log2PaddedHeight.bfield_codec_discriminant()
226 );
227 assert_eq!(
228 ProofItem::FriCodeword(vec![]).bfield_codec_discriminant(),
229 ProofItemVariant::FriCodeword.bfield_codec_discriminant()
230 );
231 }
232
233 #[test]
234 fn proof_item_variants_payload_type_has_expected_format() {
235 assert_eq!("Digest", ProofItemVariant::MerkleRoot.payload_type());
236 }
237}