sp1_core_machine/control_flow/jump/
air.rs

1use std::borrow::Borrow;
2
3use p3_air::{Air, AirBuilder};
4use p3_field::AbstractField;
5use p3_matrix::Matrix;
6use sp1_core_executor::{Opcode, DEFAULT_PC_INC, UNUSED_PC};
7use sp1_stark::air::{BaseAirBuilder, SP1AirBuilder};
8
9use crate::operations::BabyBearWordRangeChecker;
10
11use super::{JumpChip, JumpColumns};
12
13impl<AB> Air<AB> for JumpChip
14where
15    AB: SP1AirBuilder,
16    AB::Var: Sized,
17{
18    #[inline(never)]
19    fn eval(&self, builder: &mut AB) {
20        let main = builder.main();
21        let local = main.row_slice(0);
22        let local: &JumpColumns<AB::Var> = (*local).borrow();
23
24        // SAFETY: All selectors `is_jal`, `is_jalr` are checked to be boolean.
25        // Each "real" row has exactly one selector turned on, as `is_real = is_jal + is_jalr` is
26        // boolean. Therefore, the `opcode` matches the corresponding opcode.
27        builder.assert_bool(local.is_jal);
28        builder.assert_bool(local.is_jalr);
29        let is_real = local.is_jal + local.is_jalr;
30        builder.assert_bool(is_real.clone());
31
32        let opcode = local.is_jal * Opcode::JAL.as_field::<AB::F>() +
33            local.is_jalr * Opcode::JALR.as_field::<AB::F>();
34
35        // SAFETY: This checks the following.
36        // - `num_extra_cycles = 0`
37        // - `op_a_immutable = 0`
38        // - `is_memory = 0`
39        // - `is_syscall = 0`
40        // - `is_halt = 0`
41        // `next_pc` and `op_a_value` still has to be constrained, and this is done below.
42        builder.receive_instruction(
43            AB::Expr::zero(),
44            AB::Expr::zero(),
45            local.pc.reduce::<AB>(),
46            local.next_pc.reduce::<AB>(),
47            AB::Expr::zero(),
48            opcode,
49            local.op_a_value,
50            local.op_b_value,
51            local.op_c_value,
52            local.op_a_0,
53            AB::Expr::zero(),
54            AB::Expr::zero(),
55            AB::Expr::zero(),
56            AB::Expr::zero(),
57            is_real.clone(),
58        );
59
60        // Verify that the local.pc + 4 is saved in op_a for both jump instructions.
61        // When op_a is set to register X0, the RISC-V spec states that the jump instruction will
62        // not have a return destination address (it is effectively a GOTO command).  In this case,
63        // we shouldn't verify the return address.
64        builder.when(is_real.clone()).when_not(local.op_a_0).assert_eq(
65            local.op_a_value.reduce::<AB>(),
66            local.pc.reduce::<AB>() + AB::F::from_canonical_u32(DEFAULT_PC_INC),
67        );
68
69        // Range check op_a, pc, and next_pc.
70        // SAFETY: `is_real` is already checked to be boolean.
71        // `op_a_value` is checked to be a valid word, as it matches the one in the CpuChip.
72        // In the CpuChip's `eval_registers`, it's checked that this is a valid word.
73        // Combined with the `op_a_value = pc + 4` check above when `op_a_0 = 0`, this fully
74        // constrains `op_a_value`.
75        BabyBearWordRangeChecker::<AB::F>::range_check(
76            builder,
77            local.op_a_value,
78            local.op_a_range_checker,
79            is_real.clone(),
80        );
81        // SAFETY: `is_real` is already checked to be boolean.
82        // `local.pc`, `local.next_pc` are checked to a valid word when relevant.
83        // This is due to the ADD ALU table checking all inputs and outputs are valid words.
84        // This is done when the `AddOperation` is invoked in the ADD ALU table.
85        BabyBearWordRangeChecker::<AB::F>::range_check(
86            builder,
87            local.pc,
88            local.pc_range_checker,
89            is_real.clone(),
90        );
91        BabyBearWordRangeChecker::<AB::F>::range_check(
92            builder,
93            local.next_pc,
94            local.next_pc_range_checker,
95            is_real,
96        );
97
98        // We now constrain `next_pc`.
99
100        // Verify that the new pc is calculated correctly for JAL instructions.
101        // SAFETY: `is_jal` is boolean, and zero for padding rows.
102        builder.send_instruction(
103            AB::Expr::zero(),
104            AB::Expr::zero(),
105            AB::Expr::from_canonical_u32(UNUSED_PC),
106            AB::Expr::from_canonical_u32(UNUSED_PC + DEFAULT_PC_INC),
107            AB::Expr::zero(),
108            AB::Expr::from_canonical_u32(Opcode::ADD as u32),
109            local.next_pc,
110            local.pc,
111            local.op_b_value,
112            AB::Expr::zero(),
113            AB::Expr::zero(),
114            AB::Expr::zero(),
115            AB::Expr::zero(),
116            AB::Expr::zero(),
117            local.is_jal,
118        );
119
120        // Verify that the new pc is calculated correctly for JALR instructions.
121        // SAFETY: `is_jalr` is boolean, and zero for padding rows.
122        builder.send_instruction(
123            AB::Expr::zero(),
124            AB::Expr::zero(),
125            AB::Expr::from_canonical_u32(UNUSED_PC),
126            AB::Expr::from_canonical_u32(UNUSED_PC + DEFAULT_PC_INC),
127            AB::Expr::zero(),
128            AB::Expr::from_canonical_u32(Opcode::ADD as u32),
129            local.next_pc,
130            local.op_b_value,
131            local.op_c_value,
132            AB::Expr::zero(),
133            AB::Expr::zero(),
134            AB::Expr::zero(),
135            AB::Expr::zero(),
136            AB::Expr::zero(),
137            local.is_jalr,
138        );
139    }
140}
sp1_core_machine/control_flow/jump/air.rs

sp1_core_machine/control_flow/jump/
air.rs
