openvm_native_recursion/challenger/
multi_field32.rs

1use openvm_native_compiler::ir::{Array, Builder, Config, Ext, Felt, RVar, Var};
2use openvm_stark_backend::p3_field::{Field, FieldAlgebra};
3
4use crate::{
5    challenger::{
6        CanCheckWitness, CanObserveDigest, CanObserveVariable, CanSampleBitsVariable,
7        CanSampleVariable, ChallengerVariable, FeltChallenger,
8    },
9    digest::DigestVariable,
10    outer_poseidon2::{Poseidon2CircuitBuilder, RATE, SPONGE_SIZE},
11    utils::{reduce_32, split_32},
12    vars::OuterDigestVariable,
13};
14
15#[derive(Clone)]
16pub struct MultiField32ChallengerVariable<C: Config> {
17    sponge_state: [Var<C::N>; SPONGE_SIZE],
18    input_buffer: Vec<Felt<C::F>>,
19    output_buffer: Vec<Felt<C::F>>,
20    num_f_elms: usize,
21}
22
23impl<C: Config> MultiField32ChallengerVariable<C> {
24    pub fn new(builder: &mut Builder<C>) -> Self {
25        assert!(builder.flags.static_only);
26        MultiField32ChallengerVariable::<C> {
27            sponge_state: [
28                builder.eval(C::N::ZERO),
29                builder.eval(C::N::ZERO),
30                builder.eval(C::N::ZERO),
31            ],
32            input_buffer: vec![],
33            output_buffer: vec![],
34            num_f_elms: C::N::bits() / 64,
35        }
36    }
37
38    pub fn duplexing(&mut self, builder: &mut Builder<C>) {
39        assert!(self.input_buffer.len() <= self.num_f_elms * RATE);
40
41        for (i, f_chunk) in self.input_buffer.chunks(self.num_f_elms).enumerate() {
42            self.sponge_state[i] = reduce_32(builder, f_chunk);
43        }
44        self.input_buffer.clear();
45
46        builder.p2_permute_mut(self.sponge_state);
47
48        self.output_buffer.clear();
49        for &pf_val in self.sponge_state.iter() {
50            let f_vals = split_32(builder, pf_val, self.num_f_elms);
51            for f_val in f_vals {
52                self.output_buffer.push(f_val);
53            }
54        }
55    }
56
57    pub fn observe(&mut self, builder: &mut Builder<C>, value: Felt<C::F>) {
58        self.output_buffer.clear();
59
60        self.input_buffer.push(value);
61        if self.input_buffer.len() == self.num_f_elms * RATE {
62            self.duplexing(builder);
63        }
64    }
65
66    pub fn observe_commitment(&mut self, builder: &mut Builder<C>, value: OuterDigestVariable<C>) {
67        value.into_iter().for_each(|v| {
68            let f_vals: Vec<Felt<C::F>> = split_32(builder, v, self.num_f_elms);
69            for f_val in f_vals {
70                self.observe(builder, f_val);
71            }
72        });
73    }
74
75    pub fn sample(&mut self, builder: &mut Builder<C>) -> Felt<C::F> {
76        if !self.input_buffer.is_empty() || self.output_buffer.is_empty() {
77            self.duplexing(builder);
78        }
79
80        self.output_buffer
81            .pop()
82            .expect("output buffer should be non-empty")
83    }
84
85    pub fn sample_ext(&mut self, builder: &mut Builder<C>) -> Ext<C::F, C::EF> {
86        let a = self.sample(builder);
87        let b = self.sample(builder);
88        let c = self.sample(builder);
89        let d = self.sample(builder);
90        builder.felts2ext(&[a, b, c, d])
91    }
92
93    pub fn sample_bits(&mut self, builder: &mut Builder<C>, bits: usize) -> Var<C::N> {
94        let rand_f = self.sample(builder);
95        let rand_f_bits = builder.num2bits_f_circuit(rand_f);
96        builder.bits2num_v_circuit(&rand_f_bits[0..bits])
97    }
98
99    pub fn check_witness(&mut self, builder: &mut Builder<C>, bits: usize, witness: Felt<C::F>) {
100        self.observe(builder, witness);
101        let element = self.sample_bits(builder, bits);
102        builder.assert_var_eq(element, C::N::from_canonical_usize(0));
103    }
104}
105
106impl<C: Config> CanObserveVariable<C, Felt<C::F>> for MultiField32ChallengerVariable<C> {
107    fn observe(&mut self, builder: &mut Builder<C>, value: Felt<C::F>) {
108        MultiField32ChallengerVariable::observe(self, builder, value);
109    }
110
111    fn observe_slice(&mut self, builder: &mut Builder<C>, values: Array<C, Felt<C::F>>) {
112        values.vec().into_iter().for_each(|value| {
113            self.observe(builder, value);
114        });
115    }
116}
117
118impl<C: Config> CanSampleVariable<C, Felt<C::F>> for MultiField32ChallengerVariable<C> {
119    fn sample(&mut self, builder: &mut Builder<C>) -> Felt<C::F> {
120        MultiField32ChallengerVariable::sample(self, builder)
121    }
122}
123
124impl<C: Config> CanSampleBitsVariable<C> for MultiField32ChallengerVariable<C> {
125    fn sample_bits(
126        &mut self,
127        builder: &mut Builder<C>,
128        nb_bits: RVar<C::N>,
129    ) -> Array<C, Var<C::N>> {
130        let rand_f = self.sample(builder);
131        let rand_f_bits = builder.num2bits_f_circuit(rand_f);
132        builder.vec(rand_f_bits[..nb_bits.value()].to_vec())
133    }
134}
135
136impl<C: Config> CanObserveDigest<C> for MultiField32ChallengerVariable<C> {
137    fn observe_digest(&mut self, builder: &mut Builder<C>, commitment: DigestVariable<C>) {
138        if let DigestVariable::Var(v_commit) = commitment {
139            MultiField32ChallengerVariable::observe_commitment(
140                self,
141                builder,
142                v_commit.vec().try_into().unwrap(),
143            );
144        } else {
145            panic!("MultiField32ChallengerVariable expects Var commitment");
146        }
147    }
148}
149
150impl<C: Config> FeltChallenger<C> for MultiField32ChallengerVariable<C> {
151    fn sample_ext(&mut self, builder: &mut Builder<C>) -> Ext<C::F, C::EF> {
152        MultiField32ChallengerVariable::sample_ext(self, builder)
153    }
154}
155
156impl<C: Config> CanCheckWitness<C> for MultiField32ChallengerVariable<C> {
157    fn check_witness(&mut self, builder: &mut Builder<C>, nb_bits: usize, witness: Felt<C::F>) {
158        MultiField32ChallengerVariable::check_witness(self, builder, nb_bits, witness);
159    }
160}
161
162impl<C: Config> ChallengerVariable<C> for MultiField32ChallengerVariable<C> {
163    fn new(builder: &mut Builder<C>) -> Self {
164        MultiField32ChallengerVariable::new(builder)
165    }
166}
167// Testing depends on halo2. Put it inside src/halo2/tests/multi_field32.rs