openvm_native_recursion/
hints.rs

Help
1use std::cmp::Reverse;
2
3use itertools::Itertools;
4use openvm_native_compiler::ir::{
5    unsafe_array_transmute, Array, ArrayLike, Builder, Config, Ext, Felt, MemVariable, Usize, Var,
6    DIGEST_SIZE,
7};
8use openvm_native_compiler_derive::iter_zip;
9use openvm_stark_backend::{
10    config::{Com, PcsProof},
11    keygen::types::TraceWidth,
12    p3_commit::ExtensionMmcs,
13    p3_field::{extension::BinomialExtensionField, Field, FieldAlgebra, FieldExtensionAlgebra},
14    p3_util::log2_strict_usize,
15    proof::{AdjacentOpenedValues, AirProofData, Commitments, OpenedValues, OpeningProof, Proof},
16};
17use openvm_stark_sdk::{
18    config::baby_bear_poseidon2::BabyBearPoseidon2Config,
19    p3_baby_bear::{BabyBear, Poseidon2BabyBear},
20};
21use p3_fri::{BatchOpening, CommitPhaseProofStep, FriProof, QueryProof};
22use p3_merkle_tree::MerkleTreeMmcs;
23use p3_symmetric::{PaddingFreeSponge, TruncatedPermutation};
24
25use crate::{
26    types::InnerConfig,
27    vars::{
28        AdjacentOpenedValuesVariable, AirProofDataVariable, CommitmentsVariable,
29        OpenedValuesVariable, OpeningProofVariable, StarkProofVariable, TraceWidthVariable,
30    },
31};
32
33pub type InnerVal = BabyBear;
34pub type InnerChallenge = BinomialExtensionField<InnerVal, 4>;
35pub type InnerPerm = Poseidon2BabyBear<16>;
36pub type InnerHash = PaddingFreeSponge<InnerPerm, 16, 8, 8>;
37pub type InnerDigest = [InnerVal; DIGEST_SIZE];
38pub type InnerCompress = TruncatedPermutation<InnerPerm, 2, 8, 16>;
39pub type InnerValMmcs = MerkleTreeMmcs<
40    <InnerVal as Field>::Packing,
41    <InnerVal as Field>::Packing,
42    InnerHash,
43    InnerCompress,
44    8,
45>;
46pub type InnerChallengeMmcs = ExtensionMmcs<InnerVal, InnerChallenge, InnerValMmcs>;
47pub type InnerInputProof = Vec<InnerBatchOpening>;
48pub type InnerQueryProof = QueryProof<InnerChallenge, InnerChallengeMmcs, InnerInputProof>;
49pub type InnerCommitPhaseStep = CommitPhaseProofStep<InnerChallenge, InnerChallengeMmcs>;
50pub type InnerFriProof = FriProof<InnerChallenge, InnerChallengeMmcs, InnerVal, InnerInputProof>;
51pub type InnerBatchOpening = BatchOpening<InnerVal, InnerValMmcs>;
52
53pub trait Hintable<C: Config> {
54    type HintVariable: MemVariable<C>;
55
56    fn read(builder: &mut Builder<C>) -> Self::HintVariable;
57
58    fn write(&self) -> Vec<Vec<C::N>>;
59
60    fn witness(variable: &Self::HintVariable, builder: &mut Builder<C>) {
61        let target = Self::read(builder);
62        builder.assign(variable, target);
63    }
64}
65
66impl<C: Config> Hintable<C> for usize {
67    type HintVariable = Var<C::N>;
68
69    fn read(builder: &mut Builder<C>) -> Self::HintVariable {
70        builder.hint_var()
71    }
72
73    fn write(&self) -> Vec<Vec<C::N>> {
74        vec![vec![FieldAlgebra::from_canonical_usize(*self)]]
75    }
76}
77
78// Assumes F = N
79impl Hintable<InnerConfig> for InnerVal {
80    type HintVariable = Felt<InnerVal>;
81
82    fn read(builder: &mut Builder<InnerConfig>) -> Self::HintVariable {
83        builder.hint_felt()
84    }
85
86    fn write(&self) -> Vec<Vec<<InnerConfig as Config>::N>> {
87        vec![vec![*self]]
88    }
89}
90
91// Assumes F = N
92impl Hintable<InnerConfig> for InnerChallenge {
93    type HintVariable = Ext<InnerVal, InnerChallenge>;
94
95    fn read(builder: &mut Builder<InnerConfig>) -> Self::HintVariable {
96        builder.hint_ext()
97    }
98
99    fn write(&self) -> Vec<Vec<<InnerConfig as Config>::N>> {
100        self.as_base_slice()
101            .iter()
102            .copied()
103            .map(|x| vec![x])
104            .collect()
105    }
106}
107
108/// Implement this on a type `T` that also implements `Hintable<C: Config>`
109/// so that `Hintable<C>` is auto implemented on `Vec<T>`
110pub trait VecAutoHintable {}
111
112impl VecAutoHintable for Vec<usize> {}
113
114impl VecAutoHintable for Vec<InnerVal> {}
115
116impl VecAutoHintable for Vec<Vec<InnerChallenge>> {}
117
118impl VecAutoHintable for AdjacentOpenedValues<InnerChallenge> {}
119
120impl VecAutoHintable for Vec<AdjacentOpenedValues<InnerChallenge>> {}
121
122impl VecAutoHintable for Vec<Vec<AdjacentOpenedValues<InnerChallenge>>> {}
123
124impl VecAutoHintable for AirProofData<InnerVal, InnerChallenge> {}
125impl VecAutoHintable for Proof<BabyBearPoseidon2Config> {}
126
127impl<C: Config, I: VecAutoHintable + Hintable<C>> Hintable<C> for Vec<I> {
128    type HintVariable = Array<C, I::HintVariable>;
129
130    fn read(builder: &mut Builder<C>) -> Self::HintVariable {
131        let len = builder.hint_var();
132        let arr = builder.dyn_array(len);
133        iter_zip!(builder, arr).for_each(|idx_vec, builder| {
134            let hint = I::read(builder);
135            let ptr = idx_vec[0];
136            builder.iter_ptr_set(&arr, ptr, hint);
137        });
138        arr
139    }
140
141    fn write(&self) -> Vec<Vec<<C as Config>::N>> {
142        let mut stream = Vec::new();
143
144        let len = C::N::from_canonical_usize(self.len());
145        stream.push(vec![len]);
146
147        self.iter().for_each(|i| {
148            let comm = I::write(i);
149            stream.extend(comm);
150        });
151
152        stream
153    }
154}
155
156impl Hintable<InnerConfig> for Vec<usize> {
157    type HintVariable = Array<InnerConfig, Var<InnerVal>>;
158
159    fn read(builder: &mut Builder<InnerConfig>) -> Self::HintVariable {
160        builder.hint_vars()
161    }
162
163    fn write(&self) -> Vec<Vec<InnerVal>> {
164        vec![self
165            .iter()
166            .map(|x| InnerVal::from_canonical_usize(*x))
167            .collect()]
168    }
169}
170
171impl<C: Config> Hintable<C> for Vec<u8> {
172    type HintVariable = Array<C, Var<C::N>>;
173
174    fn read(builder: &mut Builder<C>) -> Self::HintVariable {
175        builder.hint_vars()
176    }
177
178    fn write(&self) -> Vec<Vec<C::N>> {
179        vec![self
180            .iter()
181            .map(|x| FieldAlgebra::from_canonical_u8(*x))
182            .collect()]
183    }
184}
185
186impl Hintable<InnerConfig> for Vec<InnerVal> {
187    type HintVariable = Array<InnerConfig, Felt<InnerVal>>;
188
189    fn read(builder: &mut Builder<InnerConfig>) -> Self::HintVariable {
190        builder.hint_felts()
191    }
192
193    fn write(&self) -> Vec<Vec<<InnerConfig as Config>::N>> {
194        vec![self.clone()]
195    }
196}
197
198impl Hintable<InnerConfig> for Vec<InnerChallenge> {
199    type HintVariable = Array<InnerConfig, Ext<InnerVal, InnerChallenge>>;
200
201    fn read(builder: &mut Builder<InnerConfig>) -> Self::HintVariable {
202        builder.hint_exts()
203    }
204
205    fn write(&self) -> Vec<Vec<<InnerConfig as Config>::N>> {
206        vec![
207            vec![InnerVal::from_canonical_usize(self.len())],
208            self.iter()
209                .flat_map(|x| (*x).as_base_slice().to_vec())
210                .collect(),
211        ]
212    }
213}
214
215impl Hintable<InnerConfig> for Vec<Vec<InnerChallenge>> {
216    type HintVariable = Array<InnerConfig, Array<InnerConfig, Ext<InnerVal, InnerChallenge>>>;
217
218    fn read(builder: &mut Builder<InnerConfig>) -> Self::HintVariable {
219        let len = builder.hint_var();
220        let arr = builder.dyn_array(len);
221        iter_zip!(builder, arr).for_each(|idx_vec, builder| {
222            let hint = Vec::<InnerChallenge>::read(builder);
223            builder.iter_ptr_set(&arr, idx_vec[0], hint);
224        });
225        arr
226    }
227
228    fn write(&self) -> Vec<Vec<<InnerConfig as Config>::N>> {
229        let mut stream = Vec::new();
230
231        let len = InnerVal::from_canonical_usize(self.len());
232        stream.push(vec![len]);
233
234        self.iter().for_each(|arr| {
235            let comm = Vec::<InnerChallenge>::write(arr);
236            stream.extend(comm);
237        });
238
239        stream
240    }
241}
242
243impl Hintable<InnerConfig> for TraceWidth {
244    type HintVariable = TraceWidthVariable<InnerConfig>;
245
246    fn read(builder: &mut Builder<InnerConfig>) -> Self::HintVariable {
247        let preprocessed = unsafe_array_transmute(Vec::<usize>::read(builder));
248        let cached_mains = unsafe_array_transmute(Vec::<usize>::read(builder));
249        let common_main = Usize::Var(usize::read(builder));
250        let after_challenge = unsafe_array_transmute(Vec::<usize>::read(builder));
251
252        TraceWidthVariable {
253            preprocessed,
254            cached_mains,
255            common_main,
256            after_challenge,
257        }
258    }
259
260    fn write(&self) -> Vec<Vec<<InnerConfig as Config>::N>> {
261        let mut stream = Vec::new();
262
263        stream.extend(self.preprocessed.into_iter().collect::<Vec<_>>().write());
264        stream.extend(self.cached_mains.write());
265        stream.extend(<usize as Hintable<InnerConfig>>::write(&self.common_main));
266        stream.extend(self.after_challenge.write());
267
268        stream
269    }
270}
271
272impl Hintable<InnerConfig> for Proof<BabyBearPoseidon2Config> {
273    type HintVariable = StarkProofVariable<InnerConfig>;
274
275    fn read(builder: &mut Builder<InnerConfig>) -> Self::HintVariable {
276        let commitments = Commitments::<Com<BabyBearPoseidon2Config>>::read(builder);
277        let opening =
278            OpeningProof::<PcsProof<BabyBearPoseidon2Config>, InnerChallenge>::read(builder);
279        let per_air = Vec::<AirProofData<InnerVal, InnerChallenge>>::read(builder);
280        let raw_air_perm_by_height = Vec::<usize>::read(builder);
281        // A hacky way to transmute from Array of Var to Array of Usize.
282        let air_perm_by_height = unsafe_array_transmute(raw_air_perm_by_height);
283        let log_up_pow_witness = builder.hint_felt();
284
285        StarkProofVariable {
286            commitments,
287            opening,
288            per_air,
289            air_perm_by_height,
290            log_up_pow_witness,
291        }
292    }
293
294    fn write(&self) -> Vec<Vec<<InnerConfig as Config>::N>> {
295        let mut stream = Vec::new();
296
297        stream.extend(self.commitments.write());
298        stream.extend(self.opening.write());
299        stream.extend(<Vec<AirProofData<_, _>> as Hintable<_>>::write(
300            &self.per_air,
301        ));
302        let air_perm_by_height: Vec<_> = (0..self.per_air.len())
303            .sorted_by_key(|i| Reverse(self.per_air[*i].degree))
304            .collect();
305        stream.extend(air_perm_by_height.write());
306        stream.extend(
307            self.rap_phase_seq_proof
308                .as_ref()
309                .map(|p| p.logup_pow_witness)
310                .unwrap_or_default()
311                .write(),
312        );
313
314        stream
315    }
316}
317
318impl Hintable<InnerConfig> for AirProofData<InnerVal, InnerChallenge> {
319    type HintVariable = AirProofDataVariable<InnerConfig>;
320    fn read(builder: &mut Builder<InnerConfig>) -> Self::HintVariable {
321        let air_id = Usize::Var(usize::read(builder));
322        let log_degree = Usize::Var(usize::read(builder));
323        let exposed_values_after_challenge = Vec::<Vec<InnerChallenge>>::read(builder);
324        let public_values = Vec::<InnerVal>::read(builder);
325        Self::HintVariable {
326            air_id,
327            log_degree,
328            exposed_values_after_challenge,
329            public_values,
330        }
331    }
332    fn write(&self) -> Vec<Vec<<InnerConfig as Config>::N>> {
333        let mut stream = Vec::new();
334
335        stream.extend(<usize as Hintable<InnerConfig>>::write(&self.air_id));
336        stream.extend(<usize as Hintable<InnerConfig>>::write(&log2_strict_usize(
337            self.degree,
338        )));
339        stream.extend(self.exposed_values_after_challenge.write());
340        stream.extend(self.public_values.write());
341
342        stream
343    }
344}
345
346impl Hintable<InnerConfig> for OpeningProof<PcsProof<BabyBearPoseidon2Config>, InnerChallenge> {
347    type HintVariable = OpeningProofVariable<InnerConfig>;
348
349    fn read(builder: &mut Builder<InnerConfig>) -> Self::HintVariable {
350        builder.cycle_tracker_start("HintOpeningProof");
351        let proof = InnerFriProof::read(builder);
352        builder.cycle_tracker_end("HintOpeningProof");
353        builder.cycle_tracker_start("HintOpeningValues");
354        let values = OpenedValues::read(builder);
355        builder.cycle_tracker_end("HintOpeningValues");
356
357        OpeningProofVariable { proof, values }
358    }
359
360    fn write(&self) -> Vec<Vec<<InnerConfig as Config>::N>> {
361        let mut stream = Vec::new();
362
363        stream.extend(self.proof.write());
364        stream.extend(self.values.write());
365
366        stream
367    }
368}
369
370impl Hintable<InnerConfig> for OpenedValues<InnerChallenge> {
371    type HintVariable = OpenedValuesVariable<InnerConfig>;
372
373    fn read(builder: &mut Builder<InnerConfig>) -> Self::HintVariable {
374        let preprocessed = Vec::<AdjacentOpenedValues<InnerChallenge>>::read(builder);
375        let main = Vec::<Vec<AdjacentOpenedValues<InnerChallenge>>>::read(builder);
376        let quotient = Vec::<Vec<Vec<InnerChallenge>>>::read(builder);
377        let after_challenge = Vec::<Vec<AdjacentOpenedValues<InnerChallenge>>>::read(builder);
378
379        OpenedValuesVariable {
380            preprocessed,
381            main,
382            quotient,
383            after_challenge,
384        }
385    }
386
387    fn write(&self) -> Vec<Vec<<InnerConfig as Config>::N>> {
388        let mut stream = Vec::new();
389
390        stream.extend(self.preprocessed.write());
391        stream.extend(self.main.write());
392        stream.extend(self.quotient.write());
393        stream.extend(self.after_challenge.write());
394
395        stream
396    }
397}
398
399impl Hintable<InnerConfig> for AdjacentOpenedValues<InnerChallenge> {
400    type HintVariable = AdjacentOpenedValuesVariable<InnerConfig>;
401
402    fn read(builder: &mut Builder<InnerConfig>) -> Self::HintVariable {
403        let local = Vec::<InnerChallenge>::read(builder);
404        let next = Vec::<InnerChallenge>::read(builder);
405        AdjacentOpenedValuesVariable { local, next }
406    }
407
408    fn write(&self) -> Vec<Vec<<InnerConfig as Config>::N>> {
409        let mut stream = Vec::new();
410        stream.extend(self.local.write());
411        stream.extend(self.next.write());
412        stream
413    }
414}
415
416impl Hintable<InnerConfig> for Commitments<Com<BabyBearPoseidon2Config>> {
417    type HintVariable = CommitmentsVariable<InnerConfig>;
418
419    fn read(builder: &mut Builder<InnerConfig>) -> Self::HintVariable {
420        let main_trace = Vec::<InnerDigest>::read(builder);
421        let after_challenge = Vec::<InnerDigest>::read(builder);
422        let quotient = InnerDigest::read(builder);
423
424        CommitmentsVariable {
425            main_trace,
426            after_challenge,
427            quotient,
428        }
429    }
430
431    fn write(&self) -> Vec<Vec<<InnerConfig as Config>::N>> {
432        let mut stream = Vec::new();
433
434        stream.extend(Vec::<InnerDigest>::write(
435            &self.main_trace.iter().map(|&x| x.into()).collect(),
436        ));
437        stream.extend(Vec::<InnerDigest>::write(
438            &self.after_challenge.iter().map(|&x| x.into()).collect(),
439        ));
440        let h: InnerDigest = self.quotient.into();
441        stream.extend(h.write());
442
443        stream
444    }
445}
446
447#[cfg(test)]
448mod test {
449    use openvm_native_circuit::execute_program;
450    use openvm_native_compiler::{
451        asm::AsmBuilder,
452        ir::{Ext, Felt, Var},
453    };
454    use openvm_stark_backend::p3_field::FieldAlgebra;
455
456    use crate::hints::{Hintable, InnerChallenge, InnerVal};
457
458    #[test]
459    fn test_var_array() {
460        let x = vec![
461            InnerVal::from_canonical_usize(1),
462            InnerVal::from_canonical_usize(2),
463            InnerVal::from_canonical_usize(3),
464        ];
465        let stream = Vec::<InnerVal>::write(&x);
466        assert_eq!(stream, vec![x.clone()]);
467
468        let mut builder = AsmBuilder::<InnerVal, InnerChallenge>::default();
469        let arr = Vec::<InnerVal>::read(&mut builder);
470
471        let expected: Var<_> = builder.constant(InnerVal::from_canonical_usize(3));
472        builder.assert_var_eq(arr.len(), expected);
473
474        for (i, &val) in x.iter().enumerate() {
475            let actual = builder.get(&arr, i);
476            let expected: Felt<InnerVal> = builder.constant(val);
477            builder.assert_felt_eq(actual, expected);
478        }
479
480        builder.halt();
481
482        let program = builder.compile_isa();
483        execute_program(program, stream);
484    }
485
486    #[test]
487    fn test_ext_array() {
488        let x = vec![
489            InnerChallenge::from_canonical_usize(1),
490            InnerChallenge::from_canonical_usize(2),
491            InnerChallenge::from_canonical_usize(3),
492        ];
493        let stream = Vec::<InnerChallenge>::write(&x);
494        assert_eq!(
495            stream,
496            vec![
497                vec![InnerVal::from_canonical_usize(x.len())],
498                vec![
499                    InnerVal::from_canonical_usize(1),
500                    InnerVal::from_canonical_usize(0),
501                    InnerVal::from_canonical_usize(0),
502                    InnerVal::from_canonical_usize(0),
503                    InnerVal::from_canonical_usize(2),
504                    InnerVal::from_canonical_usize(0),
505                    InnerVal::from_canonical_usize(0),
506                    InnerVal::from_canonical_usize(0),
507                    InnerVal::from_canonical_usize(3),
508                    InnerVal::from_canonical_usize(0),
509                    InnerVal::from_canonical_usize(0),
510                    InnerVal::from_canonical_usize(0),
511                ],
512            ]
513        );
514
515        let mut builder = AsmBuilder::<InnerVal, InnerChallenge>::default();
516        let arr = Vec::<InnerChallenge>::read(&mut builder);
517
518        let expected: Var<_> = builder.constant(InnerVal::from_canonical_usize(3));
519        builder.assert_var_eq(arr.len(), expected);
520
521        for (i, &val) in x.iter().enumerate() {
522            let actual = builder.get(&arr, i);
523            let expected: Ext<InnerVal, InnerChallenge> = builder.constant(val);
524            builder.assert_ext_eq(actual, expected);
525        }
526
527        builder.halt();
528
529        let program = builder.compile_isa();
530        execute_program(program, stream);
531    }
532}
openvm_native_recursion/hints.rs

openvm_native_recursion/
hints.rs
