openvm_sdk/
codec.rs

1use std::io::{self, Cursor, Read, Result, Write};
2
3use openvm_circuit::{
4    arch::ContinuationVmProof, system::memory::merkle::public_values::UserPublicValuesProof,
5};
6use openvm_continuations::verifier::{
7    internal::types::VmStarkProof, root::types::RootVmVerifierInput,
8};
9use openvm_native_compiler::ir::DIGEST_SIZE;
10use openvm_native_recursion::hints::{InnerBatchOpening, InnerFriProof, InnerQueryProof};
11use openvm_stark_backend::{
12    config::{Com, PcsProof},
13    interaction::{fri_log_up::FriLogUpPartialProof, RapPhaseSeqKind},
14    p3_field::{
15        extension::BinomialExtensionField, FieldAlgebra, FieldExtensionAlgebra, PrimeField32,
16    },
17    proof::{AdjacentOpenedValues, AirProofData, Commitments, OpenedValues, OpeningProof, Proof},
18};
19use p3_fri::CommitPhaseProofStep;
20
21use super::{F, SC};
22
23type Challenge = BinomialExtensionField<F, 4>;
24
25/// Codec version should change only when proof system or proof format changes.
26/// It does correspond to the main openvm version (which may change more frequently).
27const CODEC_VERSION: u32 = 1;
28
29/// Hardware and language independent encoding.
30/// Uses the Writer pattern for more efficient encoding without intermediate buffers.
31// @dev Trait just for implementation sanity
32pub trait Encode {
33    /// Writes the encoded representation of `self` to the given writer.
34    fn encode<W: Write>(&self, writer: &mut W) -> Result<()>;
35
36    /// Convenience method to encode into a `Vec<u8>`
37    fn encode_to_vec(&self) -> Result<Vec<u8>> {
38        let mut buffer = Vec::new();
39        self.encode(&mut buffer)?;
40        Ok(buffer)
41    }
42}
43
44/// Hardware and language independent decoding.
45/// Uses the Reader pattern for efficient decoding.
46pub trait Decode: Sized {
47    /// Reads and decodes a value from the given reader.
48    fn decode<R: Read>(reader: &mut R) -> Result<Self>;
49    fn decode_from_bytes(bytes: &[u8]) -> Result<Self> {
50        let mut reader = Cursor::new(bytes);
51        Self::decode(&mut reader)
52    }
53}
54
55// ==================== Encode implementation ====================
56
57impl Encode for ContinuationVmProof<SC> {
58    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
59        encode_slice(&self.per_segment, writer)?;
60        self.user_public_values.encode(writer)
61    }
62}
63
64impl Encode for VmStarkProof<SC> {
65    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
66        self.inner.encode(writer)?;
67        encode_slice(&self.user_public_values, writer)
68    }
69}
70
71impl Encode for UserPublicValuesProof<DIGEST_SIZE, F> {
72    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
73        encode_slice(&self.proof, writer)?;
74        encode_slice(&self.public_values, writer)?;
75        self.public_values_commit.encode(writer)
76    }
77}
78
79impl Encode for RootVmVerifierInput<SC> {
80    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
81        encode_slice(&self.proofs, writer)?;
82        encode_slice(&self.public_values, writer)
83    }
84}
85
86impl Encode for Proof<SC> {
87    // We need to know:
88    // - Pcs is TwoAdicFriPcs
89    // - Com<SC>: Into<[F; 8]>
90    // For simplicity, we only implement for fixed `BabyBearPoseidon2Config`
91    //
92    /// Encode a proof using FRI as the PCS with `BabyBearPoseidon2Config`.
93    /// The Merkle tree hashes have digest `[F; 8]`.
94    /// ```
95    /// pub struct Proof<SC: StarkGenericConfig> {
96    ///     pub commitments: Commitments<Com<SC>>,
97    ///     pub opening: OpeningProof<PcsProof<SC>, SC::Challenge>,
98    ///     pub per_air: Vec<AirProofData<Val<SC>, SC::Challenge>>,
99    ///     pub rap_phase_seq_proof: Option<RapPhaseSeqPartialProof<SC>>,
100    /// }
101    /// ```
102    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
103        writer.write_all(&CODEC_VERSION.to_le_bytes())?;
104        // Encode commitments
105        encode_commitments(&self.commitments.main_trace, writer)?;
106        encode_commitments(&self.commitments.after_challenge, writer)?;
107        let quotient_commit: [F; DIGEST_SIZE] = self.commitments.quotient.into();
108        quotient_commit.encode(writer)?;
109
110        // Encode OpeningProof
111        encode_opening_proof(&self.opening, writer)?;
112
113        // Encode per_air data
114        encode_slice(&self.per_air, writer)?;
115
116        writer.write_all(&[RapPhaseSeqKind::FriLogUp as u8])?;
117        // Encode logup witness
118        self.rap_phase_seq_proof.encode(writer)?;
119
120        Ok(())
121    }
122}
123
124// Helper function to encode OpeningProof
125// ```
126// pub struct OpeningProof<PcsProof, Challenge> {
127//     pub proof: PcsProof,
128//     pub values: OpenedValues<Challenge>,
129// }
130// ```
131fn encode_opening_proof<W: Write>(
132    opening: &OpeningProof<PcsProof<SC>, Challenge>,
133    writer: &mut W,
134) -> Result<()> {
135    // Encode FRI proof
136    opening.proof.encode(writer)?;
137    encode_opened_values(&opening.values, writer)?;
138    Ok(())
139}
140
141/// [OpenedValues] is a typedef for `Vec<Vec<Vec<Vec<F>>>>` for
142/// - each round
143///   - each matrix
144///     - each point to open at
145///       - evaluations for each column of matrix at that point
146fn encode_opened_values<W: Write>(
147    opened_values: &OpenedValues<Challenge>,
148    writer: &mut W,
149) -> Result<()> {
150    encode_slice(&opened_values.preprocessed, writer)?;
151    opened_values.main.len().encode(writer)?;
152    for part in &opened_values.main {
153        encode_slice(part, writer)?;
154    }
155    opened_values.after_challenge.len().encode(writer)?;
156    for phase in &opened_values.after_challenge {
157        encode_slice(phase, writer)?;
158    }
159    opened_values.quotient.len().encode(writer)?;
160    for per_air in &opened_values.quotient {
161        per_air.len().encode(writer)?;
162        for chunk in per_air {
163            encode_slice(chunk, writer)?;
164        }
165    }
166
167    Ok(())
168}
169
170impl Encode for AdjacentOpenedValues<Challenge> {
171    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
172        encode_slice(&self.local, writer)?;
173        encode_slice(&self.next, writer)?;
174        Ok(())
175    }
176}
177
178impl Encode for AirProofData<F, Challenge> {
179    /// Encodes the struct
180    /// ```
181    /// pub struct OpenedValues<Challenge> {
182    ///     pub preprocessed: Vec<AdjacentOpenedValues<Challenge>>,
183    ///     pub main: Vec<Vec<AdjacentOpenedValues<Challenge>>>,
184    ///     pub after_challenge: Vec<Vec<AdjacentOpenedValues<Challenge>>>,
185    ///     pub quotient: Vec<Vec<Vec<Challenge>>>,
186    /// }
187    /// ```
188    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
189        self.air_id.encode(writer)?;
190        self.degree.encode(writer)?;
191        self.exposed_values_after_challenge.len().encode(writer)?;
192        for exposed_vals in &self.exposed_values_after_challenge {
193            encode_slice(exposed_vals, writer)?;
194        }
195        encode_slice(&self.public_values, writer)?;
196        Ok(())
197    }
198}
199
200// PcsProof<SC> = InnerFriProof where Pcs = TwoAdicFriPcs
201impl Encode for InnerFriProof {
202    /// Encodes the struct
203    /// ```
204    /// pub struct FriProof<Challenge, M: Mmcs<Challenge>> {
205    ///     pub commit_phase_commits: Vec<M::Commitment>,
206    ///     pub query_proofs: Vec<QueryProof<Challenge, M, Vec<BatchOpening<F>>>>,
207    ///     pub final_poly: Vec<Challenge>,
208    ///     pub pow_witness: F,
209    /// }
210    /// ```
211    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
212        encode_commitments(&self.commit_phase_commits, writer)?;
213        encode_slice(&self.query_proofs, writer)?;
214        encode_slice(&self.final_poly, writer)?;
215        self.pow_witness.encode(writer)?;
216        Ok(())
217    }
218}
219
220impl Encode for InnerQueryProof {
221    /// Encodes the struct
222    /// ```
223    /// pub struct QueryProof<Challenge, M: Mmcs<Challenge>> {
224    ///     pub input_proof: Vec<BatchOpening<F>>,
225    ///     pub commit_phase_openings: Vec<CommitPhaseProofStep<Challenge, M>>,
226    /// }
227    ///
228    /// pub struct BatchOpening<F> {
229    ///     pub opened_values: Vec<Vec<F>>,
230    ///     pub opening_proof: Vec<[F; DIGEST_SIZE]>,
231    /// }
232    ///
233    /// pub struct CommitPhaseProofStep<Challenge, M: Mmcs<Challenge>> {
234    ///     pub sibling_value: Challenge,
235    ///     pub opening_proof: Vec<[F; DIGEST_SIZE]>,
236    /// }
237    /// ```
238    // @dev [jpw]: We prefer to keep the implementation all in one function
239    // without `impl Encode` on subtypes because it obfuscates what the overall
240    // struct consists of.
241    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
242        // Input proof is Vec<BatchOpening<F>>
243        self.input_proof.len().encode(writer)?;
244        for batch_opening in &self.input_proof {
245            batch_opening.opened_values.len().encode(writer)?;
246            for vals in &batch_opening.opened_values {
247                encode_slice(vals, writer)?;
248            }
249            // Opening proof is just a vector of siblings
250            encode_slice(&batch_opening.opening_proof, writer)?;
251        }
252        self.commit_phase_openings.len().encode(writer)?;
253        for step in &self.commit_phase_openings {
254            step.sibling_value.encode(writer)?;
255            encode_slice(&step.opening_proof, writer)?;
256        }
257        Ok(())
258    }
259}
260
261impl Encode for Option<FriLogUpPartialProof<F>> {
262    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
263        match self {
264            // If exists, `F` will be < MODULUS < 2^31 so it will
265            // never collide with u32::MAX
266            Some(FriLogUpPartialProof { logup_pow_witness }) => logup_pow_witness.encode(writer),
267            None => writer.write_all(&u32::MAX.to_le_bytes()),
268        }
269    }
270}
271
272impl Encode for Challenge {
273    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
274        let base_slice: &[F] = self.as_base_slice();
275        // Fixed length slice, so don't encode length
276        for val in base_slice {
277            val.encode(writer)?;
278        }
279        Ok(())
280    }
281}
282
283/// Encodes length of slice and then each commitment
284fn encode_commitments<W: Write>(commitments: &[Com<SC>], writer: &mut W) -> Result<()> {
285    let coms: Vec<[F; DIGEST_SIZE]> = commitments.iter().copied().map(Into::into).collect();
286    encode_slice(&coms, writer)
287}
288
289// Can't implement Encode on Com<SC> because Rust complains about associated trait types when you
290// don't own the trait (in this case SC)
291impl Encode for [F; DIGEST_SIZE] {
292    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
293        for val in self {
294            val.encode(writer)?;
295        }
296        Ok(())
297    }
298}
299
300/// Encodes length of slice and then each element
301pub fn encode_slice<T: Encode, W: Write>(slice: &[T], writer: &mut W) -> Result<()> {
302    slice.len().encode(writer)?;
303    for elt in slice {
304        elt.encode(writer)?;
305    }
306    Ok(())
307}
308
309impl Encode for F {
310    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
311        writer.write_all(&self.as_canonical_u32().to_le_bytes())
312    }
313}
314
315impl Encode for usize {
316    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
317        let x: u32 = (*self).try_into().map_err(io::Error::other)?;
318        writer.write_all(&x.to_le_bytes())
319    }
320}
321
322// ============ Decode implementation =============
323
324impl Decode for ContinuationVmProof<SC> {
325    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
326        let per_segment = decode_vec(reader)?;
327        let user_public_values = UserPublicValuesProof::decode(reader)?;
328        Ok(Self {
329            per_segment,
330            user_public_values,
331        })
332    }
333}
334
335impl Decode for VmStarkProof<SC> {
336    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
337        let inner = Proof::decode(reader)?;
338        let user_public_values = decode_vec(reader)?;
339        Ok(Self {
340            inner,
341            user_public_values,
342        })
343    }
344}
345
346impl Decode for UserPublicValuesProof<DIGEST_SIZE, F> {
347    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
348        let proof = decode_vec(reader)?;
349        let public_values = decode_vec(reader)?;
350        let public_values_commit = <[F; DIGEST_SIZE]>::decode(reader)?;
351        Ok(Self {
352            proof,
353            public_values,
354            public_values_commit,
355        })
356    }
357}
358
359impl Decode for RootVmVerifierInput<SC> {
360    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
361        let proofs = decode_vec(reader)?;
362        let public_values = decode_vec(reader)?;
363        Ok(Self {
364            proofs,
365            public_values,
366        })
367    }
368}
369
370impl Decode for Proof<SC> {
371    /// Decode a proof using FRI as the PCS with `BabyBearPoseidon2Config`.
372    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
373        let mut version_bytes = [0u8; 4];
374        reader.read_exact(&mut version_bytes)?;
375        let version = u32::from_le_bytes(version_bytes);
376
377        if version != CODEC_VERSION {
378            return Err(io::Error::new(
379                io::ErrorKind::InvalidData,
380                format!(
381                    "Invalid codec version. Expected {}, got {}",
382                    CODEC_VERSION, version
383                ),
384            ));
385        }
386
387        // Decode commitments
388        let main_trace = decode_commitments(reader)?;
389        let after_challenge = decode_commitments(reader)?;
390        let quotient = decode_commitment(reader)?;
391
392        let commitments = Commitments {
393            main_trace,
394            after_challenge,
395            quotient,
396        };
397
398        // Decode OpeningProof
399        let opening = decode_opening_proof(reader)?;
400
401        // Decode per_air data
402        let per_air = decode_vec(reader)?;
403
404        // Decode RAP phase sequence kind
405        let mut kind_byte = [0u8; 1];
406        reader.read_exact(&mut kind_byte)?;
407        if kind_byte[0] != RapPhaseSeqKind::FriLogUp as u8 {
408            return Err(io::Error::new(
409                io::ErrorKind::InvalidData,
410                format!("Unknown RapPhaseSeqKind: {}", kind_byte[0]),
411            ));
412        }
413
414        // Decode logup witness
415        let rap_phase_seq_proof = Option::<FriLogUpPartialProof<F>>::decode(reader)?;
416
417        Ok(Proof {
418            commitments,
419            opening,
420            per_air,
421            rap_phase_seq_proof,
422        })
423    }
424}
425
426fn decode_commitment<R: Read>(reader: &mut R) -> Result<Com<SC>> {
427    let digest = <[F; DIGEST_SIZE]>::decode(reader)?;
428    // Convert [F; DIGEST_SIZE] to Com<SC>
429    Ok(digest.into())
430}
431
432fn decode_commitments<R: Read>(reader: &mut R) -> Result<Vec<Com<SC>>> {
433    let coms_count = usize::decode(reader)?;
434    let mut coms = Vec::with_capacity(coms_count);
435
436    for _ in 0..coms_count {
437        coms.push(decode_commitment(reader)?);
438    }
439
440    Ok(coms)
441}
442
443fn decode_opening_proof<R: Read>(reader: &mut R) -> Result<OpeningProof<PcsProof<SC>, Challenge>> {
444    // Decode FRI proof
445    let proof = InnerFriProof::decode(reader)?;
446    let values = decode_opened_values(reader)?;
447
448    Ok(OpeningProof { proof, values })
449}
450
451fn decode_opened_values<R: Read>(reader: &mut R) -> Result<OpenedValues<Challenge>> {
452    let preprocessed = decode_vec(reader)?;
453
454    let main_count = usize::decode(reader)?;
455    let mut main = Vec::with_capacity(main_count);
456    for _ in 0..main_count {
457        main.push(decode_vec(reader)?);
458    }
459
460    let after_challenge_count = usize::decode(reader)?;
461    let mut after_challenge = Vec::with_capacity(after_challenge_count);
462    for _ in 0..after_challenge_count {
463        after_challenge.push(decode_vec(reader)?);
464    }
465
466    let quotient_count = usize::decode(reader)?;
467    let mut quotient = Vec::with_capacity(quotient_count);
468    for _ in 0..quotient_count {
469        let per_air_count = usize::decode(reader)?;
470        let mut per_air = Vec::with_capacity(per_air_count);
471        for _ in 0..per_air_count {
472            per_air.push(decode_vec(reader)?);
473        }
474        quotient.push(per_air);
475    }
476
477    Ok(OpenedValues {
478        preprocessed,
479        main,
480        after_challenge,
481        quotient,
482    })
483}
484
485impl Decode for AdjacentOpenedValues<Challenge> {
486    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
487        let local = decode_vec(reader)?;
488        let next = decode_vec(reader)?;
489
490        Ok(AdjacentOpenedValues { local, next })
491    }
492}
493
494impl Decode for AirProofData<F, Challenge> {
495    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
496        let air_id = usize::decode(reader)?;
497        let degree = usize::decode(reader)?;
498
499        let exposed_values_count = usize::decode(reader)?;
500        let mut exposed_values_after_challenge = Vec::with_capacity(exposed_values_count);
501        for _ in 0..exposed_values_count {
502            exposed_values_after_challenge.push(decode_vec(reader)?);
503        }
504
505        let public_values = decode_vec(reader)?;
506
507        Ok(AirProofData {
508            air_id,
509            degree,
510            exposed_values_after_challenge,
511            public_values,
512        })
513    }
514}
515
516impl Decode for InnerFriProof {
517    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
518        let commit_phase_commits = decode_commitments(reader)?;
519        let query_proofs = decode_vec(reader)?;
520        let final_poly = decode_vec(reader)?;
521        let pow_witness = F::decode(reader)?;
522
523        Ok(InnerFriProof {
524            commit_phase_commits,
525            query_proofs,
526            final_poly,
527            pow_witness,
528        })
529    }
530}
531
532impl Decode for InnerQueryProof {
533    /// See [InnerQueryProof::encode].
534    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
535        let batch_opening_count = usize::decode(reader)?;
536        let mut input_proof = Vec::with_capacity(batch_opening_count);
537        for _ in 0..batch_opening_count {
538            let opened_values_len = usize::decode(reader)?;
539            let mut opened_values = Vec::with_capacity(opened_values_len);
540            for _ in 0..opened_values_len {
541                opened_values.push(decode_vec(reader)?);
542            }
543            let opening_proof = decode_vec(reader)?;
544
545            let batch_opening = InnerBatchOpening {
546                opened_values,
547                opening_proof,
548            };
549            input_proof.push(batch_opening);
550        }
551
552        let commit_phase_openings_count = usize::decode(reader)?;
553        let mut commit_phase_openings = Vec::with_capacity(commit_phase_openings_count);
554
555        for _ in 0..commit_phase_openings_count {
556            let sibling_value = Challenge::decode(reader)?;
557            let opening_proof = decode_vec(reader)?;
558
559            commit_phase_openings.push(CommitPhaseProofStep {
560                sibling_value,
561                opening_proof,
562            });
563        }
564
565        Ok(InnerQueryProof {
566            input_proof,
567            commit_phase_openings,
568        })
569    }
570}
571
572impl Decode for Option<FriLogUpPartialProof<F>> {
573    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
574        let mut bytes = [0u8; 4];
575        reader.read_exact(&mut bytes)?;
576
577        let value = u32::from_le_bytes(bytes);
578        // When `Option<FriLogUpPartialProof<F>>` is None, it's encoded as `u32::max`.
579        if value == u32::MAX {
580            return Ok(None);
581        }
582
583        // Reconstruct the field element from the u32 value
584        let logup_pow_witness = F::from_canonical_u32(value);
585        Ok(Some(FriLogUpPartialProof { logup_pow_witness }))
586    }
587}
588
589impl Decode for Challenge {
590    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
591        // For a BinomialExtensionField<F, 4>, we need to read 4 F elements
592        let mut base_elements = [F::ZERO; 4];
593        for base_element in &mut base_elements {
594            *base_element = F::decode(reader)?;
595        }
596
597        // Construct the extension field from base elements
598        Ok(Challenge::from_base_slice(&base_elements))
599    }
600}
601
602impl Decode for [F; DIGEST_SIZE] {
603    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
604        let mut result = [F::ZERO; DIGEST_SIZE];
605        for elt in &mut result {
606            *elt = F::decode(reader)?;
607        }
608        Ok(result)
609    }
610}
611
612/// Decodes a vector of elements
613pub(crate) fn decode_vec<T: Decode, R: Read>(reader: &mut R) -> Result<Vec<T>> {
614    let len = usize::decode(reader)?;
615    let mut vec = Vec::with_capacity(len);
616
617    for _ in 0..len {
618        vec.push(T::decode(reader)?);
619    }
620
621    Ok(vec)
622}
623
624impl Decode for F {
625    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
626        let mut bytes = [0u8; 4];
627        reader.read_exact(&mut bytes)?;
628
629        let value = u32::from_le_bytes(bytes);
630        Ok(F::from_canonical_u32(value))
631    }
632}
633
634impl Decode for usize {
635    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
636        let mut bytes = [0u8; 4];
637        reader.read_exact(&mut bytes)?;
638
639        let value = u32::from_le_bytes(bytes);
640        Ok(value as usize)
641    }
642}