openvm_sha2_circuit/sha2_chips/main_chip/
air.rs

1use std::marker::PhantomData;
2
3use itertools::izip;
4use ndarray::s;
5use openvm_circuit::{
6    arch::ExecutionBridge,
7    system::{
8        memory::{offline_checker::MemoryBridge, MemoryAddress},
9        SystemPort,
10    },
11};
12use openvm_circuit_primitives::{
13    bitwise_op_lookup::BitwiseOperationLookupBus, utils::compose, ColumnsAir,
14};
15use openvm_instructions::riscv::{
16    RV32_CELL_BITS, RV32_MEMORY_AS, RV32_REGISTER_AS, RV32_REGISTER_NUM_LIMBS,
17};
18use openvm_sha2_air::Sha2BlockHasherSubairConfig;
19use openvm_stark_backend::{
20    interaction::{BusIndex, InteractionBuilder, PermutationCheckBus},
21    p3_air::{Air, AirBuilder, BaseAir},
22    p3_field::PrimeCharacteristicRing,
23    p3_matrix::Matrix,
24    BaseAirWithPublicValues, PartitionedBaseAir,
25};
26
27use super::config::Sha2MainChipConfig;
28use crate::{MessageType, Sha2ColsRef, SHA2_READ_SIZE, SHA2_WRITE_SIZE};
29
30#[derive(Clone, Debug)]
31pub struct Sha2MainAir<C: Sha2MainChipConfig> {
32    pub execution_bridge: ExecutionBridge,
33    pub memory_bridge: MemoryBridge,
34    pub bitwise_lookup_bus: BitwiseOperationLookupBus,
35    pub sha2_bus: PermutationCheckBus,
36    /// Maximum number of bits allowed for an address pointer
37    /// Must be at least 24
38    pub ptr_max_bits: usize,
39    pub offset: usize,
40    _phantom: PhantomData<C>,
41}
42
43// No columns provided: width is the config-dependent `C::MAIN_CHIP_WIDTH` and rows are accessed
44// via `Sha2ColsRef` (a slice-borrowing ref struct, no static `Cols`).
45impl<C: Sha2MainChipConfig> ColumnsAir for Sha2MainAir<C> {}
46
47impl<C: Sha2MainChipConfig> Sha2MainAir<C> {
48    pub fn new(
49        SystemPort {
50            execution_bus,
51            program_bus,
52            memory_bridge,
53        }: SystemPort,
54        bitwise_lookup_bus: BitwiseOperationLookupBus,
55        ptr_max_bits: usize,
56        self_bus_idx: BusIndex,
57        offset: usize,
58    ) -> Self {
59        Self {
60            execution_bridge: ExecutionBridge::new(execution_bus, program_bus),
61            memory_bridge,
62            bitwise_lookup_bus,
63            sha2_bus: PermutationCheckBus::new(self_bus_idx),
64            ptr_max_bits,
65            offset,
66            _phantom: PhantomData,
67        }
68    }
69}
70
71impl<F, C: Sha2MainChipConfig> BaseAirWithPublicValues<F> for Sha2MainAir<C> {}
72impl<F, C: Sha2MainChipConfig> PartitionedBaseAir<F> for Sha2MainAir<C> {}
73impl<F, C: Sha2MainChipConfig> BaseAir<F> for Sha2MainAir<C> {
74    fn width(&self) -> usize {
75        C::MAIN_CHIP_WIDTH
76    }
77}
78
79impl<AB: InteractionBuilder, C: Sha2MainChipConfig + Sha2BlockHasherSubairConfig> Air<AB>
80    for Sha2MainAir<C>
81{
82    fn eval(&self, builder: &mut AB) {
83        let main = builder.main();
84        let (local, next) = (main.row_slice(0).unwrap(), main.row_slice(1).unwrap());
85
86        let local: Sha2ColsRef<AB::Var> = Sha2ColsRef::from::<C>(&local[..C::MAIN_CHIP_WIDTH]);
87        let next: Sha2ColsRef<AB::Var> = Sha2ColsRef::from::<C>(&next[..C::MAIN_CHIP_WIDTH]);
88
89        let mut timestamp_delta = 0;
90        let mut timestamp_pp = || {
91            timestamp_delta += 1;
92            local.instruction.from_state.timestamp + AB::F::from_usize(timestamp_delta - 1)
93        };
94
95        self.eval_block(builder, &local, &next);
96        self.eval_instruction(builder, &local, &mut timestamp_pp);
97        self.eval_reads(builder, &local, &mut timestamp_pp);
98        self.eval_writes(builder, &local, &mut timestamp_pp);
99    }
100}
101
102impl<C: Sha2MainChipConfig + Sha2BlockHasherSubairConfig> Sha2MainAir<C> {
103    pub fn eval_block<AB: InteractionBuilder>(
104        &self,
105        builder: &mut AB,
106        local: &Sha2ColsRef<AB::Var>,
107        next: &Sha2ColsRef<AB::Var>,
108    ) {
109        builder.assert_bool(*local.instruction.is_enabled);
110        builder
111            .when_transition()
112            .when_ne(*local.instruction.is_enabled, AB::Expr::ONE)
113            .assert_zero(*next.instruction.is_enabled);
114
115        builder
116            .when_first_row()
117            .when(*local.instruction.is_enabled)
118            .assert_zero(*local.block.request_id);
119
120        builder
121            .when_transition()
122            .when(*next.instruction.is_enabled)
123            .assert_eq(
124                *next.block.request_id,
125                *local.block.request_id + AB::Expr::ONE,
126            );
127
128        let prev_state_as_u16s: Vec<AB::Expr> = local
129            .block
130            .prev_state
131            .exact_chunks(C::WORD_U8S)
132            .into_iter()
133            .flat_map(|word| {
134                word.as_slice()
135                    .unwrap()
136                    .chunks_exact(2)
137                    .map(|x| x[1] * AB::F::from_u64(1 << 8) + x[0])
138                    .collect::<Vec<_>>()
139            })
140            .collect();
141
142        // Send (STATE, request_id, prev_state_as_u16s, new_state) to the sha2 bus
143        self.sha2_bus.send(
144            builder,
145            [
146                AB::Expr::from_u8(MessageType::State as u8),
147                (*local.block.request_id).into(),
148            ]
149            .into_iter()
150            .chain(prev_state_as_u16s)
151            .chain(local.block.new_state.into_iter().copied().map(|x| x.into())),
152            *local.instruction.is_enabled,
153        );
154
155        // Send (MESSAGE_1, request_id, first_half_of_message) to the sha2 bus
156        self.sha2_bus.send(
157            builder,
158            [
159                AB::Expr::from_u8(MessageType::Message1 as u8),
160                (*local.block.request_id).into(),
161            ]
162            .into_iter()
163            .chain(
164                local
165                    .block
166                    .message_bytes
167                    .iter()
168                    .take(C::BLOCK_BYTES / 2)
169                    .map(|x| (*x).into()),
170            ),
171            *local.instruction.is_enabled,
172        );
173
174        // Send (MESSAGE_2, request_id, second_half_of_message) to the sha2 bus
175        self.sha2_bus.send(
176            builder,
177            [
178                AB::Expr::from_u8(MessageType::Message2 as u8),
179                (*local.block.request_id).into(),
180            ]
181            .into_iter()
182            .chain(
183                local
184                    .block
185                    .message_bytes
186                    .iter()
187                    .skip(C::BLOCK_BYTES / 2)
188                    .map(|x| (*x).into()),
189            ),
190            *local.instruction.is_enabled,
191        );
192    }
193
194    pub fn eval_instruction<AB: InteractionBuilder>(
195        &self,
196        builder: &mut AB,
197        local: &Sha2ColsRef<AB::Var>,
198        timestamp_pp: &mut impl FnMut() -> AB::Expr,
199    ) {
200        for (&ptr, val, aux) in izip!(
201            [
202                local.instruction.dst_reg_ptr,
203                local.instruction.state_reg_ptr,
204                local.instruction.input_reg_ptr
205            ],
206            [
207                local.instruction.dst_ptr_limbs,
208                local.instruction.state_ptr_limbs,
209                local.instruction.input_ptr_limbs
210            ],
211            &local.mem.register_aux,
212        ) {
213            self.memory_bridge
214                .read::<_, _, SHA2_READ_SIZE>(
215                    MemoryAddress::new(AB::Expr::from_u32(RV32_REGISTER_AS), ptr),
216                    val.to_vec().try_into().unwrap_or_else(|_| panic!()), // can't unwrap because AB::Var doesn't impl Debug
217                    timestamp_pp(),
218                    aux,
219                )
220                .eval(builder, *local.instruction.is_enabled);
221        }
222
223        // range check the memory pointers
224        assert!(self.ptr_max_bits >= RV32_CELL_BITS * (RV32_REGISTER_NUM_LIMBS - 1));
225        let shift = AB::Expr::from_usize(
226            1 << (RV32_REGISTER_NUM_LIMBS * RV32_CELL_BITS - self.ptr_max_bits),
227        );
228        let needs_range_check = [
229            local.instruction.dst_ptr_limbs[RV32_REGISTER_NUM_LIMBS - 1],
230            local.instruction.state_ptr_limbs[RV32_REGISTER_NUM_LIMBS - 1],
231            local.instruction.input_ptr_limbs[RV32_REGISTER_NUM_LIMBS - 1],
232            local.instruction.input_ptr_limbs[RV32_REGISTER_NUM_LIMBS - 1],
233        ];
234        for pair in needs_range_check.chunks_exact(2) {
235            self.bitwise_lookup_bus
236                .send_range(pair[0] * shift.clone(), pair[1] * shift.clone())
237                .eval(builder, *local.instruction.is_enabled);
238        }
239
240        self.execution_bridge
241            .execute_and_increment_pc(
242                AB::Expr::from_usize(C::OPCODE as usize + self.offset),
243                [
244                    (*local.instruction.dst_reg_ptr).into(),
245                    (*local.instruction.state_reg_ptr).into(),
246                    (*local.instruction.input_reg_ptr).into(),
247                    AB::Expr::from_u32(RV32_REGISTER_AS),
248                    AB::Expr::from_u32(RV32_MEMORY_AS),
249                ],
250                *local.instruction.from_state,
251                AB::F::from_usize(C::TIMESTAMP_DELTA),
252            )
253            .eval(builder, *local.instruction.is_enabled);
254    }
255
256    pub fn eval_reads<AB: InteractionBuilder>(
257        &self,
258        builder: &mut AB,
259        local: &Sha2ColsRef<AB::Var>,
260        timestamp_pp: &mut impl FnMut() -> AB::Expr,
261    ) {
262        let input_ptr_val = compose(&local.instruction.input_ptr_limbs.to_vec(), RV32_CELL_BITS);
263        for i in 0..C::BLOCK_READS {
264            self.memory_bridge
265                .read::<_, _, SHA2_READ_SIZE>(
266                    MemoryAddress::new(
267                        AB::Expr::from_u32(RV32_MEMORY_AS),
268                        input_ptr_val.clone() + AB::F::from_usize(i * SHA2_READ_SIZE),
269                    ),
270                    local
271                        .block
272                        .message_bytes
273                        .slice(s![i * SHA2_READ_SIZE..(i + 1) * SHA2_READ_SIZE])
274                        .to_vec()
275                        .try_into()
276                        .unwrap_or_else(|_| {
277                            panic!("message bytes is not the correct size");
278                        }),
279                    timestamp_pp(),
280                    &local.mem.input_reads[i],
281                )
282                .eval(builder, *local.instruction.is_enabled);
283        }
284
285        let state_ptr_val = compose(&local.instruction.state_ptr_limbs.to_vec(), RV32_CELL_BITS);
286        for i in 0..C::STATE_READS {
287            self.memory_bridge
288                .read::<_, _, SHA2_READ_SIZE>(
289                    MemoryAddress::new(
290                        AB::Expr::from_u32(RV32_MEMORY_AS),
291                        state_ptr_val.clone() + AB::F::from_usize(i * SHA2_READ_SIZE),
292                    ),
293                    local
294                        .block
295                        .prev_state
296                        .slice(s![i * SHA2_READ_SIZE..(i + 1) * SHA2_READ_SIZE])
297                        .to_vec()
298                        .try_into()
299                        .unwrap_or_else(|_| {
300                            panic!("prev state is not the correct size");
301                        }),
302                    timestamp_pp(),
303                    &local.mem.state_reads[i],
304                )
305                .eval(builder, *local.instruction.is_enabled);
306        }
307    }
308
309    pub fn eval_writes<AB: InteractionBuilder>(
310        &self,
311        builder: &mut AB,
312        local: &Sha2ColsRef<AB::Var>,
313        timestamp_pp: &mut impl FnMut() -> AB::Expr,
314    ) {
315        let dst_ptr_val = compose(&local.instruction.dst_ptr_limbs.to_vec(), RV32_CELL_BITS);
316        for i in 0..C::STATE_READS {
317            self.memory_bridge
318                .write::<_, _, SHA2_WRITE_SIZE>(
319                    MemoryAddress::new(
320                        AB::Expr::from_u32(RV32_MEMORY_AS),
321                        dst_ptr_val.clone() + AB::F::from_usize(i * SHA2_READ_SIZE),
322                    ),
323                    local
324                        .block
325                        .new_state
326                        .slice(s![i * SHA2_READ_SIZE..(i + 1) * SHA2_READ_SIZE])
327                        .to_vec()
328                        .try_into()
329                        .unwrap_or_else(|_| {
330                            panic!("new state is not the correct size");
331                        }),
332                    timestamp_pp(),
333                    &local.mem.write_aux[i],
334                )
335                .eval(builder, *local.instruction.is_enabled);
336        }
337    }
338}
openvm_sha2_circuit/sha2_chips/main_chip/air.rs

openvm_sha2_circuit/sha2_chips/main_chip/
air.rs
