sp1_recursion_compiler/ir/
types.rs

1use alloc::format;
2
3use p3_field::{AbstractExtensionField, AbstractField, ExtensionField, Field};
4use serde::{Deserialize, Serialize};
5
6use super::{
7    Builder, Config, DslIr, ExtConst, ExtHandle, FeltHandle, FromConstant, MemIndex, MemVariable,
8    Ptr, SymbolicExt, SymbolicFelt, SymbolicUsize, SymbolicVar, VarHandle, Variable,
9};
10
11/// A variable that represents a native field element.
12///
13/// Used for counters, simple loops, etc.
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
15pub struct Var<N> {
16    pub idx: u32,
17    pub(crate) handle: *mut VarHandle<N>,
18}
19
20/// A variable that represents an emulated field element.
21///
22/// Used to do field arithmetic for recursive verification.
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
24pub struct Felt<F> {
25    pub idx: u32,
26    pub(crate) handle: *mut FeltHandle<F>,
27}
28
29/// A variable that represents an emulated extension field element.
30///
31/// Used to do extension field arithmetic for recursive verification.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
33pub struct Ext<F, EF> {
34    pub idx: u32,
35    pub(crate) handle: *mut ExtHandle<F, EF>,
36}
37
38unsafe impl<N> Send for Var<N> {}
39unsafe impl<F, EF> Send for Ext<F, EF> {}
40unsafe impl<F> Send for Felt<F> {}
41
42unsafe impl<N> Sync for Var<N> {}
43unsafe impl<F, EF> Sync for Ext<F, EF> {}
44unsafe impl<F> Sync for Felt<F> {}
45
46/// A variable that represents either a constant or variable counter.
47#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
48pub enum Usize<N> {
49    Const(usize),
50    Var(Var<N>),
51}
52
53#[derive(Debug, Clone, Default, Serialize, Deserialize)]
54pub struct Witness<C: Config> {
55    pub vars: Vec<C::N>,
56    pub felts: Vec<C::F>,
57    pub exts: Vec<C::EF>,
58    pub vkey_hash: C::N,
59    pub committed_values_digest: C::N,
60}
61
62impl<C: Config> Witness<C> {
63    pub fn size(&self) -> usize {
64        self.vars.len() + self.felts.len() + self.exts.len() + 2
65    }
66
67    pub fn write_vkey_hash(&mut self, vkey_hash: C::N) {
68        self.vars.push(vkey_hash);
69        self.vkey_hash = vkey_hash;
70    }
71
72    pub fn write_committed_values_digest(&mut self, committed_values_digest: C::N) {
73        self.vars.push(committed_values_digest);
74        self.committed_values_digest = committed_values_digest
75    }
76}
77
78impl<N: Field> Usize<N> {
79    pub fn value(&self) -> usize {
80        match self {
81            Usize::Const(c) => *c,
82            Usize::Var(_) => panic!("Cannot get the value of a variable"),
83        }
84    }
85
86    pub fn materialize<C: Config<N = N>>(&self, builder: &mut Builder<C>) -> Var<C::N> {
87        match self {
88            Usize::Const(c) => builder.eval(C::N::from_canonical_usize(*c)),
89            Usize::Var(v) => *v,
90        }
91    }
92}
93
94impl<N> From<Var<N>> for Usize<N> {
95    fn from(v: Var<N>) -> Self {
96        Usize::Var(v)
97    }
98}
99
100impl<N> From<usize> for Usize<N> {
101    fn from(c: usize) -> Self {
102        Usize::Const(c)
103    }
104}
105
106impl<N> Var<N> {
107    pub const fn new(idx: u32, handle: *mut VarHandle<N>) -> Self {
108        Self { idx, handle }
109    }
110
111    pub fn id(&self) -> String {
112        format!("var{}", self.idx)
113    }
114
115    pub fn loc(&self) -> String {
116        self.idx.to_string()
117    }
118}
119
120impl<F> Felt<F> {
121    pub const fn new(id: u32, handle: *mut FeltHandle<F>) -> Self {
122        Self { idx: id, handle }
123    }
124
125    pub fn id(&self) -> String {
126        format!("felt{}", self.idx)
127    }
128
129    pub fn loc(&self) -> String {
130        self.idx.to_string()
131    }
132
133    pub fn inverse(&self) -> SymbolicFelt<F>
134    where
135        F: Field,
136    {
137        SymbolicFelt::<F>::one() / *self
138    }
139}
140
141impl<F, EF> Ext<F, EF> {
142    pub const fn new(id: u32, handle: *mut ExtHandle<F, EF>) -> Self {
143        Self { idx: id, handle }
144    }
145
146    pub fn id(&self) -> String {
147        format!("ext{}", self.idx)
148    }
149
150    pub fn loc(&self) -> String {
151        self.idx.to_string()
152    }
153
154    pub fn inverse(&self) -> SymbolicExt<F, EF>
155    where
156        F: Field,
157        EF: ExtensionField<F>,
158    {
159        SymbolicExt::<F, EF>::one() / *self
160    }
161}
162
163impl<C: Config> Variable<C> for Usize<C::N> {
164    type Expression = SymbolicUsize<C::N>;
165
166    fn uninit(builder: &mut Builder<C>) -> Self {
167        builder.uninit::<Var<C::N>>().into()
168    }
169
170    fn assign(&self, src: Self::Expression, builder: &mut Builder<C>) {
171        match self {
172            Usize::Const(_) => {
173                panic!("cannot assign to a constant usize")
174            }
175            Usize::Var(v) => match src {
176                SymbolicUsize::Const(src) => {
177                    builder.assign(*v, C::N::from_canonical_usize(src));
178                }
179                SymbolicUsize::Var(src) => {
180                    builder.assign(*v, src);
181                }
182            },
183        }
184    }
185
186    fn assert_eq(
187        lhs: impl Into<Self::Expression>,
188        rhs: impl Into<Self::Expression>,
189        builder: &mut Builder<C>,
190    ) {
191        let lhs = lhs.into();
192        let rhs = rhs.into();
193
194        match (lhs, rhs) {
195            (SymbolicUsize::Const(lhs), SymbolicUsize::Const(rhs)) => {
196                assert_eq!(lhs, rhs, "constant usizes do not match");
197            }
198            (SymbolicUsize::Const(lhs), SymbolicUsize::Var(rhs)) => {
199                builder.assert_var_eq(C::N::from_canonical_usize(lhs), rhs);
200            }
201            (SymbolicUsize::Var(lhs), SymbolicUsize::Const(rhs)) => {
202                builder.assert_var_eq(lhs, C::N::from_canonical_usize(rhs));
203            }
204            (SymbolicUsize::Var(lhs), SymbolicUsize::Var(rhs)) => builder.assert_var_eq(lhs, rhs),
205        }
206    }
207
208    fn assert_ne(
209        lhs: impl Into<Self::Expression>,
210        rhs: impl Into<Self::Expression>,
211        builder: &mut Builder<C>,
212    ) {
213        let lhs = lhs.into();
214        let rhs = rhs.into();
215
216        match (lhs, rhs) {
217            (SymbolicUsize::Const(lhs), SymbolicUsize::Const(rhs)) => {
218                assert_ne!(lhs, rhs, "constant usizes do not match");
219            }
220            (SymbolicUsize::Const(lhs), SymbolicUsize::Var(rhs)) => {
221                builder.assert_var_ne(C::N::from_canonical_usize(lhs), rhs);
222            }
223            (SymbolicUsize::Var(lhs), SymbolicUsize::Const(rhs)) => {
224                builder.assert_var_ne(lhs, C::N::from_canonical_usize(rhs));
225            }
226            (SymbolicUsize::Var(lhs), SymbolicUsize::Var(rhs)) => {
227                builder.assert_var_ne(lhs, rhs);
228            }
229        }
230    }
231}
232
233impl<C: Config> Variable<C> for Var<C::N> {
234    type Expression = SymbolicVar<C::N>;
235
236    fn uninit(builder: &mut Builder<C>) -> Self {
237        let id = builder.variable_count();
238        let var = Var::new(id, builder.var_handle.as_mut());
239        builder.inner.get_mut().variable_count += 1;
240        var
241    }
242
243    fn assign(&self, src: Self::Expression, builder: &mut Builder<C>) {
244        match src {
245            SymbolicVar::Const(src) => {
246                builder.push_op(DslIr::ImmV(*self, src));
247            }
248            SymbolicVar::Val(src) => {
249                builder.push_op(DslIr::AddVI(*self, src, C::N::zero()));
250            }
251        }
252    }
253
254    fn assert_eq(
255        lhs: impl Into<Self::Expression>,
256        rhs: impl Into<Self::Expression>,
257        builder: &mut Builder<C>,
258    ) {
259        let lhs = lhs.into();
260        let rhs = rhs.into();
261
262        match (lhs, rhs) {
263            (SymbolicVar::Const(lhs), SymbolicVar::Const(rhs)) => {
264                assert_eq!(lhs, rhs, "Assertion failed at compile time");
265            }
266            (SymbolicVar::Const(lhs), SymbolicVar::Val(rhs)) => {
267                builder.push_traced_op(DslIr::AssertEqVI(rhs, lhs));
268            }
269            (SymbolicVar::Val(lhs), SymbolicVar::Const(rhs)) => {
270                builder.push_traced_op(DslIr::AssertEqVI(lhs, rhs));
271            }
272            (SymbolicVar::Val(lhs), SymbolicVar::Val(rhs)) => {
273                builder.push_traced_op(DslIr::AssertEqV(lhs, rhs));
274            }
275        }
276    }
277
278    fn assert_ne(
279        lhs: impl Into<Self::Expression>,
280        rhs: impl Into<Self::Expression>,
281        builder: &mut Builder<C>,
282    ) {
283        let lhs = lhs.into();
284        let rhs = rhs.into();
285
286        match (lhs, rhs) {
287            (SymbolicVar::Const(lhs), SymbolicVar::Const(rhs)) => {
288                assert_ne!(lhs, rhs, "Assertion failed at compile time");
289            }
290            (SymbolicVar::Const(lhs), SymbolicVar::Val(rhs)) => {
291                builder.push_traced_op(DslIr::AssertNeVI(rhs, lhs));
292            }
293            (SymbolicVar::Val(lhs), SymbolicVar::Const(rhs)) => {
294                builder.push_traced_op(DslIr::AssertNeVI(lhs, rhs));
295            }
296            (SymbolicVar::Val(lhs), SymbolicVar::Val(rhs)) => {
297                builder.push_traced_op(DslIr::AssertNeV(lhs, rhs));
298            }
299        }
300    }
301}
302
303impl<C: Config> MemVariable<C> for Var<C::N> {
304    fn size_of() -> usize {
305        1
306    }
307
308    fn load(&self, ptr: Ptr<C::N>, index: MemIndex<C::N>, builder: &mut Builder<C>) {
309        builder.push_op(DslIr::LoadV(*self, ptr, index));
310    }
311
312    fn store(&self, ptr: Ptr<<C as Config>::N>, index: MemIndex<C::N>, builder: &mut Builder<C>) {
313        builder.push_op(DslIr::StoreV(*self, ptr, index));
314    }
315}
316
317impl<C: Config> Variable<C> for Felt<C::F> {
318    type Expression = SymbolicFelt<C::F>;
319
320    fn uninit(builder: &mut Builder<C>) -> Self {
321        let idx = builder.variable_count();
322        let felt = Felt::<C::F>::new(idx, builder.felt_handle.as_mut());
323        builder.inner.get_mut().variable_count += 1;
324        felt
325    }
326
327    fn assign(&self, src: Self::Expression, builder: &mut Builder<C>) {
328        match src {
329            SymbolicFelt::Const(src) => {
330                builder.push_op(DslIr::ImmF(*self, src));
331            }
332            SymbolicFelt::Val(src) => {
333                builder.push_op(DslIr::AddFI(*self, src, C::F::zero()));
334            }
335        }
336    }
337
338    fn assert_eq(
339        lhs: impl Into<Self::Expression>,
340        rhs: impl Into<Self::Expression>,
341        builder: &mut Builder<C>,
342    ) {
343        let lhs = lhs.into();
344        let rhs = rhs.into();
345
346        match (lhs, rhs) {
347            (SymbolicFelt::Const(lhs), SymbolicFelt::Const(rhs)) => {
348                assert_eq!(lhs, rhs, "Assertion failed at compile time");
349            }
350            (SymbolicFelt::Const(lhs), SymbolicFelt::Val(rhs)) => {
351                builder.push_traced_op(DslIr::AssertEqFI(rhs, lhs));
352            }
353            (SymbolicFelt::Val(lhs), SymbolicFelt::Const(rhs)) => {
354                builder.push_traced_op(DslIr::AssertEqFI(lhs, rhs));
355            }
356            (SymbolicFelt::Val(lhs), SymbolicFelt::Val(rhs)) => {
357                builder.push_traced_op(DslIr::AssertEqF(lhs, rhs));
358            }
359        }
360    }
361
362    fn assert_ne(
363        lhs: impl Into<Self::Expression>,
364        rhs: impl Into<Self::Expression>,
365        builder: &mut Builder<C>,
366    ) {
367        let lhs = lhs.into();
368        let rhs = rhs.into();
369
370        match (lhs, rhs) {
371            (SymbolicFelt::Const(lhs), SymbolicFelt::Const(rhs)) => {
372                assert_ne!(lhs, rhs, "Assertion failed at compile time");
373            }
374            (SymbolicFelt::Const(lhs), SymbolicFelt::Val(rhs)) => {
375                builder.push_traced_op(DslIr::AssertNeFI(rhs, lhs));
376            }
377            (SymbolicFelt::Val(lhs), SymbolicFelt::Const(rhs)) => {
378                builder.push_traced_op(DslIr::AssertNeFI(lhs, rhs));
379            }
380            (SymbolicFelt::Val(lhs), SymbolicFelt::Val(rhs)) => {
381                builder.push_traced_op(DslIr::AssertNeF(lhs, rhs));
382            }
383        }
384    }
385}
386
387impl<C: Config> MemVariable<C> for Felt<C::F> {
388    fn size_of() -> usize {
389        1
390    }
391
392    fn load(&self, ptr: Ptr<C::N>, index: MemIndex<C::N>, builder: &mut Builder<C>) {
393        builder.push_op(DslIr::LoadF(*self, ptr, index));
394    }
395
396    fn store(&self, ptr: Ptr<<C as Config>::N>, index: MemIndex<C::N>, builder: &mut Builder<C>) {
397        builder.push_op(DslIr::StoreF(*self, ptr, index));
398    }
399}
400
401impl<C: Config> Variable<C> for Ext<C::F, C::EF> {
402    type Expression = SymbolicExt<C::F, C::EF>;
403
404    fn uninit(builder: &mut Builder<C>) -> Self {
405        let idx = builder.variable_count();
406        let ext = Ext::<C::F, C::EF>::new(idx, builder.ext_handle.as_mut());
407        builder.inner.get_mut().variable_count += 1;
408        ext
409    }
410
411    fn assign(&self, src: Self::Expression, builder: &mut Builder<C>) {
412        match src {
413            SymbolicExt::Const(src) => {
414                builder.push_op(DslIr::ImmE(*self, src));
415            }
416            SymbolicExt::Base(src) => match src {
417                SymbolicFelt::Const(src) => {
418                    builder.push_op(DslIr::ImmE(*self, C::EF::from_base(src)));
419                }
420                SymbolicFelt::Val(src) => {
421                    builder.push_op(DslIr::AddEFFI(*self, src, C::EF::zero()));
422                }
423            },
424            SymbolicExt::Val(src) => {
425                builder.push_op(DslIr::AddEI(*self, src, C::EF::zero()));
426            }
427        }
428    }
429
430    fn assert_eq(
431        lhs: impl Into<Self::Expression>,
432        rhs: impl Into<Self::Expression>,
433        builder: &mut Builder<C>,
434    ) {
435        let lhs = lhs.into();
436        let rhs = rhs.into();
437
438        match (lhs, rhs) {
439            (SymbolicExt::Const(lhs), SymbolicExt::Const(rhs)) => {
440                assert_eq!(lhs, rhs, "Assertion failed at compile time");
441            }
442            (SymbolicExt::Const(lhs), SymbolicExt::Val(rhs)) => {
443                builder.push_traced_op(DslIr::AssertEqEI(rhs, lhs));
444            }
445            (SymbolicExt::Const(lhs), rhs) => {
446                let rhs_value = Self::uninit(builder);
447                rhs_value.assign(rhs, builder);
448                builder.push_traced_op(DslIr::AssertEqEI(rhs_value, lhs));
449            }
450            (SymbolicExt::Val(lhs), SymbolicExt::Const(rhs)) => {
451                builder.push_traced_op(DslIr::AssertEqEI(lhs, rhs));
452            }
453            (SymbolicExt::Val(lhs), SymbolicExt::Val(rhs)) => {
454                builder.push_traced_op(DslIr::AssertEqE(lhs, rhs));
455            }
456            (SymbolicExt::Val(lhs), rhs) => {
457                let rhs_value = Self::uninit(builder);
458                rhs_value.assign(rhs, builder);
459                builder.push_traced_op(DslIr::AssertEqE(lhs, rhs_value));
460            }
461            (lhs, rhs) => {
462                let lhs_value = Self::uninit(builder);
463                lhs_value.assign(lhs, builder);
464                let rhs_value = Self::uninit(builder);
465                rhs_value.assign(rhs, builder);
466                builder.push_traced_op(DslIr::AssertEqE(lhs_value, rhs_value));
467            }
468        }
469    }
470
471    fn assert_ne(
472        lhs: impl Into<Self::Expression>,
473        rhs: impl Into<Self::Expression>,
474        builder: &mut Builder<C>,
475    ) {
476        let lhs = lhs.into();
477        let rhs = rhs.into();
478
479        match (lhs, rhs) {
480            (SymbolicExt::Const(lhs), SymbolicExt::Const(rhs)) => {
481                assert_ne!(lhs, rhs, "Assertion failed at compile time");
482            }
483            (SymbolicExt::Const(lhs), SymbolicExt::Val(rhs)) => {
484                builder.push_traced_op(DslIr::AssertNeEI(rhs, lhs));
485            }
486            (SymbolicExt::Const(lhs), rhs) => {
487                let rhs_value = Self::uninit(builder);
488                rhs_value.assign(rhs, builder);
489                builder.push_traced_op(DslIr::AssertNeEI(rhs_value, lhs));
490            }
491            (SymbolicExt::Val(lhs), SymbolicExt::Const(rhs)) => {
492                builder.push_traced_op(DslIr::AssertNeEI(lhs, rhs));
493            }
494            (SymbolicExt::Val(lhs), SymbolicExt::Val(rhs)) => {
495                builder.push_traced_op(DslIr::AssertNeE(lhs, rhs));
496            }
497            (SymbolicExt::Val(lhs), rhs) => {
498                let rhs_value = Self::uninit(builder);
499                rhs_value.assign(rhs, builder);
500                builder.push_traced_op(DslIr::AssertNeE(lhs, rhs_value));
501            }
502            (lhs, rhs) => {
503                let lhs_value = Self::uninit(builder);
504                lhs_value.assign(lhs, builder);
505                let rhs_value = Self::uninit(builder);
506                rhs_value.assign(rhs, builder);
507                builder.push_traced_op(DslIr::AssertNeE(lhs_value, rhs_value));
508            }
509        }
510    }
511}
512
513impl<C: Config> MemVariable<C> for Ext<C::F, C::EF> {
514    fn size_of() -> usize {
515        1
516    }
517
518    fn load(&self, ptr: Ptr<C::N>, index: MemIndex<C::N>, builder: &mut Builder<C>) {
519        builder.push_op(DslIr::LoadE(*self, ptr, index));
520    }
521
522    fn store(&self, ptr: Ptr<<C as Config>::N>, index: MemIndex<C::N>, builder: &mut Builder<C>) {
523        builder.push_op(DslIr::StoreE(*self, ptr, index));
524    }
525}
526
527impl<C: Config> FromConstant<C> for Var<C::N> {
528    type Constant = C::N;
529
530    fn constant(value: Self::Constant, builder: &mut Builder<C>) -> Self {
531        builder.eval(value)
532    }
533}
534
535impl<C: Config> FromConstant<C> for Felt<C::F> {
536    type Constant = C::F;
537
538    fn constant(value: Self::Constant, builder: &mut Builder<C>) -> Self {
539        builder.eval(value)
540    }
541}
542
543impl<C: Config> FromConstant<C> for Ext<C::F, C::EF> {
544    type Constant = C::EF;
545
546    fn constant(value: Self::Constant, builder: &mut Builder<C>) -> Self {
547        builder.eval(value.cons())
548    }
549}
sp1_recursion_compiler/ir/types.rs

sp1_recursion_compiler/ir/
types.rs
