openvm_native_recursion/
witness.rs

1use core::borrow::Borrow;
2
3use openvm_native_compiler::ir::{
4    Array, Builder, Config, Ext, Felt, MemVariable, Usize, Var, Witness,
5};
6use openvm_stark_backend::{
7    config::Com,
8    p3_util::log2_strict_usize,
9    proof::{AdjacentOpenedValues, AirProofData, Commitments, OpenedValues, OpeningProof, Proof},
10};
11use openvm_stark_sdk::{
12    config::baby_bear_poseidon2_root::BabyBearPoseidon2RootConfig, p3_baby_bear::BabyBear,
13    p3_bn254::Bn254,
14};
15use p3_symmetric::Hash;
16
17use crate::{
18    config::outer::{OuterChallenge, OuterConfig, OuterVal},
19    digest::{DigestVal, DigestVariable},
20    hints::{InnerChallenge, InnerVal},
21    vars::{
22        AdjacentOpenedValuesVariable, AirProofDataVariable, CommitmentsVariable,
23        OpenedValuesVariable, OpeningProofVariable, StarkProofVariable,
24    },
25};
26
27pub trait Witnessable<C: Config> {
28    type WitnessVariable: MemVariable<C>;
29
30    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable;
31
32    fn write(&self, witness: &mut Witness<C>);
33}
34
35type C = OuterConfig;
36type OuterCom = Hash<BabyBear, Bn254, 1>;
37
38impl Witnessable<C> for Bn254 {
39    type WitnessVariable = Var<Bn254>;
40
41    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
42        builder.witness_var()
43    }
44
45    fn write(&self, witness: &mut Witness<C>) {
46        witness.vars.push(*self);
47    }
48}
49
50impl Witnessable<C> for OuterVal {
51    type WitnessVariable = Felt<OuterVal>;
52
53    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
54        builder.witness_felt()
55    }
56
57    fn write(&self, witness: &mut Witness<C>) {
58        witness.felts.push(*self);
59    }
60}
61
62impl Witnessable<C> for OuterChallenge {
63    type WitnessVariable = Ext<OuterVal, OuterChallenge>;
64
65    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
66        builder.witness_ext()
67    }
68
69    fn write(&self, witness: &mut Witness<C>) {
70        witness.exts.push(*self);
71    }
72}
73
74impl Witnessable<C> for OuterCom {
75    type WitnessVariable = DigestVariable<C>;
76
77    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
78        let bv: &[Bn254; 1] = self.borrow();
79        let v = vec![bv[0].read(builder)];
80        DigestVariable::Var(builder.vec(v))
81    }
82
83    fn write(&self, witness: &mut Witness<C>) {
84        let bv: &[Bn254; 1] = self.borrow();
85        witness.vars.push(bv[0]);
86    }
87}
88
89// In static mode, usize is hardcoded.
90impl Witnessable<C> for usize {
91    type WitnessVariable = Usize<<C as Config>::N>;
92
93    fn read(&self, _builder: &mut Builder<C>) -> Self::WitnessVariable {
94        Usize::from(*self)
95    }
96
97    fn write(&self, _witness: &mut Witness<C>) {
98        // Do nothing
99    }
100}
101
102pub trait VectorWitnessable<C: Config>: Witnessable<C> {}
103impl VectorWitnessable<C> for Bn254 {}
104impl VectorWitnessable<C> for OuterVal {}
105impl VectorWitnessable<C> for OuterChallenge {}
106impl VectorWitnessable<C> for OuterCom {}
107impl VectorWitnessable<C> for usize {}
108impl VectorWitnessable<C> for Vec<OuterChallenge> {}
109impl VectorWitnessable<C> for Vec<Vec<OuterChallenge>> {}
110impl VectorWitnessable<C> for Vec<OuterVal> {}
111impl VectorWitnessable<C> for Vec<Vec<OuterVal>> {}
112
113impl<I: VectorWitnessable<C>> Witnessable<C> for Vec<I> {
114    type WitnessVariable = Array<C, I::WitnessVariable>;
115
116    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
117        let raw_vec: Vec<I::WitnessVariable> = self.iter().map(|x| x.read(builder)).collect();
118        builder.vec(raw_vec)
119    }
120
121    fn write(&self, witness: &mut Witness<C>) {
122        self.iter().for_each(|x| x.write(witness));
123    }
124}
125
126impl Witnessable<C> for DigestVal<C> {
127    type WitnessVariable = DigestVariable<C>;
128
129    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
130        let result = vec![builder.witness_var()];
131        DigestVariable::Var(builder.vec(result))
132    }
133
134    fn write(&self, witness: &mut Witness<C>) {
135        if let DigestVal::N(v) = self {
136            assert_eq!(v.len(), 1);
137            witness.vars.push(v[0]);
138        } else {
139            panic!("DigestVal should use N in static mode")
140        }
141    }
142}
143impl VectorWitnessable<C> for DigestVal<C> {}
144
145impl VectorWitnessable<C> for AirProofData<InnerVal, InnerChallenge> {}
146
147impl Witnessable<C> for Proof<BabyBearPoseidon2RootConfig> {
148    type WitnessVariable = StarkProofVariable<C>;
149
150    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
151        if builder.flags.static_only {
152            // Check that the trace heights are sorted
153            assert!(
154                self.per_air.windows(2).all(|w| w[0].degree >= w[1].degree),
155                "Static verifier requires trace heights to be sorted descending"
156            );
157        }
158        let commitments = self.commitments.read(builder);
159        let opening = self.opening.read(builder);
160        let per_air = self.per_air.read(builder);
161        // This reads nothing because air_perm_by_height is a constant.
162        let air_perm_by_height = builder.array(0);
163        let log_up_pow_witness = self
164            .rap_phase_seq_proof
165            .as_ref()
166            .map(|proof| proof.logup_pow_witness)
167            .unwrap_or_default()
168            .read(builder);
169
170        StarkProofVariable {
171            commitments,
172            opening,
173            per_air,
174            air_perm_by_height,
175            log_up_pow_witness,
176        }
177    }
178
179    fn write(&self, witness: &mut Witness<C>) {
180        self.commitments.write(witness);
181        self.opening.write(witness);
182        self.per_air.write(witness);
183        // air_perm_by_height is a constant so we write nothing.
184        let logup_pow_witness = self
185            .rap_phase_seq_proof
186            .as_ref()
187            .map(|p| p.logup_pow_witness)
188            .unwrap_or_default();
189        logup_pow_witness.write(witness);
190    }
191}
192
193impl Witnessable<C> for AirProofData<OuterVal, OuterChallenge> {
194    type WitnessVariable = AirProofDataVariable<C>;
195    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
196        // air_id is constant, skip
197        let air_id = Usize::from(0);
198        // log_degree is constant, skip
199        let log_degree = Usize::from(log2_strict_usize(self.degree));
200        let exposed_values_after_challenge = self.exposed_values_after_challenge.read(builder);
201        let public_values = self.public_values.read(builder);
202        Self::WitnessVariable {
203            air_id,
204            log_degree,
205            exposed_values_after_challenge,
206            public_values,
207        }
208    }
209    fn write(&self, witness: &mut Witness<C>) {
210        // air_id is constant, skip
211        // log_degree is constant, skip
212        <_ as Witnessable<_>>::write(&self.exposed_values_after_challenge, witness);
213        <_ as Witnessable<_>>::write(&self.public_values, witness);
214    }
215}
216
217impl Witnessable<OuterConfig> for Commitments<Com<BabyBearPoseidon2RootConfig>> {
218    type WitnessVariable = CommitmentsVariable<OuterConfig>;
219
220    fn read(&self, builder: &mut Builder<OuterConfig>) -> Self::WitnessVariable {
221        let after_challenge = self.after_challenge.read(builder);
222        let main_trace = self.main_trace.read(builder);
223        let quotient = self.quotient.read(builder);
224        Self::WitnessVariable {
225            after_challenge,
226            main_trace,
227            quotient,
228        }
229    }
230
231    fn write(&self, witness: &mut Witness<OuterConfig>) {
232        self.after_challenge.write(witness);
233        self.main_trace.write(witness);
234        self.quotient.write(witness);
235    }
236}
237
238impl Witnessable<C> for OpeningProof<BabyBearPoseidon2RootConfig> {
239    type WitnessVariable = OpeningProofVariable<C>;
240
241    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
242        let proof = self.proof.read(builder);
243        let values = self.values.read(builder);
244        let deep_pow_witness = self.deep_pow_witness.read(builder);
245        OpeningProofVariable {
246            proof,
247            values,
248            deep_pow_witness,
249        }
250    }
251
252    fn write(&self, witness: &mut Witness<C>) {
253        self.proof.write(witness);
254        <_ as Witnessable<C>>::write(&self.values, witness);
255        self.deep_pow_witness.write(witness);
256    }
257}
258
259impl Witnessable<C> for OpenedValues<OuterChallenge> {
260    type WitnessVariable = OpenedValuesVariable<C>;
261
262    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
263        let preprocessed = self.preprocessed.read(builder);
264        let main = self.main.read(builder);
265        let quotient = self.quotient.read(builder);
266        let after_challenge = self.after_challenge.read(builder);
267
268        OpenedValuesVariable {
269            preprocessed,
270            main,
271            quotient,
272            after_challenge,
273        }
274    }
275
276    fn write(&self, witness: &mut Witness<C>) {
277        <Vec<_> as Witnessable<C>>::write(&self.preprocessed, witness);
278        <Vec<_> as Witnessable<C>>::write(&self.main, witness);
279        <Vec<_> as Witnessable<C>>::write(&self.quotient, witness);
280        <Vec<_> as Witnessable<C>>::write(&self.after_challenge, witness);
281    }
282}
283
284impl Witnessable<C> for AdjacentOpenedValues<OuterChallenge> {
285    type WitnessVariable = AdjacentOpenedValuesVariable<C>;
286
287    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
288        let local = self.local.read(builder);
289        let next = self.next.read(builder);
290        AdjacentOpenedValuesVariable { local, next }
291    }
292
293    fn write(&self, witness: &mut Witness<C>) {
294        <Vec<_> as Witnessable<C>>::write(&self.local, witness);
295        <Vec<_> as Witnessable<C>>::write(&self.next, witness);
296    }
297}
298impl VectorWitnessable<C> for AdjacentOpenedValues<OuterChallenge> {}
299impl VectorWitnessable<C> for Vec<AdjacentOpenedValues<OuterChallenge>> {}