openvm_keccak256_circuit/keccakf_op/
air.rs

1use std::{borrow::Borrow, iter};
2
3use itertools::izip;
4use openvm_circuit::{
5    arch::{ExecutionBridge, ExecutionState},
6    system::memory::{
7        offline_checker::{MemoryBridge, MemoryWriteAuxCols},
8        MemoryAddress,
9    },
10};
11use openvm_circuit_primitives::{bitwise_op_lookup::BitwiseOperationLookupBus, ColumnsAir};
12use openvm_instructions::riscv::{
13    RV32_CELL_BITS, RV32_MEMORY_AS, RV32_REGISTER_AS, RV32_REGISTER_NUM_LIMBS,
14};
15use openvm_keccak256_transpiler::KeccakfOpcode;
16use openvm_rv32im_circuit::adapters::abstract_compose;
17use openvm_stark_backend::{
18    interaction::{InteractionBuilder, PermutationCheckBus},
19    p3_air::{Air, BaseAir},
20    p3_field::PrimeCharacteristicRing,
21    p3_matrix::Matrix,
22    BaseAirWithPublicValues, PartitionedBaseAir,
23};
24
25use crate::{
26    keccakf_op::columns::{KeccakfOpCols, NUM_KECCAKF_OP_COLS},
27    KECCAK_WORD_SIZE,
28};
29
30#[derive(Clone, Copy, Debug, derive_new::new, ColumnsAir)]
31#[columns_via(KeccakfOpCols<u8>)]
32pub struct KeccakfOpAir {
33    pub execution_bridge: ExecutionBridge,
34    pub memory_bridge: MemoryBridge,
35    pub bitwise_lookup_bus: BitwiseOperationLookupBus,
36    /// Direct bus with keccakf pre- or post-state. Bus message is
37    /// ```text
38    /// is_post || timestamp || state_u16_limbs
39    /// ```
40    pub keccakf_state_bus: PermutationCheckBus,
41    pub ptr_max_bits: usize,
42    pub(super) offset: usize,
43}
44
45impl<F> BaseAirWithPublicValues<F> for KeccakfOpAir {}
46impl<F> PartitionedBaseAir<F> for KeccakfOpAir {}
47impl<F> BaseAir<F> for KeccakfOpAir {
48    fn width(&self) -> usize {
49        NUM_KECCAKF_OP_COLS
50    }
51}
52
53impl<AB: InteractionBuilder> Air<AB> for KeccakfOpAir {
54    fn eval(&self, builder: &mut AB) {
55        let main = builder.main();
56
57        let local = main.row_slice(0).unwrap();
58        let local: &KeccakfOpCols<_> = (*local).borrow();
59
60        let is_valid = local.is_valid;
61        builder.assert_bool(is_valid);
62
63        let start_timestamp = local.timestamp;
64        let mut timestamp_delta = 0usize;
65        let mut timestamp_pp = || {
66            timestamp_delta += 1;
67            start_timestamp + AB::F::from_usize(timestamp_delta - 1)
68        };
69        // ======== Read `rd` =========
70        let rd_ptr = local.rd_ptr;
71        let buffer_ptr_limbs = local.buffer_ptr_limbs;
72        self.memory_bridge
73            .read(
74                MemoryAddress::new(AB::F::from_u32(RV32_REGISTER_AS), rd_ptr),
75                buffer_ptr_limbs,
76                timestamp_pp(),
77                &local.rd_aux,
78            )
79            .eval(builder, is_valid);
80        // Range check that buffer_ptr_limbs fits in [0, 2^ptr_max_bits) as u32
81        {
82            assert!(self.ptr_max_bits >= RV32_CELL_BITS * (RV32_REGISTER_NUM_LIMBS - 1));
83            let limb_shift = AB::F::from_usize(
84                1 << (RV32_CELL_BITS * RV32_REGISTER_NUM_LIMBS - self.ptr_max_bits),
85            );
86            let need_range_check = [
87                buffer_ptr_limbs[RV32_REGISTER_NUM_LIMBS - 1],
88                buffer_ptr_limbs[RV32_REGISTER_NUM_LIMBS - 1],
89            ];
90            for pair in need_range_check.chunks_exact(2) {
91                self.bitwise_lookup_bus
92                    .send_range(pair[0] * limb_shift, pair[1] * limb_shift)
93                    .eval(builder, is_valid);
94            }
95        }
96        // Now it is safe to cast buffer_ptr to F
97        let buffer_ptr: AB::Expr = abstract_compose(local.buffer_ptr_limbs);
98
99        // ======== Constrain that post-state consists of bytes =========
100        // We know that the pre-state buffer consists of bytes due to the invariant of Address Space
101        // 2 in memory. The keccakf_state_bus guarantees that the post-state consists of
102        // u16, but we still need to constrain that each pair actually consists of bytes.
103        // NOTE[jpw]: this can be removed if AS2 cells are changed to u16s
104        for pair in local.postimage.chunks_exact(2) {
105            self.bitwise_lookup_bus
106                .send_range(pair[0], pair[1])
107                .eval(builder, is_valid);
108        }
109
110        // ======== Constrain new writes of `buffer` to memory =========
111        // NOTE: we use the _next_ row's `buffer` as the pre-state
112        for (word_idx, (prev_word, post_word, base_aux)) in izip!(
113            local.preimage.chunks_exact(KECCAK_WORD_SIZE),
114            local.postimage.chunks_exact(KECCAK_WORD_SIZE),
115            local.buffer_word_aux
116        )
117        .enumerate()
118        {
119            // Safety:
120            // - we range checked that buffer_ptr < 2^ptr_max_bits but not that buffer_ptr +
121            //   KECCAK_WIDTH_BYTES is in range.
122            // - the previous range check implies `buffer_ptr + KECCAK_WIDTH_BYTES` does not
123            //   overflow the field `F` hence it is safe to consider `ptr` as a field element.
124            // - the memory_bridge.write at `ptr` consists of a receive on memory bus at a previous
125            //   timestamp. The only way this bus interaction could balance is if there was already
126            //   a previous valid write at `ptr`. Assuming the invariant that all previous memory
127            //   accesses are valid and timestamp always moves forward, the new write to `ptr` must
128            //   be valid as well.
129            let ptr = buffer_ptr.clone() + AB::F::from_usize(word_idx * KECCAK_WORD_SIZE);
130            let prev_data: &[_; KECCAK_WORD_SIZE] = prev_word.try_into().unwrap();
131            // post_word consists of bytes due to range checks above
132            let data: &[_; KECCAK_WORD_SIZE] = post_word.try_into().unwrap();
133            let write_aux = MemoryWriteAuxCols {
134                base: base_aux,
135                prev_data: *prev_data,
136            };
137            self.memory_bridge
138                .write(
139                    MemoryAddress::new(AB::F::from_u32(RV32_MEMORY_AS), ptr),
140                    *data,
141                    timestamp_pp(),
142                    &write_aux,
143                )
144                .eval(builder, is_valid);
145        }
146
147        // ======== Execution bus =========
148        self.execution_bridge
149            .execute_and_increment_pc(
150                AB::Expr::from_usize(KeccakfOpcode::KECCAKF as usize + self.offset),
151                [
152                    rd_ptr.into(),
153                    AB::Expr::ZERO,
154                    AB::Expr::ZERO,
155                    AB::Expr::from_u32(RV32_REGISTER_AS),
156                    AB::Expr::from_u32(RV32_MEMORY_AS),
157                ],
158                ExecutionState::new(local.pc, local.timestamp),
159                AB::F::from_usize(timestamp_delta),
160            )
161            .eval(builder, is_valid);
162
163        // ======== KeccakF State Interaction =======
164        // Now we actually constrain that the pre- and post- buffer values are valid, but doing a
165        // permutation check with the KeccakFPeripheryAir. We compose two u8 into a u16
166        // since the keccakf periphery air uses u16 limbs
167        //
168        // We use two interactions bound with the same timestamp to avoid having a really large
169        // message length.
170        self.keccakf_state_bus.send(
171            builder,
172            iter::empty()
173                .chain([AB::Expr::ZERO, local.timestamp.into()])
174                .chain(
175                    local
176                        .preimage
177                        .chunks(2)
178                        .map(|pair| pair[0] + pair[1] * AB::F::from_u32(256)),
179                ),
180            is_valid,
181        );
182        self.keccakf_state_bus.send(
183            builder,
184            iter::empty()
185                .chain([AB::Expr::ONE, local.timestamp.into()])
186                .chain(
187                    local
188                        .postimage
189                        .chunks(2)
190                        .map(|pair| pair[0] + pair[1] * AB::F::from_u32(256)),
191                ),
192            is_valid,
193        );
194    }
195}
openvm_keccak256_circuit/keccakf_op/air.rs

openvm_keccak256_circuit/keccakf_op/
air.rs
