Skip to main content

sp1_hypercube/air/
builder.rs

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/// The scope of an interaction.
20#[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 scope.
36    Global = 0,
37    /// Local scope.
38    Local,
39}
40
41/// A builder that can send and receive messages (or interactions) with other AIRs.
42pub trait MessageBuilder<M> {
43    /// Sends a message.
44    fn send(&mut self, message: M, scope: InteractionScope);
45
46    /// Receives a message.
47    fn receive(&mut self, message: M, scope: InteractionScope);
48}
49
50/// A message builder for which sending and receiving messages is a no-op.
51pub 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
59/// A trait which contains basic methods for building an AIR.
60pub trait BaseAirBuilder: AirBuilder + MessageBuilder<AirInteraction<Self::Expr>> {
61    /// Returns a sub-builder whose constraints are enforced only when `condition` is not one.
62    fn when_not<I: Into<Self::Expr>>(&mut self, condition: I) -> FilteredAirBuilder<'_, Self> {
63        self.when_ne(condition, Self::F::one())
64    }
65
66    /// Asserts that an iterator of expressions are all equal.
67    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    /// Asserts that an iterator of expressions are all zero.
78    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    /// Will return `a` if `condition` is 1, else `b`.  This assumes that `condition` is already
83    /// checked to be a boolean.
84    #[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    /// Index an array of expressions using an index bitmap.  This function assumes that the
95    /// `EIndex` type is a boolean and that `index_bitmap`'s entries sum to 1.
96    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
111/// A trait which contains methods for byte interactions in an AIR.
112pub trait ByteAirBuilder: BaseAirBuilder {
113    /// Sends a byte operation to be processed.
114    #[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    /// Receives a byte operation to be processed.
134    #[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
154/// A trait which contains methods related to RISC-V instruction interactions in an AIR.
155pub trait InstructionAirBuilder: BaseAirBuilder {
156    /// Sends the current CPU state.
157    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    /// Receives the current CPU state.
178    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    /// Sends an syscall operation to be processed (with "ECALL" opcode).
199    #[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    /// Receives a syscall operation to be processed.
224    #[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
249/// A builder that can operation on extension elements.
250pub trait ExtensionAirBuilder: BaseAirBuilder {
251    /// Asserts that the two field extensions are equal.
252    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    /// Checks if an extension element is a base element.
263    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    /// Performs an if else on extension elements.
275    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
287/// A builder that can operation on septic extension elements.
288pub trait SepticExtensionAirBuilder: BaseAirBuilder {
289    /// Asserts that the two field extensions are equal.
290    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
301/// A trait that contains the common helper methods for building `SP1 recursion` and SP1 machine
302/// AIRs.
303pub trait MachineAirBuilder:
304    BaseAirBuilder + ExtensionAirBuilder + SepticExtensionAirBuilder + AirBuilderWithPublicValues
305{
306    /// Extract public values from the air builder and convert them to the proper type.
307    /// This is commonly used throughout the codebase to access public values in AIR
308    /// implementations.
309    #[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
333/// A trait which contains all helper methods for building SP1 machine AIRs.
334pub 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> {}