1use itertools::Itertools;
2use num_traits::Zero;
3use triton_vm::prelude::*;
4use twenty_first::error::BFieldCodecError;
5use twenty_first::error::PolynomialBFieldCodecError;
6use twenty_first::math::x_field_element::EXTENSION_DEGREE;
7use twenty_first::prelude::*;
8
9use super::tasm_object::Result;
10use crate::prelude::*;
11
12impl<const N: usize, T> TasmObject for [T; N]
13where
14 T: TasmObject,
15{
16 fn label_friendly_name() -> String {
17 format!("array{N}___{}", T::label_friendly_name())
18 }
19
20 fn compute_size_and_assert_valid_size_indicator(_: &mut Library) -> Vec<LabelledInstruction> {
21 if let Some(static_size) = Self::static_length() {
22 return triton_asm!(
23 pop 1
26 push {static_size}
27 );
29 }
30
31 todo!()
32 }
33
34 fn decode_iter<Itr: Iterator<Item = BFieldElement>>(iterator: &mut Itr) -> Result<Box<Self>> {
35 let mut vector = Vec::with_capacity(N);
36 for _ in 0..N {
37 if T::static_length().is_none() {
38 iterator.next().ok_or(BFieldCodecError::SequenceTooShort)?;
39 };
40 vector.push(*T::decode_iter(iterator)?);
41 }
42
43 let Ok(array) = vector.try_into() else {
45 unreachable!()
46 };
47
48 Ok(Box::new(array))
49 }
50}
51
52impl<T> TasmObject for Vec<T>
53where
54 T: TasmObject,
55{
56 fn label_friendly_name() -> String {
57 format!("vec___{}", T::label_friendly_name())
58 }
59
60 fn compute_size_and_assert_valid_size_indicator(
61 library: &mut Library,
62 ) -> Vec<LabelledInstruction> {
63 if T::static_length().is_some() {
64 return triton_asm!(
65 read_mem 1
68 {&T::compute_size_and_assert_valid_size_indicator(library)}
71 mul
74 addi 1
75 );
78 }
79
80 let loop_label = format!(
81 "tasmlib_structure_tasmobject_verify_size_indicators_dyn_elem_sizes___{}",
82 T::label_friendly_name()
83 );
84
85 let loop_code = triton_asm!(
86 {loop_label}:
88
89 dup 2
90 push 0
91 eq
92 skiz
93 return
94 read_mem 1
97 push {T::MAX_OFFSET}
101 dup 2
102 lt
103 assert error_id 210
104 addi 2
107 dup 0
110 {&T::compute_size_and_assert_valid_size_indicator(library)}
111 dup 2
114 eq
115 assert error_id 211
116 pick 2
119 dup 2
120 add
121 addi 1
125 place 2
128 add
131 pick 2
134 addi -1
135 place 2
136 recurse
139 );
140
141 library.explicit_import(&loop_label, &loop_code);
142 triton_asm!(
143 read_mem 1
146 hint remaining_elements = stack[1]
147 addi 2
150 hint elem_si_ptr = stack[0]
151 push 0
155 hint acc_size = stack[0]
156 place 1
159 call {loop_label}
162 pick 2
165 pop 2
166 addi 1
170 )
171 }
172
173 fn decode_iter<Itr: Iterator<Item = BFieldElement>>(iterator: &mut Itr) -> Result<Box<Self>> {
174 let vec_length = iterator.next().ok_or(BFieldCodecError::SequenceTooShort)?;
175 let mut vector = vec![];
176 for _ in 0..vec_length.value() {
177 if T::static_length().is_none() {
178 iterator.next().ok_or(BFieldCodecError::SequenceTooShort)?;
179 };
180 vector.push(*T::decode_iter(iterator)?);
181 }
182
183 Ok(Box::new(vector))
184 }
185}
186
187impl TasmObject for BFieldElement {
188 fn label_friendly_name() -> String {
189 DataType::Bfe.label_friendly_name()
190 }
191
192 fn compute_size_and_assert_valid_size_indicator(_: &mut Library) -> Vec<LabelledInstruction> {
193 triton_asm!(pop 1 push {Self::static_length().unwrap()})
194 }
195
196 fn decode_iter<Itr: Iterator<Item = BFieldElement>>(iterator: &mut Itr) -> Result<Box<Self>> {
197 let word = iterator.next().ok_or(BFieldCodecError::SequenceTooShort)?;
198 Ok(Box::new(word))
199 }
200}
201
202impl TasmObject for XFieldElement {
203 fn label_friendly_name() -> String {
204 DataType::Xfe.label_friendly_name()
205 }
206
207 fn compute_size_and_assert_valid_size_indicator(_: &mut Library) -> Vec<LabelledInstruction> {
208 triton_asm!(pop 1 push {Self::static_length().unwrap()})
209 }
210
211 fn decode_iter<Itr: Iterator<Item = BFieldElement>>(iterator: &mut Itr) -> Result<Box<Self>> {
212 let (c_0, c_1, c_2) = iterator
213 .next_tuple()
214 .ok_or(BFieldCodecError::SequenceTooShort)?;
215
216 Ok(Box::new(xfe!([c_0, c_1, c_2])))
217 }
218}
219
220impl TasmObject for Digest {
221 fn label_friendly_name() -> String {
222 DataType::Digest.label_friendly_name()
223 }
224
225 fn compute_size_and_assert_valid_size_indicator(_: &mut Library) -> Vec<LabelledInstruction> {
226 triton_asm!(pop 1 push {Self::static_length().unwrap()})
227 }
228
229 fn decode_iter<Itr: Iterator<Item = BFieldElement>>(iterator: &mut Itr) -> Result<Box<Self>> {
230 let (d_0, d_1, d_2, d_3, d_4) = iterator
231 .next_tuple()
232 .ok_or(BFieldCodecError::SequenceTooShort)?;
233
234 Ok(Box::new(Digest::new([d_0, d_1, d_2, d_3, d_4])))
235 }
236}
237
238impl TasmObject for bool {
239 fn label_friendly_name() -> String {
240 DataType::Bool.label_friendly_name()
241 }
242
243 fn compute_size_and_assert_valid_size_indicator(_: &mut Library) -> Vec<LabelledInstruction> {
244 triton_asm!(pop 1 push {Self::static_length().unwrap()})
245 }
246
247 fn decode_iter<Itr: Iterator<Item = BFieldElement>>(iterator: &mut Itr) -> Result<Box<Self>> {
248 let the_bool = iterator.next().ok_or(BFieldCodecError::SequenceTooShort)?;
249 match the_bool.value() {
250 0 => Ok(Box::new(false)),
251 1 => Ok(Box::new(true)),
252 _ => Err(Box::new(BFieldCodecError::ElementOutOfRange)),
253 }
254 }
255}
256
257impl TasmObject for u32 {
258 fn label_friendly_name() -> String {
259 DataType::U32.label_friendly_name()
260 }
261
262 fn compute_size_and_assert_valid_size_indicator(_: &mut Library) -> Vec<LabelledInstruction> {
263 triton_asm!(pop 1 push {Self::static_length().unwrap()})
264 }
265
266 fn decode_iter<Itr: Iterator<Item = BFieldElement>>(iterator: &mut Itr) -> Result<Box<Self>> {
267 let word = iterator
268 .next()
269 .ok_or(BFieldCodecError::SequenceTooShort)?
270 .try_into()
271 .map_err(|_| BFieldCodecError::ElementOutOfRange)?;
272
273 Ok(Box::new(word))
274 }
275}
276
277impl TasmObject for u64 {
278 fn label_friendly_name() -> String {
279 DataType::U64.label_friendly_name()
280 }
281
282 fn compute_size_and_assert_valid_size_indicator(_: &mut Library) -> Vec<LabelledInstruction> {
283 triton_asm!(pop 1 push {Self::static_length().unwrap()})
284 }
285
286 fn decode_iter<Itr: Iterator<Item = BFieldElement>>(iterator: &mut Itr) -> Result<Box<Self>> {
287 let (val_lo, val_hi) = iterator
288 .next_tuple()
289 .ok_or(BFieldCodecError::SequenceTooShort)?;
290
291 Ok(Self::decode(&[val_lo, val_hi])?)
292 }
293}
294
295impl TasmObject for u128 {
296 fn label_friendly_name() -> String {
297 DataType::U128.label_friendly_name()
298 }
299
300 fn compute_size_and_assert_valid_size_indicator(_: &mut Library) -> Vec<LabelledInstruction> {
301 triton_asm!(pop 1 push {Self::static_length().unwrap()})
302 }
303
304 fn decode_iter<Itr: Iterator<Item = BFieldElement>>(iterator: &mut Itr) -> Result<Box<Self>> {
305 let (lo_lo, lo_hi, hi_lo, hi_hi) = iterator
306 .next_tuple()
307 .ok_or(BFieldCodecError::SequenceTooShort)?;
308
309 Ok(Self::decode(&[lo_lo, lo_hi, hi_lo, hi_hi])?)
310 }
311}
312
313impl<T, S> TasmObject for (T, S)
314where
315 T: TasmObject,
316 S: TasmObject,
317{
318 fn label_friendly_name() -> String {
319 format!(
320 "tuple_L_{}_{}_R",
321 T::label_friendly_name(),
322 S::label_friendly_name()
323 )
324 }
325
326 fn compute_size_and_assert_valid_size_indicator(
327 library: &mut Library,
328 ) -> Vec<LabelledInstruction> {
329 let size_left = if T::static_length().is_some() {
330 T::compute_size_and_assert_valid_size_indicator(library)
331 } else {
332 triton_asm!(
333 hint left_si_ptr = stack[0]
335
336 read_mem 1
337 addi 2
338 {&T::compute_size_and_assert_valid_size_indicator(library)}
341 hint calculated_left = stack[0]
342 dup 1
345 eq
346 assert error_id 220
347 addi 1
350 )
352 };
353
354 let size_right = if S::static_length().is_some() {
355 S::compute_size_and_assert_valid_size_indicator(library)
356 } else {
357 triton_asm!(
358 hint right_si_ptr = stack[0]
360
361 read_mem 1
362 hint right_si = stack[1]
363 addi 2
367 {&S::compute_size_and_assert_valid_size_indicator(library)}
370 hint calculated_right = stack[0]
371 dup 1
374 eq
375 assert error_id 221
376 addi 1
380 )
382 };
383
384 triton_asm!(
385 dup 0
388 {&size_right}
389 hint right_ptr_or_right_si = stack[1]
390 hint right_size_incl_pot_si = stack[0]
391 pick 1
394 dup 1
395 add
396 hint left: Pointer = stack[0]
397 {&size_left}
401 hint left_size_incl_si = stack[0]
402 add
405 )
407 }
408
409 fn decode_iter<Itr: Iterator<Item = BFieldElement>>(iterator: &mut Itr) -> Result<Box<Self>> {
410 if S::static_length().is_none() {
411 iterator.next().ok_or(BFieldCodecError::SequenceTooShort)?;
412 }
413 let s = *S::decode_iter(iterator)?;
414
415 if T::static_length().is_none() {
416 iterator.next().ok_or(BFieldCodecError::SequenceTooShort)?;
417 }
418 let t = *T::decode_iter(iterator)?;
419
420 Ok(Box::new((t, s)))
421 }
422}
423
424impl TasmObject for Polynomial<'static, XFieldElement> {
425 fn label_friendly_name() -> String {
426 "polynomial_xfe".to_owned()
427 }
428
429 fn compute_size_and_assert_valid_size_indicator(_: &mut Library) -> Vec<LabelledInstruction> {
430 triton_asm!(
431 addi 1
434 read_mem 2
437 pop 1
438 push {Self::MAX_OFFSET}
441 dup 1
442 lt
443 assert
444 pick 1
447 push {EXTENSION_DEGREE}
448 mul
449 addi 1
450 dup 1
453 eq
454 assert
455 addi 1
459 )
460 }
461
462 fn decode_iter<Itr: Iterator<Item = BFieldElement>>(iterator: &mut Itr) -> Result<Box<Self>> {
463 iterator.next().ok_or(BFieldCodecError::SequenceTooShort)?;
464 let coefficients = *Vec::<XFieldElement>::decode_iter(iterator)?;
465 if coefficients.last().is_some_and(|c| c.is_zero()) {
466 return Err(PolynomialBFieldCodecError::TrailingZerosInPolynomialEncoding)?;
467 }
468
469 Ok(Box::new(Self::new(coefficients)))
470 }
471}
472
473impl TasmObject for Proof {
474 fn label_friendly_name() -> String {
475 "tvm_proof".to_owned()
476 }
477
478 fn compute_size_and_assert_valid_size_indicator(_: &mut Library) -> Vec<LabelledInstruction> {
479 triton_asm!(
485 read_mem 1
488 pop 1
489 addi 1
492 )
494 }
495
496 fn decode_iter<Itr: Iterator<Item = BFieldElement>>(iterator: &mut Itr) -> Result<Box<Self>> {
497 let _ = iterator.next().ok_or(BFieldCodecError::SequenceTooShort)?;
498 Ok(Box::new(Self(*Vec::decode_iter(iterator)?)))
499 }
500}
501
502impl<T> TasmObject for Option<T>
503where
504 T: TasmObject,
505{
506 fn label_friendly_name() -> String {
507 format!("option_L_{}_R", T::label_friendly_name())
508 }
509
510 fn compute_size_and_assert_valid_size_indicator(
511 library: &mut Library,
512 ) -> Vec<LabelledInstruction> {
513 let some_branch_label = format!(
514 "tasmlib_tasmobject_size_verifier_option_some_branch___{}",
515 T::label_friendly_name()
516 );
517 let some_branch = triton_asm!(
518 {some_branch_label}:
519
520 pop 1
522
523 {&T::compute_size_and_assert_valid_size_indicator(library)}
524 push 0
528
529 return
530 );
531
532 let none_branch_label = "tasmlib_tasmobject_size_verifier_option_none".to_owned();
533 let none_branch = triton_asm!(
534 {none_branch_label}:
535 pop 1
538 push 0
539 return
542 );
543
544 library.explicit_import(&some_branch_label, &some_branch);
545 library.explicit_import(&none_branch_label, &none_branch);
546
547 triton_asm!(
548 read_mem 1
550 addi 2
551 dup 1
555 push 0
556 eq
557 dup 2
558 push 1
559 eq
560 add
561 assert error_id 200
564 push 1
567 pick 2
570 push 1
571 eq
572 skiz
575 call {some_branch_label}
576 skiz
577 call {none_branch_label}
578
579 addi 1
581
582 )
584 }
585
586 fn decode_iter<Itr: Iterator<Item = BFieldElement>>(iterator: &mut Itr) -> Result<Box<Self>> {
587 let is_some = *bool::decode_iter(iterator)?;
588 let the_option = is_some.then(|| T::decode_iter(iterator)).transpose()?;
589 Ok(Box::new(the_option.map(|t| *t)))
590 }
591}
592
593#[cfg(test)]
594mod tests {
595 use std::collections::HashMap;
596 use std::fmt::Debug;
597
598 use proptest::prelude::*;
599 use proptest::test_runner::TestCaseResult;
600 use proptest_arbitrary_interop::arb;
601 use test_strategy::proptest;
602 use test_strategy::Arbitrary;
603
604 use super::*;
605 use crate::memory::encode_to_memory;
606
607 #[derive(Debug, Clone, Arbitrary)]
608 struct TestObject<T>
609 where
610 T: TasmObject + Debug + Clone + Eq + for<'a> arbitrary::Arbitrary<'a> + 'static,
611 {
612 #[strategy(arb())]
613 t: T,
614
615 #[strategy(arb())]
616 address: BFieldElement,
617 }
618
619 impl<T> TestObject<T>
620 where
621 T: TasmObject + Debug + Clone + Eq + for<'a> arbitrary::Arbitrary<'a>,
622 {
623 fn verify_decoding_properties(&self) -> TestCaseResult {
624 let mut memory = HashMap::default();
625 encode_to_memory(&mut memory, self.address, &self.t);
626 let decoding_result = T::decode_from_memory(&memory, self.address);
627 prop_assert!(decoding_result.is_ok());
628
629 let obj_read = *decoding_result.unwrap();
630 prop_assert_eq!(&self.t, &obj_read);
631
632 Ok(())
633 }
634 }
635
636 macro_rules! gen_test {
638 (fn $test_fn:ident for $type:ty) => {
639 #[proptest]
640 fn $test_fn(test_object: TestObject<$type>) {
641 test_object.verify_decoding_properties()?;
642 }
643 };
644 }
645
646 gen_test!(fn decode_bfe for BFieldElement);
647 gen_test!(fn decode_xfe for XFieldElement);
648 gen_test!(fn decode_digest for Digest);
649 gen_test!(fn decode_bool for bool);
650 gen_test!(fn decode_u32 for u32);
651 gen_test!(fn decode_u64 for u64);
652 gen_test!(fn decode_u128 for u128);
653 gen_test!(fn decode_poly_xfe for Polynomial<'static, XFieldElement>);
654 gen_test!(fn decode_proof for Proof);
655
656 gen_test!(fn decode_array_bfe_0 for [BFieldElement; 0]);
657 gen_test!(fn decode_array_bfe_2 for [BFieldElement; 2]);
658 gen_test!(fn decode_array_bfe_10 for [BFieldElement; 10]);
659 gen_test!(fn decode_array_bfe_25 for [BFieldElement; 25]);
660 gen_test!(fn decode_array_xfe_0 for [BFieldElement; 0]);
661 gen_test!(fn decode_array_xfe_2 for [BFieldElement; 2]);
662 gen_test!(fn decode_array_xfe_10 for [BFieldElement; 10]);
663 gen_test!(fn decode_array_xfe_25 for [BFieldElement; 25]);
664 gen_test!(fn decode_array_vec_bfe_0 for [Vec<BFieldElement>; 0]);
665 gen_test!(fn decode_array_vec_bfe_2 for [Vec<BFieldElement>; 2]);
666 gen_test!(fn decode_array_vec_vec_bfe_2 for [Vec<Vec<BFieldElement>>; 2]);
667 gen_test!(fn decode_array_tuple_l_bfe_bfe_r_0 for [(BFieldElement, BFieldElement); 0]);
668 gen_test!(fn decode_array_tuple_l_bfe_bfe_r_2 for [(BFieldElement, BFieldElement); 2]);
669 gen_test!(fn decode_array_tuple_l_u64_u128_r_0 for [(u64, u128); 0]);
670 gen_test!(fn decode_array_tuple_l_u64_u128_r_2 for [(u64, u128); 2]);
671 gen_test!(fn decode_array_tuple_l_vec_u64_vec_u64_r_2 for [(Vec<u64>, Vec<u64>); 2]);
672 gen_test!(fn decode_array_option_bfe_0 for [Option<BFieldElement>; 0]);
673 gen_test!(fn decode_array_option_bfe_5 for [Option<BFieldElement>; 5]);
674 gen_test!(fn decode_array_option_vec_u64_5 for [Option<Vec<u64>>; 5]);
675
676 gen_test!(fn decode_vec_bfe for Vec<BFieldElement>);
677 gen_test!(fn decode_vec_xfe for Vec<XFieldElement>);
678 gen_test!(fn decode_vec_digest for Vec<Digest>);
679 gen_test!(fn decode_vec_array_u32_0 for Vec<[u32; 0]>);
680 gen_test!(fn decode_vec_array_u32_7 for Vec<[u32; 7]>);
681 gen_test!(fn decode_vec_array_u64_0 for Vec<[u64; 0]>);
682 gen_test!(fn decode_vec_array_u64_7 for Vec<[u64; 7]>);
683 gen_test!(fn decode_vec_vec_digest for Vec<Vec<Digest>>);
684 gen_test!(fn decode_vec_tuple_l_u32_u64_r for Vec<(u32, u64)>);
685 gen_test!(fn decode_vec_vec_tuple_l_u32_u64_r for Vec<Vec<(u32, u64)>>);
686 gen_test!(fn decode_vec_tuple_l_vec_u32_vec_u64_r for Vec<(Vec<u32>, Vec<u64>)>);
687 gen_test!(fn decode_vec_option_bfe for Vec<Option<BFieldElement>>);
688 gen_test!(fn decode_vec_option_tuple_l_u128_vec_digest_r for Vec<Option<(u128, Vec<Digest>)>>);
689
690 gen_test!(fn decode_tuple_l_bfe_bfe_r for (BFieldElement, BFieldElement));
691 gen_test!(fn decode_tuple_l_bfe_xfe_r for (BFieldElement, XFieldElement));
692 gen_test!(fn decode_tuple_l_digest_xfe_r for (Digest, XFieldElement));
693 gen_test!(fn decode_tuple_l_digest_vec_xfe_r for (Digest, Vec<XFieldElement>));
694 gen_test!(fn decode_tuple_l_vec_digest_xfe_r for (Vec<Digest>, XFieldElement));
695 gen_test!(fn decode_tuple_l_digest_array_xfe_0_r for (Digest, [XFieldElement; 0]));
696 gen_test!(fn decode_tuple_l_digest_array_xfe_5_r for (Digest, [XFieldElement; 5]));
697 gen_test!(fn decode_tuple_l_array_xfe_0_digest_r for ([XFieldElement; 0], Digest));
698 gen_test!(fn decode_tuple_l_array_xfe_5_digest_r for ([XFieldElement; 5], Digest));
699 gen_test!(fn decode_tuple_l_array_u128_5_u64_5_digest_r for ([u128; 5], [u64; 5]));
700 gen_test!(fn decode_tuple_l_vec_u64_array_xfe_5_r for (Vec<u64>, [XFieldElement; 5]));
701 gen_test!(fn decode_tuple_l_tuple_l_u32_u64_r_u128_r for ((u32, u64), u128));
702 gen_test!(fn decode_tuple_l_u32_tuple_l_u64_u128_r_r for (u32, (u64, u128)));
703 gen_test!(fn decode_tuple_l_tuple_l_u32_u64_r_tuple_l_u64_u32_r_r for ((u32, u64), (u64, u32)));
704 gen_test!(fn decode_tuple_l_vec_vec_u32_vec_vec_u128_r for (Vec<Vec<u32>>, Vec<Vec<u128>>));
705 gen_test!(fn decode_tuple_l_option_u32_option_u64_r for (Option<u32>, Option<u64>));
706 gen_test!(fn decode_tuple_l_option_vec_u32_option_u64_r for (Option<Vec<u32>>, Option<u64>));
707
708 gen_test!(fn decode_option_bfe for Option<BFieldElement>);
709 gen_test!(fn decode_option_xfe for Option<XFieldElement>);
710 gen_test!(fn decode_option_digest for Option<Digest>);
711 gen_test!(fn decode_option_vec_bfe for Option<Vec<BFieldElement>>);
712 gen_test!(fn decode_option_vec_xfe for Option<Vec<XFieldElement>>);
713 gen_test!(fn decode_option_option_bfe for Option<Option<BFieldElement>>);
714 gen_test!(fn decode_option_array_bfe_0 for Option<[BFieldElement; 0]>);
715 gen_test!(fn decode_option_array_xfe_5 for Option<[XFieldElement; 5]>);
716 gen_test!(fn decode_option_tuple_l_bfe_xfe_r for Option<(BFieldElement, XFieldElement)>);
717 gen_test!(fn decode_option_tuple_l_u32_u64_r for Option<(u32, u64)>);
718 gen_test!(fn decode_option_vec_tuple_l_bfe_u128_r for Option<Vec<(BFieldElement, u128)>>);
719 gen_test!(fn decode_option_tuple_l_vec_u32_vec_u64_r for Option<(Vec<u32>, Vec<u64>)>);
720 gen_test!(fn decode_option_tuple_l_array_u128_5_u64_5_r for Option<([u128; 5], [u64; 5])>);
721}