1use std::{
2 array,
3 iter::once,
4 ops::{Add, Mul, Sub},
5};
6
7use itertools::Itertools;
8use serde::{Deserialize, Serialize};
9use slop_air::{AirBuilder, AirBuilderWithPublicValues, FilteredAirBuilder};
10use slop_algebra::{AbstractField, Field};
11use slop_uni_stark::{
12 ProverConstraintFolder, StarkGenericConfig, SymbolicAirBuilder, VerifierConstraintFolder,
13};
14use strum::{Display, EnumIter};
15
16use super::{interaction::AirInteraction, BinomialExtension};
17use crate::{lookup::InteractionKind, septic_extension::SepticExtension, ConstraintSumcheckFolder};
18
19#[derive(
21 Debug,
22 Clone,
23 Copy,
24 PartialEq,
25 Eq,
26 Hash,
27 Display,
28 EnumIter,
29 PartialOrd,
30 Ord,
31 Serialize,
32 Deserialize,
33)]
34pub enum InteractionScope {
35 Global = 0,
37 Local,
39}
40
41pub trait MessageBuilder<M> {
43 fn send(&mut self, message: M, scope: InteractionScope);
45
46 fn receive(&mut self, message: M, scope: InteractionScope);
48}
49
50pub trait EmptyMessageBuilder: AirBuilder {}
52
53impl<AB: EmptyMessageBuilder, M> MessageBuilder<M> for AB {
54 fn send(&mut self, _message: M, _scope: InteractionScope) {}
55
56 fn receive(&mut self, _message: M, _scope: InteractionScope) {}
57}
58
59pub trait BaseAirBuilder: AirBuilder + MessageBuilder<AirInteraction<Self::Expr>> {
61 fn when_not<I: Into<Self::Expr>>(&mut self, condition: I) -> FilteredAirBuilder<'_, Self> {
63 self.when_ne(condition, Self::F::one())
64 }
65
66 fn assert_all_eq<I1: Into<Self::Expr>, I2: Into<Self::Expr>>(
68 &mut self,
69 left: impl IntoIterator<Item = I1>,
70 right: impl IntoIterator<Item = I2>,
71 ) {
72 for (left, right) in left.into_iter().zip_eq(right) {
73 self.assert_eq(left, right);
74 }
75 }
76
77 fn assert_all_zero<I: Into<Self::Expr>>(&mut self, iter: impl IntoIterator<Item = I>) {
79 iter.into_iter().for_each(|expr| self.assert_zero(expr));
80 }
81
82 #[inline]
85 fn if_else(
86 &mut self,
87 condition: impl Into<Self::Expr> + Clone,
88 a: impl Into<Self::Expr> + Clone,
89 b: impl Into<Self::Expr> + Clone,
90 ) -> Self::Expr {
91 condition.clone().into() * a.into() + (Self::Expr::one() - condition.into()) * b.into()
92 }
93
94 fn index_array(
97 &mut self,
98 array: &[impl Into<Self::Expr> + Clone],
99 index_bitmap: &[impl Into<Self::Expr> + Clone],
100 ) -> Self::Expr {
101 let mut result = Self::Expr::zero();
102
103 for (value, i) in array.iter().zip_eq(index_bitmap) {
104 result = result.clone() + value.clone().into() * i.clone().into();
105 }
106
107 result
108 }
109}
110
111pub trait ByteAirBuilder: BaseAirBuilder {
113 #[allow(clippy::too_many_arguments)]
115 fn send_byte(
116 &mut self,
117 opcode: impl Into<Self::Expr>,
118 a: impl Into<Self::Expr>,
119 b: impl Into<Self::Expr>,
120 c: impl Into<Self::Expr>,
121 multiplicity: impl Into<Self::Expr>,
122 ) {
123 self.send(
124 AirInteraction::new(
125 vec![opcode.into(), a.into(), b.into(), c.into()],
126 multiplicity.into(),
127 InteractionKind::Byte,
128 ),
129 InteractionScope::Local,
130 );
131 }
132
133 #[allow(clippy::too_many_arguments)]
135 fn receive_byte(
136 &mut self,
137 opcode: impl Into<Self::Expr>,
138 a: impl Into<Self::Expr>,
139 b: impl Into<Self::Expr>,
140 c: impl Into<Self::Expr>,
141 multiplicity: impl Into<Self::Expr>,
142 ) {
143 self.receive(
144 AirInteraction::new(
145 vec![opcode.into(), a.into(), b.into(), c.into()],
146 multiplicity.into(),
147 InteractionKind::Byte,
148 ),
149 InteractionScope::Local,
150 );
151 }
152}
153
154pub trait InstructionAirBuilder: BaseAirBuilder {
156 fn send_state(
158 &mut self,
159 clk_high: impl Into<Self::Expr>,
160 clk_low: impl Into<Self::Expr>,
161 pc: [impl Into<Self::Expr>; 3],
162 multiplicity: impl Into<Self::Expr>,
163 ) {
164 self.send(
165 AirInteraction::new(
166 once(clk_high.into())
167 .chain(once(clk_low.into()))
168 .chain(pc.map(Into::into))
169 .collect::<Vec<_>>(),
170 multiplicity.into(),
171 InteractionKind::State,
172 ),
173 InteractionScope::Local,
174 );
175 }
176
177 fn receive_state(
179 &mut self,
180 clk_high: impl Into<Self::Expr>,
181 clk_low: impl Into<Self::Expr>,
182 pc: [impl Into<Self::Expr>; 3],
183 multiplicity: impl Into<Self::Expr>,
184 ) {
185 self.receive(
186 AirInteraction::new(
187 once(clk_high.into())
188 .chain(once(clk_low.into()))
189 .chain(pc.map(Into::into))
190 .collect::<Vec<_>>(),
191 multiplicity.into(),
192 InteractionKind::State,
193 ),
194 InteractionScope::Local,
195 );
196 }
197
198 #[allow(clippy::too_many_arguments)]
200 fn send_syscall(
201 &mut self,
202 clk_high: impl Into<Self::Expr> + Clone,
203 clk_low: impl Into<Self::Expr> + Clone,
204 syscall_id: impl Into<Self::Expr> + Clone,
205 arg1: [impl Into<Self::Expr>; 3],
206 arg2: [impl Into<Self::Expr>; 3],
207 multiplicity: impl Into<Self::Expr>,
208 scope: InteractionScope,
209 ) {
210 let values = once(clk_high.into())
211 .chain(once(clk_low.into()))
212 .chain(once(syscall_id.into()))
213 .chain(arg1.map(Into::into))
214 .chain(arg2.map(Into::into))
215 .collect::<Vec<_>>();
216
217 self.send(
218 AirInteraction::new(values, multiplicity.into(), InteractionKind::Syscall),
219 scope,
220 );
221 }
222
223 #[allow(clippy::too_many_arguments)]
225 fn receive_syscall(
226 &mut self,
227 clk_high: impl Into<Self::Expr> + Clone,
228 clk_low: impl Into<Self::Expr> + Clone,
229 syscall_id: impl Into<Self::Expr> + Clone,
230 arg1: [Self::Expr; 3],
231 arg2: [Self::Expr; 3],
232 multiplicity: impl Into<Self::Expr>,
233 scope: InteractionScope,
234 ) {
235 let values = once(clk_high.into())
236 .chain(once(clk_low.into()))
237 .chain(once(syscall_id.into()))
238 .chain(arg1)
239 .chain(arg2)
240 .collect::<Vec<_>>();
241
242 self.receive(
243 AirInteraction::new(values, multiplicity.into(), InteractionKind::Syscall),
244 scope,
245 );
246 }
247}
248
249pub trait ExtensionAirBuilder: BaseAirBuilder {
251 fn assert_ext_eq<I: Into<Self::Expr>>(
253 &mut self,
254 left: BinomialExtension<I>,
255 right: BinomialExtension<I>,
256 ) {
257 for (left, right) in left.0.into_iter().zip(right.0) {
258 self.assert_eq(left, right);
259 }
260 }
261
262 fn assert_is_base_element<I: Into<Self::Expr> + Clone>(
264 &mut self,
265 element: BinomialExtension<I>,
266 ) {
267 let base_slice = element.as_base_slice();
268 let degree = base_slice.len();
269 base_slice[1..degree].iter().for_each(|coeff| {
270 self.assert_zero(coeff.clone().into());
271 });
272 }
273
274 fn if_else_ext(
276 &mut self,
277 condition: impl Into<Self::Expr> + Clone,
278 a: BinomialExtension<impl Into<Self::Expr> + Clone>,
279 b: BinomialExtension<impl Into<Self::Expr> + Clone>,
280 ) -> BinomialExtension<Self::Expr> {
281 BinomialExtension(array::from_fn(|i| {
282 self.if_else(condition.clone(), a.0[i].clone(), b.0[i].clone())
283 }))
284 }
285}
286
287pub trait SepticExtensionAirBuilder: BaseAirBuilder {
289 fn assert_septic_ext_eq<I: Into<Self::Expr>>(
291 &mut self,
292 left: SepticExtension<I>,
293 right: SepticExtension<I>,
294 ) {
295 for (left, right) in left.0.into_iter().zip(right.0) {
296 self.assert_eq(left, right);
297 }
298 }
299}
300
301pub trait MachineAirBuilder:
304 BaseAirBuilder + ExtensionAirBuilder + SepticExtensionAirBuilder + AirBuilderWithPublicValues
305{
306 #[allow(clippy::type_complexity)]
310 fn extract_public_values(
311 &self,
312 ) -> super::PublicValues<
313 [Self::PublicVar; 4],
314 [Self::PublicVar; 3],
315 [Self::PublicVar; 4],
316 Self::PublicVar,
317 > {
318 use super::{PublicValues, SP1_PROOF_NUM_PV_ELTS};
319 use std::borrow::Borrow;
320
321 let public_values_slice: [Self::PublicVar; SP1_PROOF_NUM_PV_ELTS] =
322 core::array::from_fn(|i| self.public_values()[i]);
323 let public_values: &PublicValues<
324 [Self::PublicVar; 4],
325 [Self::PublicVar; 3],
326 [Self::PublicVar; 4],
327 Self::PublicVar,
328 > = public_values_slice.as_slice().borrow();
329 *public_values
330 }
331}
332
333pub trait SP1AirBuilder: MachineAirBuilder + ByteAirBuilder + InstructionAirBuilder {}
335
336impl<AB: AirBuilder + MessageBuilder<M>, M> MessageBuilder<M> for FilteredAirBuilder<'_, AB> {
337 fn send(&mut self, message: M, scope: InteractionScope) {
338 self.inner.send(message, scope);
339 }
340
341 fn receive(&mut self, message: M, scope: InteractionScope) {
342 self.inner.receive(message, scope);
343 }
344}
345
346impl<AB: AirBuilder + MessageBuilder<AirInteraction<AB::Expr>>> BaseAirBuilder for AB {}
347impl<AB: BaseAirBuilder> ByteAirBuilder for AB {}
348impl<AB: BaseAirBuilder> InstructionAirBuilder for AB {}
349
350impl<AB: BaseAirBuilder> ExtensionAirBuilder for AB {}
351impl<AB: BaseAirBuilder> SepticExtensionAirBuilder for AB {}
352impl<AB: BaseAirBuilder + AirBuilderWithPublicValues> MachineAirBuilder for AB {}
353impl<AB: BaseAirBuilder + AirBuilderWithPublicValues> SP1AirBuilder for AB {}
354
355impl<SC: StarkGenericConfig> EmptyMessageBuilder for ProverConstraintFolder<'_, SC> {}
356impl<SC: StarkGenericConfig> EmptyMessageBuilder for VerifierConstraintFolder<'_, SC> {}
357impl<
358 F: Field,
359 K: Field + From<F> + Add<F, Output = K> + Sub<F, Output = K> + Mul<F, Output = K>,
360 EF: Field + Mul<K, Output = EF>,
361 > EmptyMessageBuilder for ConstraintSumcheckFolder<'_, F, K, EF>
362{
363}
364impl<F: Field> EmptyMessageBuilder for SymbolicAirBuilder<F> {}
365
366#[cfg(debug_assertions)]
367#[cfg(not(doctest))]
368impl<F: Field> EmptyMessageBuilder for slop_uni_stark::DebugConstraintBuilder<'_, F> {}