sp1_core_machine/control_flow/jump/
mod.rs

1mod air;
2mod columns;
3mod trace;
4
5pub use columns::*;
6use p3_air::BaseAir;
7
8#[derive(Default)]
9pub struct JumpChip;
10
11impl<F> BaseAir<F> for JumpChip {
12    fn width(&self) -> usize {
13        NUM_JUMP_COLS
14    }
15}
16
17#[cfg(test)]
18mod tests {
19    use std::borrow::BorrowMut;
20
21    use p3_baby_bear::BabyBear;
22    use p3_field::AbstractField;
23    use p3_matrix::dense::RowMajorMatrix;
24    use sp1_core_executor::{ExecutionRecord, Instruction, Opcode, Program};
25    use sp1_stark::{
26        air::MachineAir, baby_bear_poseidon2::BabyBearPoseidon2, chip_name, CpuProver,
27        MachineProver, Val,
28    };
29
30    use crate::{
31        control_flow::{JumpChip, JumpColumns},
32        io::SP1Stdin,
33        riscv::RiscvAir,
34        utils::run_malicious_test,
35    };
36
37    #[test]
38    fn test_malicious_jumps() {
39        let mut jump_instructions = [
40            vec![Instruction::new(Opcode::JAL, 29, 8, 0, true, true)],
41            vec![
42                Instruction::new(Opcode::ADD, 28, 0, 8, false, true),
43                Instruction::new(Opcode::JALR, 29, 28, 0, false, true),
44            ],
45        ];
46
47        for instructions in jump_instructions.iter_mut() {
48            instructions.extend(vec![
49                Instruction::new(Opcode::ADD, 30, 0, 5, false, true),
50                Instruction::new(Opcode::ADD, 28, 0, 5, false, true),
51                Instruction::new(Opcode::ADD, 28, 0, 5, false, true),
52            ]);
53            let program = Program::new(instructions.to_vec(), 0, 0);
54            let stdin = SP1Stdin::new();
55
56            type P = CpuProver<BabyBearPoseidon2, RiscvAir<BabyBear>>;
57
58            let malicious_trace_pv_generator =
59                |prover: &P,
60                 record: &mut ExecutionRecord|
61                 -> Vec<(String, RowMajorMatrix<Val<BabyBearPoseidon2>>)> {
62                    let mut traces = prover.generate_traces(record);
63                    let jump_chip_name = chip_name!(JumpChip, BabyBear);
64                    for (chip_name, trace) in traces.iter_mut() {
65                        if *chip_name == jump_chip_name {
66                            let first_row = trace.row_mut(0);
67                            let first_row: &mut JumpColumns<BabyBear> = first_row.borrow_mut();
68                            first_row.next_pc = 4.into();
69                        }
70                    }
71
72                    traces
73                };
74
75            let result =
76                run_malicious_test::<P>(program, stdin, Box::new(malicious_trace_pv_generator));
77            assert!(result.is_err() && result.unwrap_err().is_local_cumulative_sum_failing());
78        }
79    }
80
81    #[test]
82    fn test_malicious_multiple_opcode_flags() {
83        let instructions = vec![
84            Instruction::new(Opcode::JAL, 29, 12, 0, true, true),
85            Instruction::new(Opcode::ADD, 30, 0, 5, false, true),
86            Instruction::new(Opcode::ADD, 28, 0, 5, false, true),
87            Instruction::new(Opcode::ADD, 28, 0, 5, false, true),
88        ];
89        let program = Program::new(instructions, 0, 0);
90        let stdin = SP1Stdin::new();
91
92        type P = CpuProver<BabyBearPoseidon2, RiscvAir<BabyBear>>;
93
94        let malicious_trace_pv_generator =
95            |prover: &P,
96             record: &mut ExecutionRecord|
97             -> Vec<(String, RowMajorMatrix<Val<BabyBearPoseidon2>>)> {
98                // Modify the branch chip to have a row that has multiple opcode flags set.
99                let mut traces = prover.generate_traces(record);
100                let jump_chip_name = chip_name!(JumpChip, BabyBear);
101                for (chip_name, trace) in traces.iter_mut() {
102                    if *chip_name == jump_chip_name {
103                        let first_row = trace.row_mut(0);
104                        let first_row: &mut JumpColumns<BabyBear> = first_row.borrow_mut();
105                        assert!(first_row.is_jal == BabyBear::one());
106                        first_row.is_jalr = BabyBear::one();
107                    }
108                }
109                traces
110            };
111
112        let result =
113            run_malicious_test::<P>(program, stdin, Box::new(malicious_trace_pv_generator));
114        let jump_chip_name = chip_name!(JumpChip, BabyBear);
115        assert!(result.is_err() && result.unwrap_err().is_constraints_failing(&jump_chip_name));
116    }
117}