o1vm/interpreters/riscv32im/
constraints.rs

1use super::{
2    column::{Column, E},
3    interpreter::{Instruction, InterpreterEnv},
4    INSTRUCTION_SET_SIZE,
5};
6use crate::{
7    interpreters::riscv32im::{constraints::ConstantTerm::Literal, SCRATCH_SIZE},
8    lookups::Lookup,
9};
10use ark_ff::{Field, One};
11use kimchi::circuits::{
12    expr::{ConstantTerm, Expr, ExprInner, Operations, Variable},
13    gate::CurrOrNext,
14};
15
16pub struct Env<F: Field> {
17    pub scratch_state_idx: usize,
18    pub scratch_state_idx_inverse: usize,
19    pub lookups: Vec<Lookup<E<F>>>,
20    pub constraints: Vec<E<F>>,
21    pub selector: Option<E<F>>,
22}
23
24impl<Fp: Field> Default for Env<Fp> {
25    fn default() -> Self {
26        Self {
27            scratch_state_idx: 0,
28            scratch_state_idx_inverse: 0,
29            constraints: Vec::new(),
30            lookups: Vec::new(),
31            selector: None,
32        }
33    }
34}
35
36impl<Fp: Field> InterpreterEnv for Env<Fp> {
37    /// In the concrete implementation for the constraints, the interpreter will
38    /// work over columns. The position in this case can be seen as a new
39    /// variable/input of our circuit.
40    type Position = Column;
41
42    // Use one of the available columns. It won't create a new column every time
43    // this function is called. The number of columns is defined upfront by
44    // crate::mips::witness::SCRATCH_SIZE.
45    fn alloc_scratch(&mut self) -> Self::Position {
46        // All columns are implemented using a simple index, and a name is given
47        // to the index. See crate::SCRATCH_SIZE for the maximum number of
48        // columns the circuit can use.
49        let scratch_idx = self.scratch_state_idx;
50        self.scratch_state_idx += 1;
51        Column::ScratchState(scratch_idx)
52    }
53
54    fn alloc_scratch_inverse(&mut self) -> Self::Position {
55        let scratch_idx = self.scratch_state_idx_inverse;
56        self.scratch_state_idx_inverse += 1;
57        Column::ScratchStateInverse(scratch_idx)
58    }
59
60    type Variable = E<Fp>;
61
62    fn variable(&self, column: Self::Position) -> Self::Variable {
63        Expr::Atom(ExprInner::Cell(Variable {
64            col: column,
65            row: CurrOrNext::Curr,
66        }))
67    }
68
69    fn activate_selector(&mut self, selector: Instruction) {
70        // Sanity check: we only want to activate once per instruction
71        assert!(self.selector.is_none(), "A selector has been already activated. You might need to reset the environment if you want to start a new instruction.");
72        let n = usize::from(selector) - SCRATCH_SIZE - 1;
73        self.selector = Some(self.variable(Column::Selector(n)))
74    }
75
76    fn add_constraint(&mut self, assert_equals_zero: Self::Variable) {
77        self.constraints.push(assert_equals_zero)
78    }
79
80    fn check_is_zero(_assert_equals_zero: &Self::Variable) {
81        // No-op, witness only
82    }
83
84    fn check_equal(_x: &Self::Variable, _y: &Self::Variable) {
85        // No-op, witness only
86    }
87
88    fn assert_boolean(&mut self, x: &Self::Variable) {
89        self.add_constraint(x.clone() * x.clone() - x.clone());
90    }
91
92    fn add_lookup(&mut self, lookup: Lookup<Self::Variable>) {
93        self.lookups.push(lookup);
94    }
95
96    fn instruction_counter(&self) -> Self::Variable {
97        self.variable(Column::InstructionCounter)
98    }
99
100    fn increase_instruction_counter(&mut self) {
101        // No-op, witness only
102    }
103
104    unsafe fn fetch_register(
105        &mut self,
106        _idx: &Self::Variable,
107        output: Self::Position,
108    ) -> Self::Variable {
109        self.variable(output)
110    }
111
112    unsafe fn push_register_if(
113        &mut self,
114        _idx: &Self::Variable,
115        _value: Self::Variable,
116        _if_is_true: &Self::Variable,
117    ) {
118        // No-op, witness only
119    }
120
121    unsafe fn fetch_register_access(
122        &mut self,
123        _idx: &Self::Variable,
124        output: Self::Position,
125    ) -> Self::Variable {
126        self.variable(output)
127    }
128
129    unsafe fn push_register_access_if(
130        &mut self,
131        _idx: &Self::Variable,
132        _value: Self::Variable,
133        _if_is_true: &Self::Variable,
134    ) {
135        // No-op, witness only
136    }
137
138    unsafe fn fetch_memory(
139        &mut self,
140        _addr: &Self::Variable,
141        output: Self::Position,
142    ) -> Self::Variable {
143        self.variable(output)
144    }
145
146    unsafe fn push_memory(&mut self, _addr: &Self::Variable, _value: Self::Variable) {
147        // No-op, witness only
148    }
149
150    unsafe fn fetch_memory_access(
151        &mut self,
152        _addr: &Self::Variable,
153        output: Self::Position,
154    ) -> Self::Variable {
155        self.variable(output)
156    }
157
158    unsafe fn push_memory_access(&mut self, _addr: &Self::Variable, _value: Self::Variable) {
159        // No-op, witness only
160    }
161
162    fn constant(x: u32) -> Self::Variable {
163        Self::Variable::constant(Operations::from(Literal(Fp::from(x))))
164    }
165
166    unsafe fn bitmask(
167        &mut self,
168        _x: &Self::Variable,
169        _highest_bit: u32,
170        _lowest_bit: u32,
171        position: Self::Position,
172    ) -> Self::Variable {
173        self.variable(position)
174    }
175
176    unsafe fn shift_left(
177        &mut self,
178        _x: &Self::Variable,
179        _by: &Self::Variable,
180        position: Self::Position,
181    ) -> Self::Variable {
182        self.variable(position)
183    }
184
185    unsafe fn shift_right(
186        &mut self,
187        _x: &Self::Variable,
188        _by: &Self::Variable,
189        position: Self::Position,
190    ) -> Self::Variable {
191        self.variable(position)
192    }
193
194    unsafe fn shift_right_arithmetic(
195        &mut self,
196        _x: &Self::Variable,
197        _by: &Self::Variable,
198        position: Self::Position,
199    ) -> Self::Variable {
200        self.variable(position)
201    }
202
203    unsafe fn test_zero(
204        &mut self,
205        _x: &Self::Variable,
206        position: Self::Position,
207    ) -> Self::Variable {
208        self.variable(position)
209    }
210
211    fn is_zero(&mut self, x: &Self::Variable) -> Self::Variable {
212        let res = {
213            let pos = self.alloc_scratch();
214            unsafe { self.test_zero(x, pos) }
215        };
216        let x_inv_or_zero = {
217            let pos = self.alloc_scratch_inverse();
218            self.variable(pos)
219        };
220        // If x = 0, then res = 1 and x_inv_or_zero = 0
221        // If x <> 0, then res = 0 and x_inv_or_zero = x^(-1)
222        self.add_constraint(x.clone() * x_inv_or_zero.clone() + res.clone() - Self::constant(1));
223        self.add_constraint(x.clone() * res.clone());
224        res
225    }
226
227    fn equal(&mut self, x: &Self::Variable, y: &Self::Variable) -> Self::Variable {
228        self.is_zero(&(x.clone() - y.clone()))
229    }
230
231    unsafe fn test_less_than(
232        &mut self,
233        _x: &Self::Variable,
234        _y: &Self::Variable,
235        position: Self::Position,
236    ) -> Self::Variable {
237        self.variable(position)
238    }
239
240    unsafe fn test_less_than_signed(
241        &mut self,
242        _x: &Self::Variable,
243        _y: &Self::Variable,
244        position: Self::Position,
245    ) -> Self::Variable {
246        self.variable(position)
247    }
248
249    unsafe fn and_witness(
250        &mut self,
251        _x: &Self::Variable,
252        _y: &Self::Variable,
253        position: Self::Position,
254    ) -> Self::Variable {
255        self.variable(position)
256    }
257
258    unsafe fn nor_witness(
259        &mut self,
260        _x: &Self::Variable,
261        _y: &Self::Variable,
262        position: Self::Position,
263    ) -> Self::Variable {
264        self.variable(position)
265    }
266
267    unsafe fn or_witness(
268        &mut self,
269        _x: &Self::Variable,
270        _y: &Self::Variable,
271        position: Self::Position,
272    ) -> Self::Variable {
273        self.variable(position)
274    }
275
276    unsafe fn xor_witness(
277        &mut self,
278        _x: &Self::Variable,
279        _y: &Self::Variable,
280        position: Self::Position,
281    ) -> Self::Variable {
282        self.variable(position)
283    }
284
285    unsafe fn add_witness(
286        &mut self,
287        _y: &Self::Variable,
288        _x: &Self::Variable,
289        out_position: Self::Position,
290        overflow_position: Self::Position,
291    ) -> (Self::Variable, Self::Variable) {
292        (
293            self.variable(out_position),
294            self.variable(overflow_position),
295        )
296    }
297
298    unsafe fn sub_witness(
299        &mut self,
300        _y: &Self::Variable,
301        _x: &Self::Variable,
302        out_position: Self::Position,
303        underflow_position: Self::Position,
304    ) -> (Self::Variable, Self::Variable) {
305        (
306            self.variable(out_position),
307            self.variable(underflow_position),
308        )
309    }
310
311    unsafe fn mul_signed_witness(
312        &mut self,
313        _x: &Self::Variable,
314        _y: &Self::Variable,
315        position: Self::Position,
316    ) -> Self::Variable {
317        self.variable(position)
318    }
319
320    unsafe fn mul_hi_signed(
321        &mut self,
322        _x: &Self::Variable,
323        _y: &Self::Variable,
324        position: Self::Position,
325    ) -> Self::Variable {
326        self.variable(position)
327    }
328
329    unsafe fn mul_lo_signed(
330        &mut self,
331        _x: &Self::Variable,
332        _y: &Self::Variable,
333        position: Self::Position,
334    ) -> Self::Variable {
335        self.variable(position)
336    }
337
338    unsafe fn mul_hi(
339        &mut self,
340        _x: &Self::Variable,
341        _y: &Self::Variable,
342        position: Self::Position,
343    ) -> Self::Variable {
344        self.variable(position)
345    }
346
347    unsafe fn mul_hi_signed_unsigned(
348        &mut self,
349        _x: &Self::Variable,
350        _y: &Self::Variable,
351        position: Self::Position,
352    ) -> Self::Variable {
353        self.variable(position)
354    }
355
356    unsafe fn div_signed(
357        &mut self,
358        _x: &Self::Variable,
359        _y: &Self::Variable,
360        position: Self::Position,
361    ) -> Self::Variable {
362        self.variable(position)
363    }
364
365    unsafe fn mod_signed(
366        &mut self,
367        _x: &Self::Variable,
368        _y: &Self::Variable,
369        position: Self::Position,
370    ) -> Self::Variable {
371        self.variable(position)
372    }
373
374    unsafe fn div(
375        &mut self,
376        _x: &Self::Variable,
377        _y: &Self::Variable,
378        position: Self::Position,
379    ) -> Self::Variable {
380        self.variable(position)
381    }
382
383    unsafe fn mod_unsigned(
384        &mut self,
385        _x: &Self::Variable,
386        _y: &Self::Variable,
387        position: Self::Position,
388    ) -> Self::Variable {
389        self.variable(position)
390    }
391
392    unsafe fn mul_lo(
393        &mut self,
394        _x: &Self::Variable,
395        _y: &Self::Variable,
396        position: Self::Position,
397    ) -> Self::Variable {
398        self.variable(position)
399    }
400
401    unsafe fn count_leading_zeros(
402        &mut self,
403        _x: &Self::Variable,
404        position: Self::Position,
405    ) -> Self::Variable {
406        self.variable(position)
407    }
408
409    unsafe fn count_leading_ones(
410        &mut self,
411        _x: &Self::Variable,
412        position: Self::Position,
413    ) -> Self::Variable {
414        self.variable(position)
415    }
416
417    fn copy(&mut self, x: &Self::Variable, position: Self::Position) -> Self::Variable {
418        let res = self.variable(position);
419        self.constraints.push(x.clone() - res.clone());
420        res
421    }
422
423    fn set_halted(&mut self, _flag: Self::Variable) {
424        // TODO
425    }
426
427    fn report_exit(&mut self, _exit_code: &Self::Variable) {}
428
429    fn reset(&mut self) {
430        self.scratch_state_idx = 0;
431        self.constraints.clear();
432        self.lookups.clear();
433        self.selector = None;
434    }
435}
436
437impl<Fp: Field> Env<Fp> {
438    /// Return the constraints for the selector.
439    /// Each selector must be a boolean.
440    pub fn get_selector_constraints(&self) -> Vec<E<Fp>> {
441        let one = <Self as InterpreterEnv>::Variable::one();
442        let mut enforce_bool: Vec<E<Fp>> = (0..INSTRUCTION_SET_SIZE)
443            .map(|i| {
444                let var = self.variable(Column::Selector(i));
445                (var.clone() - one.clone()) * var.clone()
446            })
447            .collect();
448        let enforce_one_activation = (0..INSTRUCTION_SET_SIZE).fold(E::<Fp>::one(), |res, i| {
449            let var = self.variable(Column::Selector(i));
450            res - var.clone()
451        });
452
453        enforce_bool.push(enforce_one_activation);
454        enforce_bool
455    }
456
457    pub fn get_selector(&self) -> E<Fp> {
458        self.selector
459            .clone()
460            .unwrap_or_else(|| panic!("Selector is not set"))
461    }
462
463    /// Return the constraints for the current instruction, without the selector
464    pub fn get_constraints(&self) -> Vec<E<Fp>> {
465        self.constraints.clone()
466    }
467
468    pub fn get_lookups(&self) -> Vec<Lookup<E<Fp>>> {
469        self.lookups.clone()
470    }
471}