sp1_stark/air/
builder.rs

1use std::{array, iter::once};
2
3use itertools::Itertools;
4use p3_air::{AirBuilder, AirBuilderWithPublicValues, FilteredAirBuilder, PermutationAirBuilder};
5use p3_field::{AbstractField, Field};
6use p3_uni_stark::{
7    ProverConstraintFolder, StarkGenericConfig, SymbolicAirBuilder, VerifierConstraintFolder,
8};
9use serde::{Deserialize, Serialize};
10use strum::{Display, EnumIter};
11
12use super::{interaction::AirInteraction, BinomialExtension};
13use crate::{
14    lookup::InteractionKind, septic_digest::SepticDigest, septic_extension::SepticExtension, Word,
15};
16
17/// The scope of an interaction.
18#[derive(
19    Debug,
20    Clone,
21    Copy,
22    PartialEq,
23    Eq,
24    Hash,
25    Display,
26    EnumIter,
27    PartialOrd,
28    Ord,
29    Serialize,
30    Deserialize,
31)]
32pub enum InteractionScope {
33    /// Global scope.
34    Global = 0,
35    /// Local scope.
36    Local,
37}
38
39/// A builder that can send and receive messages (or interactions) with other AIRs.
40pub trait MessageBuilder<M> {
41    /// Sends a message.
42    fn send(&mut self, message: M, scope: InteractionScope);
43
44    /// Receives a message.
45    fn receive(&mut self, message: M, scope: InteractionScope);
46}
47
48/// A message builder for which sending and receiving messages is a no-op.
49pub trait EmptyMessageBuilder: AirBuilder {}
50
51impl<AB: EmptyMessageBuilder, M> MessageBuilder<M> for AB {
52    fn send(&mut self, _message: M, _scope: InteractionScope) {}
53
54    fn receive(&mut self, _message: M, _scope: InteractionScope) {}
55}
56
57/// A trait which contains basic methods for building an AIR.
58pub trait BaseAirBuilder: AirBuilder + MessageBuilder<AirInteraction<Self::Expr>> {
59    /// Returns a sub-builder whose constraints are enforced only when `condition` is not one.
60    fn when_not<I: Into<Self::Expr>>(&mut self, condition: I) -> FilteredAirBuilder<'_, Self> {
61        self.when_ne(condition, Self::F::one())
62    }
63
64    /// Asserts that an iterator of expressions are all equal.
65    fn assert_all_eq<I1: Into<Self::Expr>, I2: Into<Self::Expr>>(
66        &mut self,
67        left: impl IntoIterator<Item = I1>,
68        right: impl IntoIterator<Item = I2>,
69    ) {
70        for (left, right) in left.into_iter().zip_eq(right) {
71            self.assert_eq(left, right);
72        }
73    }
74
75    /// Asserts that an iterator of expressions are all zero.
76    fn assert_all_zero<I: Into<Self::Expr>>(&mut self, iter: impl IntoIterator<Item = I>) {
77        iter.into_iter().for_each(|expr| self.assert_zero(expr));
78    }
79
80    /// Will return `a` if `condition` is 1, else `b`.  This assumes that `condition` is already
81    /// checked to be a boolean.
82    #[inline]
83    fn if_else(
84        &mut self,
85        condition: impl Into<Self::Expr> + Clone,
86        a: impl Into<Self::Expr> + Clone,
87        b: impl Into<Self::Expr> + Clone,
88    ) -> Self::Expr {
89        condition.clone().into() * a.into() + (Self::Expr::one() - condition.into()) * b.into()
90    }
91
92    /// Index an array of expressions using an index bitmap.  This function assumes that the
93    /// `EIndex` type is a boolean and that `index_bitmap`'s entries sum to 1.
94    fn index_array(
95        &mut self,
96        array: &[impl Into<Self::Expr> + Clone],
97        index_bitmap: &[impl Into<Self::Expr> + Clone],
98    ) -> Self::Expr {
99        let mut result = Self::Expr::zero();
100
101        for (value, i) in array.iter().zip_eq(index_bitmap) {
102            result = result.clone() + value.clone().into() * i.clone().into();
103        }
104
105        result
106    }
107}
108
109/// A trait which contains methods for byte interactions in an AIR.
110pub trait ByteAirBuilder: BaseAirBuilder {
111    /// Sends a byte operation to be processed.
112    #[allow(clippy::too_many_arguments)]
113    fn send_byte(
114        &mut self,
115        opcode: impl Into<Self::Expr>,
116        a: impl Into<Self::Expr>,
117        b: impl Into<Self::Expr>,
118        c: impl Into<Self::Expr>,
119        multiplicity: impl Into<Self::Expr>,
120    ) {
121        self.send_byte_pair(opcode, a, Self::Expr::zero(), b, c, multiplicity);
122    }
123
124    /// Sends a byte operation with two outputs to be processed.
125    #[allow(clippy::too_many_arguments)]
126    fn send_byte_pair(
127        &mut self,
128        opcode: impl Into<Self::Expr>,
129        a1: impl Into<Self::Expr>,
130        a2: impl Into<Self::Expr>,
131        b: impl Into<Self::Expr>,
132        c: impl Into<Self::Expr>,
133        multiplicity: impl Into<Self::Expr>,
134    ) {
135        self.send(
136            AirInteraction::new(
137                vec![opcode.into(), a1.into(), a2.into(), b.into(), c.into()],
138                multiplicity.into(),
139                InteractionKind::Byte,
140            ),
141            InteractionScope::Local,
142        );
143    }
144
145    /// Receives a byte operation to be processed.
146    #[allow(clippy::too_many_arguments)]
147    fn receive_byte(
148        &mut self,
149        opcode: impl Into<Self::Expr>,
150        a: impl Into<Self::Expr>,
151        b: impl Into<Self::Expr>,
152        c: impl Into<Self::Expr>,
153        multiplicity: impl Into<Self::Expr>,
154    ) {
155        self.receive_byte_pair(opcode, a, Self::Expr::zero(), b, c, multiplicity);
156    }
157
158    /// Receives a byte operation with two outputs to be processed.
159    #[allow(clippy::too_many_arguments)]
160    fn receive_byte_pair(
161        &mut self,
162        opcode: impl Into<Self::Expr>,
163        a1: impl Into<Self::Expr>,
164        a2: impl Into<Self::Expr>,
165        b: impl Into<Self::Expr>,
166        c: impl Into<Self::Expr>,
167        multiplicity: impl Into<Self::Expr>,
168    ) {
169        self.receive(
170            AirInteraction::new(
171                vec![opcode.into(), a1.into(), a2.into(), b.into(), c.into()],
172                multiplicity.into(),
173                InteractionKind::Byte,
174            ),
175            InteractionScope::Local,
176        );
177    }
178}
179
180/// A trait which contains methods related to RISC-V instruction interactions in an AIR.
181pub trait InstructionAirBuilder: BaseAirBuilder {
182    /// Sends a RISC-V instruction to be processed.
183    #[allow(clippy::too_many_arguments)]
184    fn send_instruction(
185        &mut self,
186        shard: impl Into<Self::Expr> + Clone,
187        clk: impl Into<Self::Expr> + Clone,
188        pc: impl Into<Self::Expr>,
189        next_pc: impl Into<Self::Expr>,
190        num_extra_cycles: impl Into<Self::Expr>,
191        opcode: impl Into<Self::Expr>,
192        a: Word<impl Into<Self::Expr>>,
193        b: Word<impl Into<Self::Expr>>,
194        c: Word<impl Into<Self::Expr>>,
195        op_a_0: impl Into<Self::Expr>,
196        op_a_immutable: impl Into<Self::Expr>,
197        is_memory: impl Into<Self::Expr>,
198        is_syscall: impl Into<Self::Expr>,
199        is_halt: impl Into<Self::Expr>,
200        multiplicity: impl Into<Self::Expr>,
201    ) {
202        let values = once(shard.into())
203            .chain(once(clk.into()))
204            .chain(once(pc.into()))
205            .chain(once(next_pc.into()))
206            .chain(once(num_extra_cycles.into()))
207            .chain(once(opcode.into()))
208            .chain(a.0.into_iter().map(Into::into))
209            .chain(b.0.into_iter().map(Into::into))
210            .chain(c.0.into_iter().map(Into::into))
211            .chain(once(op_a_0.into()))
212            .chain(once(op_a_immutable.into()))
213            .chain(once(is_memory.into()))
214            .chain(once(is_syscall.into()))
215            .chain(once(is_halt.into()))
216            .collect();
217
218        self.send(
219            AirInteraction::new(values, multiplicity.into(), InteractionKind::Instruction),
220            InteractionScope::Local,
221        );
222    }
223
224    /// Receives an ALU operation to be processed.
225    #[allow(clippy::too_many_arguments)]
226    fn receive_instruction(
227        &mut self,
228        shard: impl Into<Self::Expr> + Clone,
229        clk: impl Into<Self::Expr> + Clone,
230        pc: impl Into<Self::Expr>,
231        next_pc: impl Into<Self::Expr>,
232        num_extra_cycles: impl Into<Self::Expr>,
233        opcode: impl Into<Self::Expr>,
234        a: Word<impl Into<Self::Expr>>,
235        b: Word<impl Into<Self::Expr>>,
236        c: Word<impl Into<Self::Expr>>,
237        op_a_0: impl Into<Self::Expr>,
238        op_a_immutable: impl Into<Self::Expr>,
239        is_memory: impl Into<Self::Expr>,
240        is_syscall: impl Into<Self::Expr>,
241        is_halt: impl Into<Self::Expr>,
242        multiplicity: impl Into<Self::Expr>,
243    ) {
244        let values = once(shard.into())
245            .chain(once(clk.into()))
246            .chain(once(pc.into()))
247            .chain(once(next_pc.into()))
248            .chain(once(num_extra_cycles.into()))
249            .chain(once(opcode.into()))
250            .chain(a.0.into_iter().map(Into::into))
251            .chain(b.0.into_iter().map(Into::into))
252            .chain(c.0.into_iter().map(Into::into))
253            .chain(once(op_a_0.into()))
254            .chain(once(op_a_immutable.into()))
255            .chain(once(is_memory.into()))
256            .chain(once(is_syscall.into()))
257            .chain(once(is_halt.into()))
258            .collect();
259
260        self.receive(
261            AirInteraction::new(values, multiplicity.into(), InteractionKind::Instruction),
262            InteractionScope::Local,
263        );
264    }
265
266    /// Sends an syscall operation to be processed (with "ECALL" opcode).
267    #[allow(clippy::too_many_arguments)]
268    fn send_syscall(
269        &mut self,
270        shard: impl Into<Self::Expr> + Clone,
271        clk: impl Into<Self::Expr> + Clone,
272        syscall_id: impl Into<Self::Expr> + Clone,
273        arg1: impl Into<Self::Expr> + Clone,
274        arg2: impl Into<Self::Expr> + Clone,
275        multiplicity: impl Into<Self::Expr>,
276        scope: InteractionScope,
277    ) {
278        self.send(
279            AirInteraction::new(
280                vec![
281                    shard.clone().into(),
282                    clk.clone().into(),
283                    syscall_id.clone().into(),
284                    arg1.clone().into(),
285                    arg2.clone().into(),
286                ],
287                multiplicity.into(),
288                InteractionKind::Syscall,
289            ),
290            scope,
291        );
292    }
293
294    /// Receives a syscall operation to be processed.
295    #[allow(clippy::too_many_arguments)]
296    fn receive_syscall(
297        &mut self,
298        shard: impl Into<Self::Expr> + Clone,
299        clk: impl Into<Self::Expr> + Clone,
300        syscall_id: impl Into<Self::Expr> + Clone,
301        arg1: impl Into<Self::Expr> + Clone,
302        arg2: impl Into<Self::Expr> + Clone,
303        multiplicity: impl Into<Self::Expr>,
304        scope: InteractionScope,
305    ) {
306        self.receive(
307            AirInteraction::new(
308                vec![
309                    shard.clone().into(),
310                    clk.clone().into(),
311                    syscall_id.clone().into(),
312                    arg1.clone().into(),
313                    arg2.clone().into(),
314                ],
315                multiplicity.into(),
316                InteractionKind::Syscall,
317            ),
318            scope,
319        );
320    }
321}
322
323/// A builder that can operation on extension elements.
324pub trait ExtensionAirBuilder: BaseAirBuilder {
325    /// Asserts that the two field extensions are equal.
326    fn assert_ext_eq<I: Into<Self::Expr>>(
327        &mut self,
328        left: BinomialExtension<I>,
329        right: BinomialExtension<I>,
330    ) {
331        for (left, right) in left.0.into_iter().zip(right.0) {
332            self.assert_eq(left, right);
333        }
334    }
335
336    /// Checks if an extension element is a base element.
337    fn assert_is_base_element<I: Into<Self::Expr> + Clone>(
338        &mut self,
339        element: BinomialExtension<I>,
340    ) {
341        let base_slice = element.as_base_slice();
342        let degree = base_slice.len();
343        base_slice[1..degree].iter().for_each(|coeff| {
344            self.assert_zero(coeff.clone().into());
345        });
346    }
347
348    /// Performs an if else on extension elements.
349    fn if_else_ext(
350        &mut self,
351        condition: impl Into<Self::Expr> + Clone,
352        a: BinomialExtension<impl Into<Self::Expr> + Clone>,
353        b: BinomialExtension<impl Into<Self::Expr> + Clone>,
354    ) -> BinomialExtension<Self::Expr> {
355        BinomialExtension(array::from_fn(|i| {
356            self.if_else(condition.clone(), a.0[i].clone(), b.0[i].clone())
357        }))
358    }
359}
360
361/// A builder that can operation on septic extension elements.
362pub trait SepticExtensionAirBuilder: BaseAirBuilder {
363    /// Asserts that the two field extensions are equal.
364    fn assert_septic_ext_eq<I: Into<Self::Expr>>(
365        &mut self,
366        left: SepticExtension<I>,
367        right: SepticExtension<I>,
368    ) {
369        for (left, right) in left.0.into_iter().zip(right.0) {
370            self.assert_eq(left, right);
371        }
372    }
373}
374
375/// A builder that implements a permutation argument.
376pub trait MultiTableAirBuilder<'a>: PermutationAirBuilder {
377    /// The type of the local cumulative sum.
378    type LocalSum: Into<Self::ExprEF> + Copy;
379
380    /// The type of the global cumulative sum;
381    type GlobalSum: Into<Self::Expr> + Copy;
382
383    /// Returns the local cumulative sum of the permutation.
384    fn local_cumulative_sum(&self) -> &'a Self::LocalSum;
385
386    /// Returns the global cumulative sum of the permutation.
387    fn global_cumulative_sum(&self) -> &'a SepticDigest<Self::GlobalSum>;
388}
389
390/// A trait that contains the common helper methods for building `SP1 recursion` and SP1 machine
391/// AIRs.
392pub trait MachineAirBuilder:
393    BaseAirBuilder + ExtensionAirBuilder + SepticExtensionAirBuilder + AirBuilderWithPublicValues
394{
395}
396
397/// A trait which contains all helper methods for building SP1 machine AIRs.
398pub trait SP1AirBuilder: MachineAirBuilder + ByteAirBuilder + InstructionAirBuilder {}
399
400impl<AB: AirBuilder + MessageBuilder<M>, M> MessageBuilder<M> for FilteredAirBuilder<'_, AB> {
401    fn send(&mut self, message: M, scope: InteractionScope) {
402        self.inner.send(message, scope);
403    }
404
405    fn receive(&mut self, message: M, scope: InteractionScope) {
406        self.inner.receive(message, scope);
407    }
408}
409
410impl<AB: AirBuilder + MessageBuilder<AirInteraction<AB::Expr>>> BaseAirBuilder for AB {}
411impl<AB: BaseAirBuilder> ByteAirBuilder for AB {}
412impl<AB: BaseAirBuilder> InstructionAirBuilder for AB {}
413
414impl<AB: BaseAirBuilder> ExtensionAirBuilder for AB {}
415impl<AB: BaseAirBuilder> SepticExtensionAirBuilder for AB {}
416impl<AB: BaseAirBuilder + AirBuilderWithPublicValues> MachineAirBuilder for AB {}
417impl<AB: BaseAirBuilder + AirBuilderWithPublicValues> SP1AirBuilder for AB {}
418
419impl<SC: StarkGenericConfig> EmptyMessageBuilder for ProverConstraintFolder<'_, SC> {}
420impl<SC: StarkGenericConfig> EmptyMessageBuilder for VerifierConstraintFolder<'_, SC> {}
421impl<F: Field> EmptyMessageBuilder for SymbolicAirBuilder<F> {}
422
423#[cfg(debug_assertions)]
424#[cfg(not(doctest))]
425impl<F: Field> EmptyMessageBuilder for p3_uni_stark::DebugConstraintBuilder<'_, F> {}
sp1_stark/air/builder.rs

sp1_stark/air/
builder.rs
