openvm_native_recursion/
hints.rs

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

openvm_native_recursion/
hints.rs
