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