1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

use std::borrow::BorrowMut;

use p3_field::{extension::BinomiallyExtendable, PrimeField32};
use p3_matrix::{dense::RowMajorMatrix, Matrix};
use sp1_core::{
    air::{BinomialExtension, MachineAir},
    utils::pad_rows_fixed,
};
use tracing::instrument;

use crate::{
    air::BinomialExtensionUtils,
    memory::MemoryCols,
    runtime::{
        get_heap_size_range_check_events, instruction_is_heap_expand, ExecutionRecord, Opcode,
        RecursionProgram, D,
    },
};

use super::{CpuChip, CpuCols, CPU_COL_MAP, NUM_CPU_COLS};

impl<F: PrimeField32 + BinomiallyExtendable<D>, const L: usize> MachineAir<F> for CpuChip<F, L> {
    type Record = ExecutionRecord<F>;
    type Program = RecursionProgram<F>;

    fn name(&self) -> String {
        "CPU".to_string()
    }

    fn generate_dependencies(&self, _: &Self::Record, _: &mut Self::Record) {
        // There are no dependencies, since we do it all in the runtime. This is just a placeholder.
    }

    #[instrument(name = "generate cpu trace", level = "debug", skip_all, fields(rows = input.cpu_events.len()))]
    fn generate_trace(
        &self,
        input: &ExecutionRecord<F>,
        _: &mut ExecutionRecord<F>,
    ) -> RowMajorMatrix<F> {
        let mut rows = input
            .cpu_events
            .iter()
            .map(|event| {
                let mut row = [F::zero(); NUM_CPU_COLS];
                let cols: &mut CpuCols<F> = row.as_mut_slice().borrow_mut();

                cols.clk = event.clk;
                cols.pc = event.pc;
                cols.fp = event.fp;

                // Populate the instruction related columns.
                cols.selectors.populate(&event.instruction);
                cols.instruction.populate(&event.instruction);

                // Populate the register columns.
                if let Some(record) = &event.a_record {
                    cols.a.populate(record);
                }
                if let Some(record) = &event.b_record {
                    cols.b.populate(record);
                } else {
                    *cols.b.value_mut() = event.instruction.op_b;
                }
                if let Some(record) = &event.c_record {
                    cols.c.populate(record);
                } else {
                    *cols.c.value_mut() = event.instruction.op_c;
                }
                if let Some(record) = &event.memory_record {
                    let memory_cols = cols.opcode_specific.memory_mut();
                    memory_cols.memory.populate(record);
                    memory_cols.memory_addr = record.addr;
                }

                // Populate the heap columns.
                if instruction_is_heap_expand(&event.instruction) {
                    let (u16_range_check, u12_range_check) =
                        get_heap_size_range_check_events(cols.a.value()[0]);

                    let heap_cols = cols.opcode_specific.heap_expand_mut();
                    heap_cols.diff_16bit_limb = F::from_canonical_u16(u16_range_check.val);
                    heap_cols.diff_12bit_limb = F::from_canonical_u16(u12_range_check.val);
                }

                // Populate the branch columns.
                if matches!(
                    event.instruction.opcode,
                    Opcode::BEQ | Opcode::BNE | Opcode::BNEINC
                ) {
                    let branch_cols = cols.opcode_specific.branch_mut();
                    let a_ext: BinomialExtension<F> =
                        BinomialExtensionUtils::from_block(*cols.a.value());
                    let b_ext: BinomialExtension<F> =
                        BinomialExtensionUtils::from_block(*cols.b.value());

                    let (comparison_diff, do_branch) = match event.instruction.opcode {
                        Opcode::BEQ => (a_ext - b_ext, a_ext == b_ext),
                        Opcode::BNE | Opcode::BNEINC => (a_ext - b_ext, a_ext != b_ext),
                        _ => unreachable!(),
                    };

                    branch_cols
                        .comparison_diff
                        .populate((comparison_diff).as_block());
                    branch_cols.comparison_diff_val = comparison_diff;
                    branch_cols.do_branch = F::from_bool(do_branch);
                    branch_cols.next_pc = if do_branch {
                        event.pc + event.instruction.op_c[0]
                    } else {
                        event.pc + F::one()
                    };
                }

                // Populate the public values columns.
                if event.instruction.opcode == Opcode::Commit {
                    let public_values_cols = cols.opcode_specific.public_values_mut();
                    let idx = cols.b.prev_value()[0].as_canonical_u32() as usize;
                    public_values_cols.idx_bitmap[idx] = F::one();
                }

                cols.is_real = F::one();
                row
            })
            .collect::<Vec<_>>();

        pad_rows_fixed(
            &mut rows,
            || {
                let mut row = [F::zero(); NUM_CPU_COLS];
                let cols: &mut CpuCols<F> = row.as_mut_slice().borrow_mut();
                cols.selectors.is_noop = F::one();
                cols.instruction.imm_b = F::one();
                cols.instruction.imm_c = F::one();
                row
            },
            self.fixed_log2_rows,
        );

        let mut trace =
            RowMajorMatrix::new(rows.into_iter().flatten().collect::<Vec<_>>(), NUM_CPU_COLS);

        for i in input.cpu_events.len()..trace.height() {
            trace.values[i * NUM_CPU_COLS + CPU_COL_MAP.clk] =
                F::from_canonical_u32(4) * F::from_canonical_usize(i);
            trace.values[i * NUM_CPU_COLS + CPU_COL_MAP.instruction.imm_b] =
                F::from_canonical_u32(1);
            trace.values[i * NUM_CPU_COLS + CPU_COL_MAP.instruction.imm_c] =
                F::from_canonical_u32(1);
        }
        trace
    }

    fn included(&self, _: &Self::Record) -> bool {
        true
    }
}