openvm_sdk/
codec.rs

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

openvm_sdk/
codec.rs
