openvm_continuations/verifier/
utils.rs

1use std::array;
2
3use openvm_native_compiler::prelude::*;
4use openvm_native_recursion::{config::outer::OuterConfig, hints::Hintable, types::InnerConfig};
5use openvm_stark_backend::p3_field::PrimeField32;
6use openvm_stark_sdk::{
7    openvm_stark_backend::p3_field::FieldAlgebra, p3_baby_bear::BabyBear, p3_bn254_fr::Bn254Fr,
8};
9
10pub fn compress_babybear_var_to_bn254(
11    builder: &mut Builder<OuterConfig>,
12    var: [Var<Bn254Fr>; DIGEST_SIZE],
13) -> Var<Bn254Fr> {
14    let mut ret = SymbolicVar::ZERO;
15    let order = Bn254Fr::from_canonical_u32(BabyBear::ORDER_U32);
16    let mut base = Bn254Fr::ONE;
17    var.iter().for_each(|&x| {
18        ret += x * base;
19        base *= order;
20    });
21    builder.eval(ret)
22}
23
24pub(crate) fn assign_array_to_slice<C: Config>(
25    builder: &mut Builder<C>,
26    dst_slice: &[Felt<C::F>],
27    src: &Array<C, Felt<C::F>>,
28    src_offset: usize,
29) {
30    for (i, dst) in dst_slice.iter().enumerate() {
31        let pv = builder.get(src, i + src_offset);
32        builder.assign(dst, pv);
33    }
34}
35
36pub(crate) fn assign_slice_to_array<C: Config>(
37    builder: &mut Builder<C>,
38    dst: &Array<C, Felt<C::F>>,
39    src_slice: &[Felt<C::F>],
40) {
41    for (i, &src) in src_slice.iter().enumerate() {
42        builder.set_value(dst, i, src);
43    }
44}
45
46pub(crate) fn write_field_slice(arr: &[BabyBear; DIGEST_SIZE]) -> Vec<Vec<BabyBear>> {
47    arr.iter()
48        .flat_map(Hintable::<InnerConfig>::write)
49        .collect()
50}
51
52/// Returns 1 if lhs == rhs, 0 otherwise.
53pub(crate) fn eq_felt_slice<C: Config, const N: usize>(
54    builder: &mut Builder<C>,
55    lhs: &[Felt<C::F>; N],
56    rhs: &[Felt<C::F>; N],
57) -> Var<C::N> {
58    let sub_res: [Felt<C::F>; N] = array::from_fn(|i| builder.eval(lhs[i] - rhs[i]));
59    let var_res = sub_res.map(|f| builder.cast_felt_to_var(f));
60    let ret: Var<C::N> = builder.eval(C::N::ONE);
61    var_res.into_iter().for_each(|v| {
62        builder
63            .if_ne(v, C::N::ZERO)
64            .then(|builder| builder.assign(&ret, C::N::ZERO))
65    });
66    ret
67}
68
69#[derive(Clone)]
70pub(crate) struct VariableP2Compressor<C: Config> {
71    state: Array<C, Felt<C::F>>,
72    lhs: Array<C, Felt<C::F>>,
73    rhs: Array<C, Felt<C::F>>,
74}
75
76impl<C: Config> VariableP2Compressor<C> {
77    pub fn new(builder: &mut Builder<C>) -> Self {
78        Self {
79            state: builder.array(PERMUTATION_WIDTH),
80            lhs: builder.array(DIGEST_SIZE),
81            rhs: builder.array(DIGEST_SIZE),
82        }
83    }
84
85    pub fn compress(
86        &self,
87        builder: &mut Builder<C>,
88        lhs: &[Felt<C::F>; DIGEST_SIZE],
89        rhs: &[Felt<C::F>; DIGEST_SIZE],
90    ) -> [Felt<C::F>; DIGEST_SIZE] {
91        assign_slice_to_array(builder, &self.lhs, lhs);
92        assign_slice_to_array(builder, &self.rhs, rhs);
93        builder.poseidon2_compress_x(&self.state, &self.lhs, &self.rhs);
94        array::from_fn(|i| builder.get(&self.state, i))
95    }
96
97    pub fn compress_array(
98        &self,
99        builder: &mut Builder<C>,
100        lhs: &Array<C, Felt<C::F>>,
101        rhs: &Array<C, Felt<C::F>>,
102    ) -> Array<C, Felt<C::F>> {
103        let ret = builder.array(DIGEST_SIZE);
104        builder.poseidon2_compress_x(&self.state, lhs, rhs);
105        for i in 0..DIGEST_SIZE {
106            let v = builder.get(&self.state, i);
107            builder.set_value(&ret, i, v);
108        }
109        ret
110    }
111}
112
113#[derive(Clone)]
114pub(crate) struct VariableP2Hasher<C: Config> {
115    pub compressor: VariableP2Compressor<C>,
116    pub const_zeros: Array<C, Felt<C::F>>,
117    pub const_zero: Felt<C::F>,
118}
119
120impl<C: Config> VariableP2Hasher<C> {
121    pub fn new(builder: &mut Builder<C>) -> Self {
122        let const_zero: Felt<C::F> = builder.eval(C::F::ZERO);
123        let const_zeros = builder.array(DIGEST_SIZE);
124        for i in 0..DIGEST_SIZE {
125            builder.set_value(&const_zeros, i, const_zero);
126        }
127        Self {
128            compressor: VariableP2Compressor::new(builder),
129            const_zeros,
130            const_zero,
131        }
132    }
133    pub fn hash(
134        &self,
135        builder: &mut Builder<C>,
136        payload: &[Felt<C::F>; DIGEST_SIZE],
137    ) -> [Felt<C::F>; DIGEST_SIZE] {
138        self.compressor
139            .compress(builder, payload, &[self.const_zero; DIGEST_SIZE])
140    }
141    pub fn hash_array(
142        &self,
143        builder: &mut Builder<C>,
144        payload: &Array<C, Felt<C::F>>,
145    ) -> Array<C, Felt<C::F>> {
146        self.compressor
147            .compress_array(builder, payload, &self.const_zeros)
148    }
149
150    pub fn merkle_root(
151        &self,
152        builder: &mut Builder<C>,
153        values: &[Felt<C::F>],
154    ) -> [Felt<C::F>; DIGEST_SIZE] {
155        assert_eq!(values.len() % DIGEST_SIZE, 0);
156        assert!((values.len() / DIGEST_SIZE).is_power_of_two());
157        let buffer = builder.array(DIGEST_SIZE);
158        let mut leaves: Vec<_> = values
159            .chunks_exact(DIGEST_SIZE)
160            .map(|chunk| {
161                assign_slice_to_array(builder, &buffer, chunk);
162                self.hash_array(builder, &buffer)
163            })
164            .collect();
165        while leaves.len() > 1 {
166            leaves = leaves
167                .chunks_exact(2)
168                .map(|chunk| {
169                    self.compressor
170                        .compress_array(builder, &chunk[0], &chunk[1])
171                })
172                .collect();
173        }
174        array::from_fn(|i| builder.get(&leaves[0], i))
175    }
176}