openvm_continuations/circuit/subair/hash_slice/
air.rs

1use std::array::from_fn;
2
3use openvm_circuit::arch::POSEIDON2_WIDTH;
4use openvm_circuit_primitives::SubAir;
5use openvm_recursion_circuit::bus::{
6    Poseidon2CompressBus, Poseidon2CompressMessage, Poseidon2PermuteBus, Poseidon2PermuteMessage,
7};
8use openvm_stark_backend::interaction::InteractionBuilder;
9use openvm_stark_sdk::config::baby_bear_poseidon2::DIGEST_SIZE;
10use p3_air::AirBuilder;
11use p3_field::PrimeCharacteristicRing;
12
13use crate::utils::digests_to_poseidon2_input;
14
15/// SubAir to hash N digest-sized elements into a single digest using a chain of Poseidon2
16/// operations: N−1 permutations followed by 1 compression. Each step absorbs one element
17/// into the capacity portion of the running state.
18#[derive(Clone, Debug, derive_new::new)]
19pub struct HashSliceSubAir {
20    pub compress_bus: Poseidon2CompressBus,
21    pub permute_bus: Poseidon2PermuteBus,
22}
23
24pub struct HashSliceCtx<'a, T> {
25    // N elements
26    pub elements: &'a [[T; DIGEST_SIZE]],
27    // N - 1 intermediate states
28    pub intermediate: &'a [[T; POSEIDON2_WIDTH]],
29    // 1 final result
30    pub result: &'a [T; DIGEST_SIZE],
31    // Boolean value that is true iff enabled
32    pub enabled: &'a T,
33}
34
35impl<AB: AirBuilder + InteractionBuilder> SubAir<AB> for HashSliceSubAir {
36    type AirContext<'a>
37        = HashSliceCtx<'a, AB::Expr>
38    where
39        AB: 'a,
40        AB::Var: 'a,
41        AB::Expr: 'a;
42
43    fn eval<'a>(&'a self, builder: &'a mut AB, ctx: Self::AirContext<'a>)
44    where
45        AB::Var: 'a,
46        AB::Expr: 'a,
47    {
48        let n = ctx.elements.len();
49
50        assert_eq!(n, ctx.intermediate.len() + 1);
51        assert!(n > 1);
52
53        self.permute_bus.lookup_key(
54            builder,
55            Poseidon2PermuteMessage {
56                input: digests_to_poseidon2_input(
57                    ctx.elements[0].clone(),
58                    [AB::Expr::ZERO; DIGEST_SIZE],
59                ),
60                output: ctx.intermediate[0].clone(),
61            },
62            ctx.enabled.clone(),
63        );
64
65        for i in 1..(n - 1) {
66            self.permute_bus.lookup_key(
67                builder,
68                Poseidon2PermuteMessage {
69                    input: digests_to_poseidon2_input(
70                        ctx.elements[i].clone(),
71                        from_fn(|j| ctx.intermediate[i - 1][j + DIGEST_SIZE].clone()),
72                    ),
73                    output: ctx.intermediate[i].clone(),
74                },
75                ctx.enabled.clone(),
76            );
77        }
78
79        self.compress_bus.lookup_key(
80            builder,
81            Poseidon2CompressMessage {
82                input: digests_to_poseidon2_input(
83                    ctx.elements[n - 1].clone(),
84                    from_fn(|j| ctx.intermediate[n - 2][j + DIGEST_SIZE].clone()),
85                ),
86                output: ctx.result.clone(),
87            },
88            ctx.enabled.clone(),
89        );
90    }
91}
openvm_continuations/circuit/subair/hash_slice/air.rs

openvm_continuations/circuit/subair/hash_slice/
air.rs
