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,
13    interaction::{fri_log_up::FriLogUpPartialProof, RapPhaseSeqKind},
14    p3_field::{
15        extension::BinomialExtensionField, BasedVectorSpace, PrimeCharacteristicRing, 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 not correspond to the main openvm version (which may change more frequently).
27const CODEC_VERSION: u32 = 2;
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>(opening: &OpeningProof<SC>, writer: &mut W) -> Result<()> {
132    // Encode FRI proof
133    opening.proof.encode(writer)?;
134    encode_opened_values(&opening.values, writer)?;
135    opening.deep_pow_witness.encode(writer)?;
136    Ok(())
137}
138
139/// [OpenedValues] is a typedef for `Vec<Vec<Vec<Vec<F>>>>` for
140/// - each round
141///   - each matrix
142///     - each point to open at
143///       - evaluations for each column of matrix at that point
144fn encode_opened_values<W: Write>(
145    opened_values: &OpenedValues<Challenge>,
146    writer: &mut W,
147) -> Result<()> {
148    encode_slice(&opened_values.preprocessed, writer)?;
149    opened_values.main.len().encode(writer)?;
150    for part in &opened_values.main {
151        encode_slice(part, writer)?;
152    }
153    opened_values.after_challenge.len().encode(writer)?;
154    for phase in &opened_values.after_challenge {
155        encode_slice(phase, writer)?;
156    }
157    opened_values.quotient.len().encode(writer)?;
158    for per_air in &opened_values.quotient {
159        per_air.len().encode(writer)?;
160        for chunk in per_air {
161            encode_slice(chunk, writer)?;
162        }
163    }
164
165    Ok(())
166}
167
168impl Encode for AdjacentOpenedValues<Challenge> {
169    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
170        encode_slice(&self.local, writer)?;
171        encode_slice(&self.next, writer)?;
172        Ok(())
173    }
174}
175
176impl Encode for AirProofData<F, Challenge> {
177    /// Encodes the struct
178    /// ```
179    /// pub struct OpenedValues<Challenge> {
180    ///     pub preprocessed: Vec<AdjacentOpenedValues<Challenge>>,
181    ///     pub main: Vec<Vec<AdjacentOpenedValues<Challenge>>>,
182    ///     pub after_challenge: Vec<Vec<AdjacentOpenedValues<Challenge>>>,
183    ///     pub quotient: Vec<Vec<Vec<Challenge>>>,
184    /// }
185    /// ```
186    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
187        self.air_id.encode(writer)?;
188        self.degree.encode(writer)?;
189        self.exposed_values_after_challenge.len().encode(writer)?;
190        for exposed_vals in &self.exposed_values_after_challenge {
191            encode_slice(exposed_vals, writer)?;
192        }
193        encode_slice(&self.public_values, writer)?;
194        Ok(())
195    }
196}
197
198// PcsProof<SC> = InnerFriProof where Pcs = TwoAdicFriPcs
199impl Encode for InnerFriProof {
200    /// Encodes the struct
201    /// ```
202    /// pub struct FriProof<Challenge, M: Mmcs<Challenge>> {
203    ///     pub commit_phase_commits: Vec<M::Commitment>,
204    ///     pub commit_pow_witnesses: Vec<F>,
205    ///     pub query_proofs: Vec<QueryProof<Challenge, M, Vec<BatchOpening<F>>>>,
206    ///     pub final_poly: Vec<Challenge>,
207    ///     pub query_pow_witness: F,
208    /// }
209    /// ```
210    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
211        encode_commitments(&self.commit_phase_commits, writer)?;
212        encode_slice(&self.commit_pow_witnesses, writer)?;
213        encode_slice(&self.query_proofs, writer)?;
214        encode_slice(&self.final_poly, writer)?;
215        self.query_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_basis_coefficients_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!("Invalid codec version. Expected {CODEC_VERSION}, got {version}"),
381            ));
382        }
383
384        // Decode commitments
385        let main_trace = decode_commitments(reader)?;
386        let after_challenge = decode_commitments(reader)?;
387        let quotient = decode_commitment(reader)?;
388
389        let commitments = Commitments {
390            main_trace,
391            after_challenge,
392            quotient,
393        };
394
395        // Decode OpeningProof
396        let opening = decode_opening_proof(reader)?;
397
398        // Decode per_air data
399        let per_air = decode_vec(reader)?;
400
401        // Decode RAP phase sequence kind
402        let mut kind_byte = [0u8; 1];
403        reader.read_exact(&mut kind_byte)?;
404        if kind_byte[0] != RapPhaseSeqKind::FriLogUp as u8 {
405            return Err(io::Error::new(
406                io::ErrorKind::InvalidData,
407                format!("Unknown RapPhaseSeqKind: {}", kind_byte[0]),
408            ));
409        }
410
411        // Decode logup witness
412        let rap_phase_seq_proof = Option::<FriLogUpPartialProof<F>>::decode(reader)?;
413
414        Ok(Proof {
415            commitments,
416            opening,
417            per_air,
418            rap_phase_seq_proof,
419        })
420    }
421}
422
423fn decode_commitment<R: Read>(reader: &mut R) -> Result<Com<SC>> {
424    let digest = <[F; DIGEST_SIZE]>::decode(reader)?;
425    // Convert [F; DIGEST_SIZE] to Com<SC>
426    Ok(digest.into())
427}
428
429fn decode_commitments<R: Read>(reader: &mut R) -> Result<Vec<Com<SC>>> {
430    let coms_count = usize::decode(reader)?;
431    let mut coms = Vec::with_capacity(coms_count);
432
433    for _ in 0..coms_count {
434        coms.push(decode_commitment(reader)?);
435    }
436
437    Ok(coms)
438}
439
440fn decode_opening_proof<R: Read>(reader: &mut R) -> Result<OpeningProof<SC>> {
441    // Decode FRI proof
442    let proof = InnerFriProof::decode(reader)?;
443    let values = decode_opened_values(reader)?;
444    let deep_pow_witness = F::decode(reader)?;
445
446    Ok(OpeningProof {
447        proof,
448        values,
449        deep_pow_witness,
450    })
451}
452
453fn decode_opened_values<R: Read>(reader: &mut R) -> Result<OpenedValues<Challenge>> {
454    let preprocessed = decode_vec(reader)?;
455
456    let main_count = usize::decode(reader)?;
457    let mut main = Vec::with_capacity(main_count);
458    for _ in 0..main_count {
459        main.push(decode_vec(reader)?);
460    }
461
462    let after_challenge_count = usize::decode(reader)?;
463    let mut after_challenge = Vec::with_capacity(after_challenge_count);
464    for _ in 0..after_challenge_count {
465        after_challenge.push(decode_vec(reader)?);
466    }
467
468    let quotient_count = usize::decode(reader)?;
469    let mut quotient = Vec::with_capacity(quotient_count);
470    for _ in 0..quotient_count {
471        let per_air_count = usize::decode(reader)?;
472        let mut per_air = Vec::with_capacity(per_air_count);
473        for _ in 0..per_air_count {
474            per_air.push(decode_vec(reader)?);
475        }
476        quotient.push(per_air);
477    }
478
479    Ok(OpenedValues {
480        preprocessed,
481        main,
482        after_challenge,
483        quotient,
484    })
485}
486
487impl Decode for AdjacentOpenedValues<Challenge> {
488    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
489        let local = decode_vec(reader)?;
490        let next = decode_vec(reader)?;
491
492        Ok(AdjacentOpenedValues { local, next })
493    }
494}
495
496impl Decode for AirProofData<F, Challenge> {
497    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
498        let air_id = usize::decode(reader)?;
499        let degree = usize::decode(reader)?;
500
501        let exposed_values_count = usize::decode(reader)?;
502        let mut exposed_values_after_challenge = Vec::with_capacity(exposed_values_count);
503        for _ in 0..exposed_values_count {
504            exposed_values_after_challenge.push(decode_vec(reader)?);
505        }
506
507        let public_values = decode_vec(reader)?;
508
509        Ok(AirProofData {
510            air_id,
511            degree,
512            exposed_values_after_challenge,
513            public_values,
514        })
515    }
516}
517
518impl Decode for InnerFriProof {
519    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
520        let commit_phase_commits = decode_commitments(reader)?;
521        let commit_pow_witnesses = decode_vec(reader)?;
522        let query_proofs = decode_vec(reader)?;
523        let final_poly = decode_vec(reader)?;
524        let query_pow_witness = F::decode(reader)?;
525
526        Ok(InnerFriProof {
527            commit_phase_commits,
528            commit_pow_witnesses,
529            query_proofs,
530            final_poly,
531            query_pow_witness,
532        })
533    }
534}
535
536impl Decode for InnerQueryProof {
537    /// See [InnerQueryProof::encode].
538    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
539        let batch_opening_count = usize::decode(reader)?;
540        let mut input_proof = Vec::with_capacity(batch_opening_count);
541        for _ in 0..batch_opening_count {
542            let opened_values_len = usize::decode(reader)?;
543            let mut opened_values = Vec::with_capacity(opened_values_len);
544            for _ in 0..opened_values_len {
545                opened_values.push(decode_vec(reader)?);
546            }
547            let opening_proof = decode_vec(reader)?;
548
549            let batch_opening = InnerBatchOpening {
550                opened_values,
551                opening_proof,
552            };
553            input_proof.push(batch_opening);
554        }
555
556        let commit_phase_openings_count = usize::decode(reader)?;
557        let mut commit_phase_openings = Vec::with_capacity(commit_phase_openings_count);
558
559        for _ in 0..commit_phase_openings_count {
560            let sibling_value = Challenge::decode(reader)?;
561            let opening_proof = decode_vec(reader)?;
562
563            commit_phase_openings.push(CommitPhaseProofStep {
564                sibling_value,
565                opening_proof,
566            });
567        }
568
569        Ok(InnerQueryProof {
570            input_proof,
571            commit_phase_openings,
572        })
573    }
574}
575
576impl Decode for Option<FriLogUpPartialProof<F>> {
577    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
578        let mut bytes = [0u8; 4];
579        reader.read_exact(&mut bytes)?;
580
581        let value = u32::from_le_bytes(bytes);
582        // When `Option<FriLogUpPartialProof<F>>` is None, it's encoded as `u32::max`.
583        if value == u32::MAX {
584            return Ok(None);
585        }
586
587        // Reconstruct the field element from the u32 value
588        let logup_pow_witness = F::from_u32(value);
589        Ok(Some(FriLogUpPartialProof { logup_pow_witness }))
590    }
591}
592
593impl Decode for Challenge {
594    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
595        // For a BinomialExtensionField<F, 4>, we need to read 4 F elements
596        let mut base_elements = [F::ZERO; 4];
597        for base_element in &mut base_elements {
598            *base_element = F::decode(reader)?;
599        }
600
601        // Construct the extension field from base elements
602        Ok(Challenge::from_basis_coefficients_slice(&base_elements).unwrap())
603    }
604}
605
606impl Decode for [F; DIGEST_SIZE] {
607    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
608        let mut result = [F::ZERO; DIGEST_SIZE];
609        for elt in &mut result {
610            *elt = F::decode(reader)?;
611        }
612        Ok(result)
613    }
614}
615
616/// Decodes a vector of elements
617pub(crate) fn decode_vec<T: Decode, R: Read>(reader: &mut R) -> Result<Vec<T>> {
618    let len = usize::decode(reader)?;
619    let mut vec = Vec::with_capacity(len);
620
621    for _ in 0..len {
622        vec.push(T::decode(reader)?);
623    }
624
625    Ok(vec)
626}
627
628impl Decode for F {
629    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
630        let mut bytes = [0u8; 4];
631        reader.read_exact(&mut bytes)?;
632
633        let value = u32::from_le_bytes(bytes);
634        Ok(F::from_u32(value))
635    }
636}
637
638impl Decode for usize {
639    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
640        let mut bytes = [0u8; 4];
641        reader.read_exact(&mut bytes)?;
642
643        let value = u32::from_le_bytes(bytes);
644        Ok(value as usize)
645    }
646}