openvm_native_recursion/challenger/
multi_field32.rs1use 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