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, PrimeCharacteristicRing};
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        if bits == 0 {
101            return;
102        }
103        self.observe(builder, witness);
104        let element = self.sample_bits(builder, bits);
105        builder.assert_var_eq(element, C::N::from_usize(0));
106    }
107}
108
109impl<C: Config> CanObserveVariable<C, Felt<C::F>> for MultiField32ChallengerVariable<C> {
110    fn observe(&mut self, builder: &mut Builder<C>, value: Felt<C::F>) {
111        MultiField32ChallengerVariable::observe(self, builder, value);
112    }
113
114    fn observe_slice(&mut self, builder: &mut Builder<C>, values: Array<C, Felt<C::F>>) {
115        values.vec().into_iter().for_each(|value| {
116            self.observe(builder, value);
117        });
118    }
119}
120
121impl<C: Config> CanSampleVariable<C, Felt<C::F>> for MultiField32ChallengerVariable<C> {
122    fn sample(&mut self, builder: &mut Builder<C>) -> Felt<C::F> {
123        MultiField32ChallengerVariable::sample(self, builder)
124    }
125}
126
127impl<C: Config> CanSampleBitsVariable<C> for MultiField32ChallengerVariable<C> {
128    fn sample_bits(
129        &mut self,
130        builder: &mut Builder<C>,
131        nb_bits: RVar<C::N>,
132    ) -> Array<C, Var<C::N>> {
133        let rand_f = self.sample(builder);
134        let rand_f_bits = builder.num2bits_f_circuit(rand_f);
135        builder.vec(rand_f_bits[..nb_bits.value()].to_vec())
136    }
137}
138
139impl<C: Config> CanObserveDigest<C> for MultiField32ChallengerVariable<C> {
140    fn observe_digest(&mut self, builder: &mut Builder<C>, commitment: DigestVariable<C>) {
141        if let DigestVariable::Var(v_commit) = commitment {
142            MultiField32ChallengerVariable::observe_commitment(
143                self,
144                builder,
145                v_commit.vec().try_into().unwrap(),
146            );
147        } else {
148            panic!("MultiField32ChallengerVariable expects Var commitment");
149        }
150    }
151}
152
153impl<C: Config> FeltChallenger<C> for MultiField32ChallengerVariable<C> {
154    fn sample_ext(&mut self, builder: &mut Builder<C>) -> Ext<C::F, C::EF> {
155        MultiField32ChallengerVariable::sample_ext(self, builder)
156    }
157}
158
159impl<C: Config> CanCheckWitness<C> for MultiField32ChallengerVariable<C> {
160    fn check_witness(&mut self, builder: &mut Builder<C>, nb_bits: usize, witness: Felt<C::F>) {
161        MultiField32ChallengerVariable::check_witness(self, builder, nb_bits, witness);
162    }
163}
164
165impl<C: Config> ChallengerVariable<C> for MultiField32ChallengerVariable<C> {
166    fn new(builder: &mut Builder<C>) -> Self {
167        MultiField32ChallengerVariable::new(builder)
168    }
169}
170// Testing depends on halo2. Put it inside src/halo2/tests/multi_field32.rs