openvm_sha2_circuit/sha2_chips/block_hasher_chip/
air.rs

1use openvm_circuit_primitives::{bitwise_op_lookup::BitwiseOperationLookupBus, ColumnsAir, SubAir};
2use openvm_sha2_air::{compose, Sha2BlockHasherSubAir};
3use openvm_stark_backend::{
4    interaction::{BusIndex, InteractionBuilder, PermutationCheckBus},
5    p3_air::{Air, AirBuilder, BaseAir},
6    p3_field::{Field, PrimeCharacteristicRing},
7    p3_matrix::Matrix,
8    BaseAirWithPublicValues, PartitionedBaseAir,
9};
10
11use crate::{
12    MessageType, Sha2BlockHasherDigestColsRef, Sha2BlockHasherRoundColsRef,
13    Sha2BlockHasherVmConfig, INNER_OFFSET,
14};
15
16pub struct Sha2BlockHasherVmAir<C: Sha2BlockHasherVmConfig> {
17    pub inner: Sha2BlockHasherSubAir<C>,
18    pub sha2_bus: PermutationCheckBus,
19}
20// No columns provided: width is the config-dependent `C::BLOCK_HASHER_WIDTH` and rows are accessed
21// via `Sha2BlockHasher{Round,Digest}ColsRef` (slice-borrowing ref structs, no static `Cols`).
22impl<C: Sha2BlockHasherVmConfig> ColumnsAir for Sha2BlockHasherVmAir<C> {}
23
24impl<C: Sha2BlockHasherVmConfig> Sha2BlockHasherVmAir<C> {
25    pub fn new(
26        bitwise_lookup_bus: BitwiseOperationLookupBus,
27        inner_bus_idx: BusIndex,
28        sha2_bus_idx: BusIndex,
29    ) -> Self {
30        Self {
31            inner: Sha2BlockHasherSubAir::new(bitwise_lookup_bus, inner_bus_idx),
32            sha2_bus: PermutationCheckBus::new(sha2_bus_idx),
33        }
34    }
35}
36
37impl<F: Field, C: Sha2BlockHasherVmConfig> BaseAirWithPublicValues<F> for Sha2BlockHasherVmAir<C> {}
38impl<F: Field, C: Sha2BlockHasherVmConfig> PartitionedBaseAir<F> for Sha2BlockHasherVmAir<C> {}
39impl<F: Field, C: Sha2BlockHasherVmConfig> BaseAir<F> for Sha2BlockHasherVmAir<C> {
40    fn width(&self) -> usize {
41        C::BLOCK_HASHER_WIDTH
42    }
43}
44
45impl<AB: InteractionBuilder, C: Sha2BlockHasherVmConfig> Air<AB> for Sha2BlockHasherVmAir<C> {
46    fn eval(&self, builder: &mut AB) {
47        self.inner.eval(builder, INNER_OFFSET);
48        self.eval_interactions(builder);
49        self.eval_request_id(builder);
50    }
51}
52
53impl<C: Sha2BlockHasherVmConfig> Sha2BlockHasherVmAir<C> {
54    fn eval_interactions<AB: InteractionBuilder>(&self, builder: &mut AB) {
55        let main = builder.main();
56        let local_slice = main.row_slice(0).unwrap();
57        let next_slice = main.row_slice(1).unwrap();
58
59        let local = Sha2BlockHasherDigestColsRef::<AB::Var>::from::<C>(
60            &local_slice[..C::BLOCK_HASHER_DIGEST_WIDTH],
61        );
62
63        // Receive (STATE, request_id, prev_state_as_u16s, new_state) on the sha2 bus
64        self.sha2_bus.receive(
65            builder,
66            [
67                AB::Expr::from_u8(MessageType::State as u8),
68                (*local.request_id).into(),
69            ]
70            .into_iter()
71            .chain(local.inner.prev_hash.flatten().map(|x| (*x).into()))
72            .chain(local.inner.final_hash.flatten().map(|x| (*x).into())),
73            *local.inner.flags.is_digest_row,
74        );
75
76        let local = Sha2BlockHasherRoundColsRef::<AB::Var>::from::<C>(
77            &local_slice[..C::BLOCK_HASHER_ROUND_WIDTH],
78        );
79        let next = Sha2BlockHasherRoundColsRef::<AB::Var>::from::<C>(
80            &next_slice[..C::BLOCK_HASHER_ROUND_WIDTH],
81        );
82
83        let is_local_first_row = self
84            .inner
85            .row_idx_encoder
86            .contains_flag::<AB>(local.inner.flags.row_idx.to_slice().unwrap(), &[0]);
87
88        // Taken from old Sha256VmChip:
89        // https://github.com/openvm-org/openvm/blob/c2e376e6059c8bbf206736cf01d04cda43dfc42d/extensions/sha256/circuit/src/sha256_chip/air.rs#L310C1-L318C1
90        let get_ith_byte = |i: usize, cols: &Sha2BlockHasherRoundColsRef<AB::Var>| {
91            debug_assert!(i < C::WORD_U8S * C::ROUNDS_PER_ROW);
92            let row_idx = i / C::WORD_U8S;
93            let word: Vec<AB::Var> = cols
94                .inner
95                .message_schedule
96                .w
97                .row(row_idx)
98                .into_iter()
99                .copied()
100                .collect::<Vec<_>>();
101            // Need to reverse the byte order to match the endianness of the memory
102            let byte_idx = C::WORD_U8S - i % C::WORD_U8S - 1;
103            compose::<AB::Expr>(&word[byte_idx * 8..(byte_idx + 1) * 8], 1)
104        };
105
106        let local_message = (0..C::WORD_U8S * C::ROUNDS_PER_ROW).map(|i| get_ith_byte(i, &local));
107        let next_message = (0..C::WORD_U8S * C::ROUNDS_PER_ROW).map(|i| get_ith_byte(i, &next));
108
109        // Receive (MESSAGE_1, request_id, first_half_of_message) on the sha2 bus
110        self.sha2_bus.receive(
111            builder,
112            [
113                AB::Expr::from_u8(MessageType::Message1 as u8),
114                (*local.request_id).into(),
115            ]
116            .into_iter()
117            .chain(local_message.clone())
118            .chain(next_message.clone()),
119            is_local_first_row * local.inner.flags.is_not_padding_row(),
120        );
121
122        let is_local_third_row = self
123            .inner
124            .row_idx_encoder
125            .contains_flag::<AB>(local.inner.flags.row_idx.to_slice().unwrap(), &[2]);
126
127        // Send (MESSAGE_2, request_id, second_half_of_message) to the sha2 bus
128        self.sha2_bus.receive(
129            builder,
130            [
131                AB::Expr::from_u8(MessageType::Message2 as u8),
132                (*local.request_id).into(),
133            ]
134            .into_iter()
135            .chain(local_message)
136            .chain(next_message),
137            is_local_third_row * local.inner.flags.is_not_padding_row(),
138        );
139    }
140
141    fn eval_request_id<AB: InteractionBuilder>(&self, builder: &mut AB) {
142        let main = builder.main();
143        let local = main.row_slice(0).unwrap();
144        let next = main.row_slice(1).unwrap();
145
146        // doesn't matter if we use round or digest cols here, since we only access
147        // request_id and inner.flags.is_last block, which are common to both
148        // field
149        let local =
150            Sha2BlockHasherRoundColsRef::<AB::Var>::from::<C>(&local[..C::BLOCK_HASHER_WIDTH]);
151        let next =
152            Sha2BlockHasherRoundColsRef::<AB::Var>::from::<C>(&next[..C::BLOCK_HASHER_WIDTH]);
153
154        builder
155            .when_transition()
156            .when(*local.inner.flags.is_round_row)
157            .assert_eq(*next.request_id, *local.request_id);
158    }
159}
openvm_sha2_circuit/sha2_chips/block_hasher_chip/air.rs

openvm_sha2_circuit/sha2_chips/block_hasher_chip/
air.rs
