Skip to main content

tasm_lib/verifier/vm_proof_iter/
dequeue_next_as.rs

1use triton_vm::fri::AuthenticationStructure;
2use triton_vm::prelude::*;
3use triton_vm::proof_item::FriResponse;
4use triton_vm::proof_item::ProofItemVariant;
5use triton_vm::table::AuxiliaryRow;
6use triton_vm::table::MainRow;
7use triton_vm::table::QuotientSegments;
8use twenty_first::math::x_field_element::EXTENSION_DEGREE;
9use twenty_first::prelude::Polynomial;
10
11use crate::hashing::sponge_hasher::pad_and_absorb_all::PadAndAbsorbAll;
12use crate::prelude::*;
13
14const MAX_SIZE_FOR_DYNAMICALLY_SIZED_PROOF_ITEMS: u32 = 1u32 << 22;
15
16/// Reads a proof item of the supplied type from the [`ProofStream`][proof_stream].
17/// Crashes Triton VM if the proof item is not of the expected type.
18/// Updates an internal pointer to the next proof item.
19///
20/// [proof_stream]: triton_vm::proof_stream::ProofStream
21#[derive(Debug, Copy, Clone, Eq, PartialEq)]
22pub struct DequeueNextAs {
23    pub proof_item: ProofItemVariant,
24}
25
26impl DequeueNextAs {
27    pub fn new(proof_item: ProofItemVariant) -> Self {
28        Self { proof_item }
29    }
30
31    fn item_name(&self) -> String {
32        self.proof_item.to_string().to_lowercase()
33    }
34
35    fn proof_item_calculate_size(&self, library: &mut Library) -> Vec<LabelledInstruction> {
36        match self.proof_item {
37            ProofItemVariant::MasterMainTableRows => {
38                Vec::<MainRow<BFieldElement>>::compute_size_and_assert_valid_size_indicator(library)
39            }
40            ProofItemVariant::MasterAuxTableRows => {
41                Vec::<AuxiliaryRow>::compute_size_and_assert_valid_size_indicator(library)
42            }
43            ProofItemVariant::QuotientSegmentsElements => {
44                Vec::<QuotientSegments>::compute_size_and_assert_valid_size_indicator(library)
45            }
46            ProofItemVariant::FriCodeword => {
47                Vec::<XFieldElement>::compute_size_and_assert_valid_size_indicator(library)
48            }
49            ProofItemVariant::FriPolynomial => {
50                Polynomial::<XFieldElement>::compute_size_and_assert_valid_size_indicator(library)
51            }
52            ProofItemVariant::FriResponse => {
53                FriResponse::compute_size_and_assert_valid_size_indicator(library)
54            }
55            ProofItemVariant::AuthenticationStructure => {
56                // All authentication structures should be stripped from loaded proofs, but for
57                // more complete tests, we support it here anyway.
58                AuthenticationStructure::compute_size_and_assert_valid_size_indicator(library)
59            }
60            _ => unreachable!(),
61        }
62    }
63
64    /// Increment the counter of proof items already dequeued.
65    ///
66    /// ```text
67    /// BEFORE: _ *vm_proof_iter
68    /// AFTER:  _ *vm_proof_iter
69    /// ```
70    fn increment_current_item_count(&self) -> Vec<LabelledInstruction> {
71        const CURRENT_ITEM_COUNT_OFFSET: BFieldElement = BFieldElement::new(4);
72
73        triton_asm!(
74            // _ *proof_item_iter
75
76            dup 0
77            addi {CURRENT_ITEM_COUNT_OFFSET}
78            // _ *proof_item_iter *current_item_count
79
80            read_mem 1
81            addi 1
82            // _ *proof_item_iter current_item_count *current_item_count
83
84            pick 1
85            addi 1
86            pick 1
87            // _ *proof_item_iter (current_item_count + 1) *current_item_count
88
89            write_mem 1
90            pop 1
91            // _ *proof_item_iter
92        )
93    }
94
95    /// ```text
96    /// BEFORE: _ *proof_item_size
97    /// AFTER:  _ *proof_item_size
98    /// ```
99    fn assert_consistent_size_indicators(&self, library: &mut Library) -> Vec<LabelledInstruction> {
100        if self.proof_item.payload_static_length().is_some() {
101            return vec![];
102        }
103
104        let calculate_own_size = self.proof_item_calculate_size(library);
105
106        triton_asm! {
107            // _ *proof_item_size
108
109            dup 0
110            addi 3
111            {&calculate_own_size}
112            hint calculated_payload_size = stack[0]
113            // _ *proof_item_size calculated_payload_size
114
115            dup 1
116            addi 2
117            read_mem 1
118            pop 1
119            hint indicated_payload_size = stack[0]
120            // _ *proof_item_size calculated_payload_size reported_field_size
121
122            dup 1
123            eq
124            assert
125            hint payload_size = stack[0]
126            // _ *proof_item_size payload_size
127
128
129            /* Account for discriminant and size-indicator */
130            addi 2
131            hint calculated_item_size = stack[0]
132            // _ *proof_item_size calculated_item_size
133
134            dup 1
135            read_mem 1
136            pop 1
137            // _ *proof_item_size calculated_own_size indicated_item_size
138
139            eq
140            assert
141            // _ *proof_item_size
142        }
143    }
144
145    /// ```text
146    /// BEFORE: _ *proof_item_size
147    /// AFTER:  _ *proof_item_size
148    /// ```
149    fn fiat_shamir_absorb_snippet(&self, library: &mut Library) -> Vec<LabelledInstruction> {
150        if !self.proof_item.include_in_fiat_shamir_heuristic() {
151            return vec![];
152        }
153
154        let pad_and_absorb_all = library.import(Box::new(PadAndAbsorbAll));
155        triton_asm! {
156            dup 0           // _ *proof_item_size *proof_item_size
157            call {pad_and_absorb_all}
158                            // _ *proof_item_size
159        }
160    }
161
162    /// ```text
163    /// BEFORE: _ *proof_item_iter *proof_item_size
164    /// AFTER:  _ *proof_item_size
165    /// ```
166    fn update_proof_item_iter_to_next_proof_item(&self) -> Vec<LabelledInstruction> {
167        triton_asm! {
168            dup 0               // _ *proof_item_iter *proof_item_size *proof_item_iter
169            {&self.advance_proof_item_pointer_to_next_item()}
170            hint next_proof_item_list_element_size_pointer = stack[0]
171                                // _ *proof_item_iter *proof_item_size *next_proof_item_size
172
173            swap 1 swap 2       // _ *pi_size *next_pi_size *pi_iter
174            write_mem 1 pop 1   // _ *pi_size
175        }
176    }
177
178    /// ```text
179    /// BEFORE: _ *proof_item_size
180    /// AFTER:  _ *next_proof_item_size
181    /// ```
182    fn advance_proof_item_pointer_to_next_item(&self) -> Vec<LabelledInstruction> {
183        let Some(static_length) = self.proof_item.payload_static_length() else {
184            return self.advance_proof_item_pointer_to_next_item_reading_length_from_memory();
185        };
186        Self::advance_proof_item_pointer_to_next_item_using_static_size(static_length)
187    }
188
189    /// ```text
190    /// BEFORE: _ *proof_item_size
191    /// AFTER:  _ *next_proof_item_size
192    /// ```
193    fn advance_proof_item_pointer_to_next_item_using_static_size(
194        length: usize,
195    ) -> Vec<LabelledInstruction> {
196        let length_plus_bookkeeping_offset = length + 2;
197        triton_asm! {
198            push {length_plus_bookkeeping_offset}
199            hint proof_item_length_plus_bookkeeping_offset = stack[0]
200            add
201        }
202    }
203
204    /// ```text
205    /// BEFORE: _ *proof_item_size
206    /// AFTER:  _ *next_proof_item_size
207    /// ```
208    fn advance_proof_item_pointer_to_next_item_reading_length_from_memory(
209        &self,
210    ) -> Vec<LabelledInstruction> {
211        triton_asm! {
212            read_mem 1          // _ proof_item_size (*proof_item_size - 1)
213            hint proof_item_length = stack[1]
214
215            // Verify that size does not exceed max allowed size
216            push {MAX_SIZE_FOR_DYNAMICALLY_SIZED_PROOF_ITEMS}
217            dup 2
218            // _ proof_item_size (*proof_item_size - 1) max_size proof_item_size
219
220            lt
221            // _ proof_item_size (*proof_item_size - 1) (max_size > proof_item_size)
222
223            assert
224            // _ proof_item_size (*proof_item_size - 1)
225
226            push 2 add add      // _ *next_proof_item_size
227        }
228    }
229
230    /// ```text
231    /// BEFORE: _ *proof_item_size
232    /// AFTER:  _ *proof_item_payload
233    /// ```
234    fn advance_list_element_pointer_to_proof_item_payload(&self) -> Vec<LabelledInstruction> {
235        let payload_length_indicator_size = match self.proof_item.payload_static_length() {
236            Some(_) => 0,
237            None => 1,
238        };
239        let list_item_length_indicator_size = 1;
240        let discriminant_size = 1;
241        let bookkeeping_offset =
242            payload_length_indicator_size + list_item_length_indicator_size + discriminant_size;
243
244        triton_asm! { push {bookkeeping_offset} hint bookkeeping_offset = stack[0] add }
245    }
246
247    /// A [`BFieldCodec`] encoding of a [`Polynomial`] may not contain trailing
248    /// zeros. This is because it must be easy to read the degree of the polynomial.
249    /// In order to be consistent with that rule, we perform this check here, before
250    /// returning a pointer to the polynomial to the caller.
251    ///
252    /// ```text
253    /// BEFORE: _ *proof_item_payload
254    /// AFTER:  _ *proof_item_payload
255    /// ```
256    fn verify_last_xfe_is_non_zero_if_payload_is_polynomial(&self) -> Vec<LabelledInstruction> {
257        match self.proof_item {
258            ProofItemVariant::FriPolynomial => triton_asm!(
259                // _ *fri_polynomial
260
261                dup 0
262                push 1
263                add
264                read_mem 1
265                pop 1
266                // _ *fri_polynomial num_coefficients
267
268                push {EXTENSION_DEGREE}
269                mul
270                push 1
271                add
272                // _ *fri_polynomial offset_last_word
273
274                dup 1
275                add
276                // _ *proof_item_payload *coefficients_last_word
277
278                read_mem {EXTENSION_DEGREE}
279                pop 1
280                push 0
281                eq
282                swap 2
283                push 0
284                eq
285                swap 1
286                push 0
287                eq
288                add
289                add
290                push {EXTENSION_DEGREE}
291                eq
292                // _ *proof_item_payload (last_coefficient == 0)
293
294                push 0
295                eq
296                // _ *proof_item_payload (last_coefficient != 0)
297
298                assert
299            ),
300            _ => triton_asm!(),
301        }
302    }
303}
304
305impl BasicSnippet for DequeueNextAs {
306    fn parameters(&self) -> Vec<(DataType, String)> {
307        vec![(DataType::VoidPointer, "*proof_item_iter".to_string())]
308    }
309
310    fn return_values(&self) -> Vec<(DataType, String)> {
311        let payload_pointer_str = format!("*{}_payload", self.item_name());
312        vec![(DataType::VoidPointer, payload_pointer_str)]
313    }
314
315    fn entrypoint(&self) -> String {
316        let proof_item_name = self.item_name();
317        format!("tasmlib_verifier_vm_proof_iter_dequeue_next_as_{proof_item_name}")
318    }
319
320    /// ```text
321    /// BEFORE: _ *vm_proof_iter
322    /// AFTER:  _ *proof_item_payload
323    /// ```
324    fn code(&self, library: &mut Library) -> Vec<LabelledInstruction> {
325        let final_hint = format!("hint {}_pointer: Pointer = stack[0]", self.item_name());
326        triton_asm! {
327            {self.entrypoint()}:
328                {&self.increment_current_item_count()}
329                                    // _ *vm_proof_iter
330
331                read_mem 1
332                hint proof_item_list_element_size_pointer = stack[1]
333
334                addi 1
335                swap 1              // _ *proof_item_iter *proof_item_size
336
337                addi 1
338                hint proof_item_discriminant_pointer = stack[0]
339
340                read_mem 1          // _ *proof_item_iter discriminant *proof_item_size
341                hint proof_item_list_element_size_pointer = stack[0]
342                hint proof_item_discriminant = stack[1]
343
344                swap 1
345                push {self.proof_item.bfield_codec_discriminant()}
346                hint expected_proof_item_discriminant = stack[0]
347
348                eq assert           // _ *proof_item_iter *proof_item_size
349
350                {&self.assert_consistent_size_indicators(library)}
351                                    // _ *proof_item_iter *proof_item_size
352
353                {&self.fiat_shamir_absorb_snippet(library)}
354                                    // _ *proof_item_iter *proof_item_size
355
356                {&self.update_proof_item_iter_to_next_proof_item()}
357                                    // _ *proof_item_size
358
359                {&self.advance_list_element_pointer_to_proof_item_payload()}
360                {final_hint}        // _ *proof_item_payload
361
362                {&self.verify_last_xfe_is_non_zero_if_payload_is_polynomial()}
363                {final_hint}        // _ *proof_item_payload
364
365                return
366        }
367    }
368}
369
370#[cfg(test)]
371mod tests {
372    use arbitrary::Arbitrary;
373    use arbitrary::Unstructured;
374    use num_traits::One;
375    use num_traits::Zero;
376    use strum::IntoEnumIterator;
377    use triton_vm::proof_item::ProofItem;
378    use triton_vm::proof_stream::ProofStream;
379    use triton_vm::table::master_table::MasterMainTable;
380    use twenty_first::prelude::*;
381
382    use super::*;
383    use crate::empty_stack;
384    use crate::execute_with_terminal_state;
385    use crate::rust_shadowing_helper_functions::dyn_malloc::dynamic_allocator;
386    use crate::structure::tasm_object::decode_from_memory_with_size;
387    use crate::test_prelude::*;
388    use crate::verifier::vm_proof_iter::shared::vm_proof_iter_struct::VmProofIter;
389
390    #[derive(Debug)]
391    struct RustShadowForDequeueNextAs<'a> {
392        dequeue_next_as: DequeueNextAs,
393
394        proof_iter_pointer: BFieldElement,
395
396        stack: &'a mut Vec<BFieldElement>,
397        memory: &'a mut HashMap<BFieldElement, BFieldElement>,
398        sponge: &'a mut Tip5,
399    }
400
401    impl RustShadowForDequeueNextAs<'_> {
402        fn execute(&mut self) -> Result<Vec<BFieldElement>, RustShadowError> {
403            self.assert_correct_discriminant()?;
404            self.maybe_alter_fiat_shamir_heuristic_with_proof_item();
405            self.update_stack_using_current_proof_item();
406            self.update_proof_item_iter();
407            self.update_current_item_count();
408
409            Ok(self.public_output())
410        }
411
412        fn current_proof_item_list_element_size_pointer(&self) -> BFieldElement {
413            let &maybe_pointer = self.memory.get(&self.proof_iter_pointer).unwrap();
414            maybe_pointer
415        }
416
417        fn current_proof_item_list_element_size(&self) -> BFieldElement {
418            let Some(item_length) = self.dequeue_next_as.proof_item.payload_static_length() else {
419                return self.fetch_current_proof_item_list_element_size_from_memory();
420            };
421            let discriminant_length = 1;
422            BFieldElement::new((item_length + discriminant_length) as u64)
423        }
424
425        /// Accounts for the item's discriminant.
426        fn fetch_current_proof_item_list_element_size_from_memory(&self) -> BFieldElement {
427            let size_pointer = self.current_proof_item_list_element_size_pointer();
428            let &element_size = self.memory.get(&size_pointer).unwrap();
429            assert!(
430                element_size.value() < MAX_SIZE_FOR_DYNAMICALLY_SIZED_PROOF_ITEMS as u64,
431                "proof item size may not exceed max allowed size"
432            );
433            element_size
434        }
435
436        fn next_proof_item_list_element_size_pointer(&self) -> BFieldElement {
437            self.current_proof_item_list_element_size_pointer()
438                + self.current_proof_item_list_element_size()
439                + BFieldElement::one()
440        }
441
442        fn discriminant_pointer(&self) -> BFieldElement {
443            self.current_proof_item_list_element_size_pointer() + BFieldElement::one()
444        }
445
446        fn discriminant(&self) -> BFieldElement {
447            let &discriminant = self.memory.get(&self.discriminant_pointer()).unwrap();
448            discriminant
449        }
450
451        fn proof_item_payload_pointer(&self) -> BFieldElement {
452            self.discriminant_pointer()
453                + BFieldElement::one()
454                + self.proof_item_payload_size_indicator_length()
455        }
456
457        fn proof_item_payload_size_indicator_length(&self) -> BFieldElement {
458            match self.dequeue_next_as.proof_item.payload_static_length() {
459                Some(_) => BFieldElement::zero(),
460                None => BFieldElement::one(),
461            }
462        }
463
464        fn assert_correct_discriminant(&self) -> Result<(), RustShadowError> {
465            let expected_discriminant = self.dequeue_next_as.proof_item.bfield_codec_discriminant();
466            if expected_discriminant != self.discriminant().value() as usize {
467                return Err(RustShadowError::Other);
468            };
469
470            Ok(())
471        }
472
473        fn proof_item(&self) -> ProofItem {
474            let maybe_item = decode_from_memory_with_size::<ProofItem>(
475                self.memory,
476                self.discriminant_pointer(),
477                self.current_proof_item_list_element_size().value() as usize,
478            );
479            *maybe_item.unwrap()
480        }
481
482        fn maybe_alter_fiat_shamir_heuristic_with_proof_item(&mut self) {
483            let proof_item = self.proof_item();
484            if proof_item.include_in_fiat_shamir_heuristic() {
485                Tip5::pad_and_absorb_all(self.sponge, &proof_item.encode());
486            }
487        }
488
489        fn update_stack_using_current_proof_item(&mut self) {
490            self.stack.push(self.proof_item_payload_pointer());
491        }
492
493        fn update_proof_item_iter(&mut self) {
494            let next_item_pointer = self.next_proof_item_list_element_size_pointer();
495            self.memory
496                .insert(self.proof_iter_pointer, next_item_pointer);
497        }
498
499        fn update_current_item_count(&mut self) {
500            let count = self
501                .memory
502                .get_mut(&(self.proof_iter_pointer + bfe!(4)))
503                .unwrap();
504            count.increment();
505        }
506
507        fn public_output(&self) -> Vec<BFieldElement> {
508            vec![]
509        }
510    }
511
512    impl Procedure for DequeueNextAs {
513        fn rust_shadow(
514            &self,
515            stack: &mut Vec<BFieldElement>,
516            memory: &mut HashMap<BFieldElement, BFieldElement>,
517            _: &NonDeterminism,
518            _: &[BFieldElement],
519            sponge: &mut Option<Tip5>,
520        ) -> Result<Vec<BFieldElement>, RustShadowError> {
521            let Some(sponge) = sponge.as_mut() else {
522                return Err(RustShadowError::SpongeUninitialized);
523            };
524
525            RustShadowForDequeueNextAs {
526                dequeue_next_as: *self,
527                proof_iter_pointer: stack.pop().ok_or(RustShadowError::StackUnderflow)?,
528                stack,
529                memory,
530                sponge,
531            }
532            .execute()
533        }
534
535        fn pseudorandom_initial_state(
536            &self,
537            seed: [u8; 32],
538            _: Option<BenchmarkCase>,
539        ) -> ProcedureInitialState {
540            let mut rng = StdRng::from_seed(seed);
541            let mut proof_stream = ProofStream::new();
542            proof_stream.enqueue(Self::pseudorandom_proof_item(self.proof_item, rng.random()));
543
544            let other_item_type = ProofItemVariant::iter().choose(&mut rng).unwrap();
545            proof_stream.enqueue(Self::pseudorandom_proof_item(other_item_type, rng.random()));
546
547            self.initial_state_from_proof_and_address(proof_stream.into(), rng.random())
548        }
549    }
550
551    impl DequeueNextAs {
552        fn initial_state_from_proof_and_address(
553            &self,
554            proof: Proof,
555            proof_address: BFieldElement,
556        ) -> ProcedureInitialState {
557            let mut ram = HashMap::new();
558            encode_to_memory(&mut ram, proof_address, &proof);
559
560            let proof_iter_address = dynamic_allocator(&mut ram);
561            let vm_proof_iter_init_state = VmProofIter::new(proof_address, &proof);
562            encode_to_memory(&mut ram, proof_iter_address, &vm_proof_iter_init_state);
563
564            ProcedureInitialState {
565                stack: [self.init_stack_for_isolated_run(), vec![proof_iter_address]].concat(),
566                nondeterminism: NonDeterminism::default().with_ram(ram),
567                public_input: vec![],
568                sponge: Some(Tip5::init()),
569            }
570        }
571
572        fn pseudorandom_new(seed: [u8; 32]) -> Self {
573            let mut unstructured = Unstructured::new(&seed);
574            let proof_item: ProofItem = Arbitrary::arbitrary(&mut unstructured).unwrap();
575            let proof_item = proof_item.into();
576            Self { proof_item }
577        }
578
579        pub(crate) fn pseudorandom_proof_stream(
580            proof_items_variants: Vec<ProofItemVariant>,
581            seed: [u8; 32],
582        ) -> ProofStream {
583            let mut rng = StdRng::from_seed(seed);
584
585            let mut proof_stream = ProofStream::new();
586            for &proof_item in &proof_items_variants {
587                let item = DequeueNextAs::pseudorandom_proof_item(proof_item, rng.random());
588                proof_stream.enqueue(item);
589            }
590
591            proof_stream
592        }
593
594        fn pseudorandom_proof_item(
595            proof_item_variant: ProofItemVariant,
596            seed: [u8; 32],
597        ) -> ProofItem {
598            let mut rng = StdRng::from_seed(seed);
599            let proof_stream_seed: [u8; 10000] = rng.random();
600            let mut unstructured = Unstructured::new(&proof_stream_seed);
601
602            use ProofItemVariant::*;
603            match proof_item_variant {
604                MerkleRoot => {
605                    ProofItem::MerkleRoot(Arbitrary::arbitrary(&mut unstructured).unwrap())
606                }
607                OutOfDomainMainRow => {
608                    ProofItem::OutOfDomainMainRow(Arbitrary::arbitrary(&mut unstructured).unwrap())
609                }
610                OutOfDomainAuxRow => {
611                    ProofItem::OutOfDomainAuxRow(Arbitrary::arbitrary(&mut unstructured).unwrap())
612                }
613                OutOfDomainQuotientSegments => ProofItem::OutOfDomainQuotientSegments(
614                    Arbitrary::arbitrary(&mut unstructured).unwrap(),
615                ),
616                AuthenticationStructure => ProofItem::AuthenticationStructure(
617                    Arbitrary::arbitrary(&mut unstructured).unwrap(),
618                ),
619                MasterMainTableRows => {
620                    ProofItem::MasterMainTableRows(Arbitrary::arbitrary(&mut unstructured).unwrap())
621                }
622                MasterAuxTableRows => {
623                    ProofItem::MasterAuxTableRows(Arbitrary::arbitrary(&mut unstructured).unwrap())
624                }
625                Log2PaddedHeight => {
626                    ProofItem::Log2PaddedHeight(Arbitrary::arbitrary(&mut unstructured).unwrap())
627                }
628                QuotientSegmentsElements => ProofItem::QuotientSegmentsElements(
629                    Arbitrary::arbitrary(&mut unstructured).unwrap(),
630                ),
631                FriCodeword => {
632                    ProofItem::FriCodeword(Arbitrary::arbitrary(&mut unstructured).unwrap())
633                }
634                FriPolynomial => {
635                    ProofItem::FriPolynomial(Arbitrary::arbitrary(&mut unstructured).unwrap())
636                }
637                FriResponse => {
638                    // This code requires the authentication paths field to be empty
639                    let fri_response = triton_vm::proof_item::FriResponse {
640                        auth_structure: vec![],
641                        revealed_leaves: Arbitrary::arbitrary(&mut unstructured).unwrap(),
642                    };
643                    ProofItem::FriResponse(fri_response)
644                }
645            }
646        }
647
648        fn test_rust_equivalence(self, initial_state: ProcedureInitialState) {
649            test_rust_equivalence_given_complete_state(
650                &ShadowedProcedure::new(self),
651                &initial_state.stack,
652                &initial_state.public_input,
653                &initial_state.nondeterminism,
654                &initial_state.sponge,
655                None,
656            );
657        }
658    }
659
660    #[macro_rules_attr::apply(test)]
661    fn disallow_too_big_dynamically_sized_proof_item() {
662        let dequeue_next_as = DequeueNextAs::new(ProofItemVariant::MasterMainTableRows);
663        let initial_state = initial_state_with_too_big_master_table_rows();
664        let tvm_result = execute_with_terminal_state(
665            Program::new(&dequeue_next_as.link_for_isolated_run()),
666            &[],
667            &initial_state.stack,
668            &initial_state.nondeterminism,
669            None,
670        );
671
672        let mut stack = initial_state.stack.clone();
673        let mut memory = initial_state.nondeterminism.ram.clone();
674        let rust_result = dequeue_next_as.rust_shadow(
675            &mut stack,
676            &mut memory,
677            &NonDeterminism::default(),
678            &[],
679            &mut None,
680        );
681
682        assert!(
683            rust_result.is_err() && tvm_result.is_err(),
684            "Test case: Too big dynamically-sized proof item must fail on both platforms."
685        );
686        let err = tvm_result.unwrap_err();
687        assert!(matches!(err, InstructionError::AssertionFailed(_)));
688    }
689
690    fn initial_state_with_too_big_master_table_rows() -> ProcedureInitialState {
691        let dummy_master_table_rows = vec![[bfe!(101); MasterMainTable::NUM_COLUMNS]; 15000];
692        let proof_item = ProofItem::MasterMainTableRows(dummy_master_table_rows);
693        let mut proof_stream = ProofStream::new();
694        proof_stream.enqueue(proof_item);
695        let proof_ptr = BFieldElement::zero();
696        DequeueNextAs::new(ProofItemVariant::MasterMainTableRows)
697            .initial_state_from_proof_and_address(proof_stream.into(), proof_ptr)
698    }
699
700    #[macro_rules_attr::apply(test)]
701    fn disallow_trailing_zeros_in_xfe_poly_encoding() {
702        let dequeue_next_as = DequeueNextAs::new(ProofItemVariant::FriPolynomial);
703        let initial_state = initial_state_with_trailing_zeros_in_xfe_poly_encoding();
704        let code = dequeue_next_as.link_for_isolated_run();
705        let tvm_result = execute_with_terminal_state(
706            Program::new(&code),
707            &[],
708            &initial_state.stack,
709            &initial_state.nondeterminism,
710            None,
711        );
712
713        let mut stack = initial_state.stack.clone();
714        let mut memory = initial_state.nondeterminism.ram.clone();
715        let rust_result = dequeue_next_as.rust_shadow(
716            &mut stack,
717            &mut memory,
718            &NonDeterminism::default(),
719            &[],
720            &mut None,
721        );
722
723        assert!(
724            rust_result.is_err() && tvm_result.is_err(),
725            "Test case: Too big dynamically-sized proof item must fail on both platforms."
726        );
727        let err = tvm_result.unwrap_err();
728        assert!(matches!(err, InstructionError::AssertionFailed(_)));
729    }
730
731    fn initial_state_with_trailing_zeros_in_xfe_poly_encoding() -> ProcedureInitialState {
732        let mut encoded_fri_poly_proof_item =
733            ProofItem::FriPolynomial(Polynomial::new(vec![xfe!(1), xfe!(2), xfe!(3), xfe!(100)]))
734                .encode();
735
736        // Manually set the leading coefficient to zero
737        let encoding_length = encoded_fri_poly_proof_item.len();
738        for i in 0..EXTENSION_DEGREE {
739            encoded_fri_poly_proof_item[encoding_length - 1 - i] = bfe!(0);
740        }
741
742        let proof_start_address = BFieldElement::zero();
743
744        let mut ram = HashMap::new();
745
746        let mut address = proof_start_address;
747
748        // Inidicate a field-size of `ProofStream`
749        ram.insert(address, bfe!(encoding_length as u64 + 2));
750        address.increment();
751
752        // Indicate that there is one proof item
753        ram.insert(address, bfe!(1));
754        address.increment();
755
756        // Indicate size of encoded proof item
757        ram.insert(address, bfe!(encoding_length as u64));
758        address.increment();
759
760        // Add proof item to ND memory
761        for word in encoded_fri_poly_proof_item {
762            ram.insert(address, word);
763            address.increment();
764        }
765
766        // uses highly specific knowledge of `BFieldCodec`
767        let address_of_first_element = proof_start_address + BFieldElement::new(2);
768        let proof_iter_address = proof_start_address - BFieldElement::one();
769        ram.insert(proof_iter_address, address_of_first_element);
770        ProcedureInitialState {
771            stack: [empty_stack(), vec![proof_iter_address]].concat(),
772            nondeterminism: NonDeterminism::default().with_ram(ram),
773            public_input: vec![],
774            sponge: Some(Tip5::init()),
775        }
776    }
777
778    fn dequeueing_is_equivalent_in_rust_and_tasm_prop(proof_item_variant: ProofItemVariant) {
779        let mut proof_stream = ProofStream::new();
780        let proof_item = DequeueNextAs::pseudorandom_proof_item(proof_item_variant, rand::random());
781        proof_stream.enqueue(proof_item);
782
783        let dequeue_next_as = DequeueNextAs::new(proof_item_variant);
784        let proof_ptr = BFieldElement::new(14);
785        let init_state =
786            dequeue_next_as.initial_state_from_proof_and_address(proof_stream.into(), proof_ptr);
787        dequeue_next_as.test_rust_equivalence(init_state);
788    }
789
790    #[macro_rules_attr::apply(test)]
791    fn dequeueing_all_proof_items_individually_is_equivalent_in_rust_and_tasm() {
792        for variant in ProofItemVariant::iter() {
793            dequeueing_is_equivalent_in_rust_and_tasm_prop(variant);
794        }
795    }
796
797    #[macro_rules_attr::apply(proptest)]
798    fn dequeuing_as_suggested_element_is_equivalent_in_rust_and_tasm(seed: [u8; 32]) {
799        let dequeue_next_as = DequeueNextAs::pseudorandom_new(seed);
800        let initial_state = dequeue_next_as.pseudorandom_initial_state(seed, None);
801        dequeue_next_as.test_rust_equivalence(initial_state);
802    }
803
804    /// Helps testing dequeuing multiple items.
805    #[derive(Debug, Default, Clone, Eq, PartialEq)]
806    struct TestHelperDequeueMultipleAs {
807        proof_items: Vec<ProofItemVariant>,
808    }
809
810    impl BasicSnippet for TestHelperDequeueMultipleAs {
811        fn parameters(&self) -> Vec<(DataType, String)> {
812            vec![(DataType::VoidPointer, "*proof_item_iter".to_string())]
813        }
814
815        fn return_values(&self) -> Vec<(DataType, String)> {
816            vec![(DataType::VoidPointer, "*proof_item_iter".to_string())]
817        }
818
819        fn entrypoint(&self) -> String {
820            "test_helper_dequeue_multiple_as".to_string()
821        }
822
823        fn code(&self, library: &mut Library) -> Vec<LabelledInstruction> {
824            let dequeue_next_as_entrypoints = self
825                .proof_items
826                .iter()
827                .map(|&item| library.import(Box::new(DequeueNextAs::new(item))))
828                .collect::<Vec<_>>();
829
830            let code_body = dequeue_next_as_entrypoints
831                .into_iter()
832                .flat_map(|snippet| triton_asm!(dup 0 call {snippet} pop 1))
833                .collect_vec();
834
835            triton_asm!({self.entrypoint()}: {&code_body} return)
836        }
837    }
838
839    impl Procedure for TestHelperDequeueMultipleAs {
840        fn rust_shadow(
841            &self,
842            stack: &mut Vec<BFieldElement>,
843            memory: &mut HashMap<BFieldElement, BFieldElement>,
844            non_determinism: &NonDeterminism,
845            std_in: &[BFieldElement],
846            sponge: &mut Option<Tip5>,
847        ) -> Result<Vec<BFieldElement>, RustShadowError> {
848            let proof_iter_pointer = stack
849                .last()
850                .copied()
851                .ok_or(RustShadowError::StackUnderflow)?;
852            for &proof_item in &self.proof_items {
853                let dequeue_next_as = DequeueNextAs { proof_item };
854                stack.push(proof_iter_pointer);
855                dequeue_next_as.rust_shadow(stack, memory, non_determinism, std_in, sponge)?;
856                stack.pop().ok_or(RustShadowError::StackUnderflow)?;
857            }
858
859            Ok(Vec::new())
860        }
861
862        fn pseudorandom_initial_state(
863            &self,
864            seed: [u8; 32],
865            _: Option<BenchmarkCase>,
866        ) -> ProcedureInitialState {
867            let mut rng = StdRng::from_seed(seed);
868            let proof_stream =
869                DequeueNextAs::pseudorandom_proof_stream(self.proof_items.clone(), rng.random());
870
871            // We just use the state-initialization method from the `[DequeueNextAs]` snippet.
872            // This is OK as long as we don't read the program hash which we don't do in this
873            // snippet.
874            let dummy_snippet_for_state_init = DequeueNextAs::new(ProofItemVariant::MerkleRoot);
875            dummy_snippet_for_state_init
876                .initial_state_from_proof_and_address(proof_stream.into(), rng.random())
877        }
878    }
879
880    impl TestHelperDequeueMultipleAs {
881        fn test_rust_equivalence(self, initial_state: ProcedureInitialState) {
882            test_rust_equivalence_given_complete_state(
883                &ShadowedProcedure::new(self),
884                &initial_state.stack,
885                &initial_state.public_input,
886                &initial_state.nondeterminism,
887                &initial_state.sponge,
888                None,
889            );
890        }
891    }
892
893    #[macro_rules_attr::apply(test)]
894    fn dequeue_two() {
895        let proof_items = vec![ProofItemVariant::MerkleRoot, ProofItemVariant::MerkleRoot];
896        let dequeue_multiple = TestHelperDequeueMultipleAs { proof_items };
897        let initial_state = dequeue_multiple.pseudorandom_initial_state([0; 32], None);
898        dequeue_multiple.test_rust_equivalence(initial_state);
899    }
900
901    #[macro_rules_attr::apply(proptest)]
902    fn dequeue_multiple(#[strategy(arb())] proof_items: Vec<ProofItem>, seed: [u8; 32]) {
903        let proof_items = proof_items.into_iter().map(|i| i.into()).collect_vec();
904        let dequeue_multiple = TestHelperDequeueMultipleAs { proof_items };
905        let initial_state = dequeue_multiple.pseudorandom_initial_state(seed, None);
906        dequeue_multiple.test_rust_equivalence(initial_state);
907    }
908}