openvm_native_recursion/
witness.rs

Help
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, PcsProof},
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_fr::Bn254Fr,
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, Bn254Fr, 1>;
37
38impl Witnessable<C> for Bn254Fr {
39    type WitnessVariable = Var<Bn254Fr>;
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: &[Bn254Fr; 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: &[Bn254Fr; 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 Bn254Fr {}
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<PcsProof<BabyBearPoseidon2RootConfig>, OuterChallenge> {
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        OpeningProofVariable { proof, values }
245    }
246
247    fn write(&self, witness: &mut Witness<C>) {
248        self.proof.write(witness);
249        <_ as Witnessable<C>>::write(&self.values, witness);
250    }
251}
252
253impl Witnessable<C> for OpenedValues<OuterChallenge> {
254    type WitnessVariable = OpenedValuesVariable<C>;
255
256    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
257        let preprocessed = self.preprocessed.read(builder);
258        let main = self.main.read(builder);
259        let quotient = self.quotient.read(builder);
260        let after_challenge = self.after_challenge.read(builder);
261
262        OpenedValuesVariable {
263            preprocessed,
264            main,
265            quotient,
266            after_challenge,
267        }
268    }
269
270    fn write(&self, witness: &mut Witness<C>) {
271        <Vec<_> as Witnessable<C>>::write(&self.preprocessed, witness);
272        <Vec<_> as Witnessable<C>>::write(&self.main, witness);
273        <Vec<_> as Witnessable<C>>::write(&self.quotient, witness);
274        <Vec<_> as Witnessable<C>>::write(&self.after_challenge, witness);
275    }
276}
277
278impl Witnessable<C> for AdjacentOpenedValues<OuterChallenge> {
279    type WitnessVariable = AdjacentOpenedValuesVariable<C>;
280
281    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
282        let local = self.local.read(builder);
283        let next = self.next.read(builder);
284        AdjacentOpenedValuesVariable { local, next }
285    }
286
287    fn write(&self, witness: &mut Witness<C>) {
288        <Vec<_> as Witnessable<C>>::write(&self.local, witness);
289        <Vec<_> as Witnessable<C>>::write(&self.next, witness);
290    }
291}
292impl VectorWitnessable<C> for AdjacentOpenedValues<OuterChallenge> {}
293impl VectorWitnessable<C> for Vec<AdjacentOpenedValues<OuterChallenge>> {}
openvm_native_recursion/witness.rs

openvm_native_recursion/
witness.rs
