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        trap_code: impl Into<Self::Expr> + Clone,
206        arg1: [impl Into<Self::Expr>; 3],
207        arg2: [impl Into<Self::Expr>; 3],
208        multiplicity: impl Into<Self::Expr>,
209        scope: InteractionScope,
210    ) {
211        let values = once(clk_high.into())
212            .chain(once(clk_low.into()))
213            .chain(once(syscall_id.into()))
214            .chain(cfg!(feature = "mprotect").then(|| trap_code.into()))
215            .chain(arg1.map(Into::into))
216            .chain(arg2.map(Into::into))
217            .collect::<Vec<_>>();
218
219        self.send(
220            AirInteraction::new(values, multiplicity.into(), InteractionKind::Syscall),
221            scope,
222        );
223    }
224
225    /// Receives a syscall operation to be processed.
226    #[allow(clippy::too_many_arguments)]
227    fn receive_syscall(
228        &mut self,
229        clk_high: impl Into<Self::Expr> + Clone,
230        clk_low: impl Into<Self::Expr> + Clone,
231        syscall_id: impl Into<Self::Expr> + Clone,
232        trap_code: impl Into<Self::Expr> + Clone,
233        arg1: [Self::Expr; 3],
234        arg2: [Self::Expr; 3],
235        multiplicity: impl Into<Self::Expr>,
236        scope: InteractionScope,
237    ) {
238        let values = once(clk_high.into())
239            .chain(once(clk_low.into()))
240            .chain(once(syscall_id.into()))
241            .chain(cfg!(feature = "mprotect").then(|| trap_code.into()))
242            .chain(arg1)
243            .chain(arg2)
244            .collect::<Vec<_>>();
245
246        self.receive(
247            AirInteraction::new(values, multiplicity.into(), InteractionKind::Syscall),
248            scope,
249        );
250    }
251}
252
253/// A builder that can operation on extension elements.
254pub trait ExtensionAirBuilder: BaseAirBuilder {
255    /// Asserts that the two field extensions are equal.
256    fn assert_ext_eq<I: Into<Self::Expr>>(
257        &mut self,
258        left: BinomialExtension<I>,
259        right: BinomialExtension<I>,
260    ) {
261        for (left, right) in left.0.into_iter().zip(right.0) {
262            self.assert_eq(left, right);
263        }
264    }
265
266    /// Checks if an extension element is a base element.
267    fn assert_is_base_element<I: Into<Self::Expr> + Clone>(
268        &mut self,
269        element: BinomialExtension<I>,
270    ) {
271        let base_slice = element.as_base_slice();
272        let degree = base_slice.len();
273        base_slice[1..degree].iter().for_each(|coeff| {
274            self.assert_zero(coeff.clone().into());
275        });
276    }
277
278    /// Performs an if else on extension elements.
279    fn if_else_ext(
280        &mut self,
281        condition: impl Into<Self::Expr> + Clone,
282        a: BinomialExtension<impl Into<Self::Expr> + Clone>,
283        b: BinomialExtension<impl Into<Self::Expr> + Clone>,
284    ) -> BinomialExtension<Self::Expr> {
285        BinomialExtension(array::from_fn(|i| {
286            self.if_else(condition.clone(), a.0[i].clone(), b.0[i].clone())
287        }))
288    }
289}
290
291/// A builder that can operation on septic extension elements.
292pub trait SepticExtensionAirBuilder: BaseAirBuilder {
293    /// Asserts that the two field extensions are equal.
294    fn assert_septic_ext_eq<I: Into<Self::Expr>>(
295        &mut self,
296        left: SepticExtension<I>,
297        right: SepticExtension<I>,
298    ) {
299        for (left, right) in left.0.into_iter().zip(right.0) {
300            self.assert_eq(left, right);
301        }
302    }
303}
304
305/// A trait that contains the common helper methods for building `SP1 recursion` and SP1 machine
306/// AIRs.
307pub trait MachineAirBuilder:
308    BaseAirBuilder + ExtensionAirBuilder + SepticExtensionAirBuilder + AirBuilderWithPublicValues
309{
310    /// Extract public values from the air builder and convert them to the proper type.
311    /// This is commonly used throughout the codebase to access public values in AIR
312    /// implementations.
313    #[allow(clippy::type_complexity)]
314    fn extract_public_values(
315        &self,
316    ) -> super::PublicValues<
317        [Self::PublicVar; 4],
318        [Self::PublicVar; 3],
319        [Self::PublicVar; 4],
320        Self::PublicVar,
321    > {
322        use super::{PublicValues, SP1_PROOF_NUM_PV_ELTS};
323        use std::borrow::Borrow;
324
325        let public_values_slice: [Self::PublicVar; SP1_PROOF_NUM_PV_ELTS] =
326            core::array::from_fn(|i| self.public_values()[i]);
327        let public_values: &PublicValues<
328            [Self::PublicVar; 4],
329            [Self::PublicVar; 3],
330            [Self::PublicVar; 4],
331            Self::PublicVar,
332        > = public_values_slice.as_slice().borrow();
333        *public_values
334    }
335}
336
337/// A trait which contains all helper methods for building SP1 machine AIRs.
338pub trait SP1AirBuilder: MachineAirBuilder + ByteAirBuilder + InstructionAirBuilder {}
339
340impl<AB: AirBuilder + MessageBuilder<M>, M> MessageBuilder<M> for FilteredAirBuilder<'_, AB> {
341    fn send(&mut self, message: M, scope: InteractionScope) {
342        self.inner.send(message, scope);
343    }
344
345    fn receive(&mut self, message: M, scope: InteractionScope) {
346        self.inner.receive(message, scope);
347    }
348}
349
350impl<AB: AirBuilder + MessageBuilder<AirInteraction<AB::Expr>>> BaseAirBuilder for AB {}
351impl<AB: BaseAirBuilder> ByteAirBuilder for AB {}
352impl<AB: BaseAirBuilder> InstructionAirBuilder for AB {}
353
354impl<AB: BaseAirBuilder> ExtensionAirBuilder for AB {}
355impl<AB: BaseAirBuilder> SepticExtensionAirBuilder for AB {}
356impl<AB: BaseAirBuilder + AirBuilderWithPublicValues> MachineAirBuilder for AB {}
357impl<AB: BaseAirBuilder + AirBuilderWithPublicValues> SP1AirBuilder for AB {}
358
359impl<SC: StarkGenericConfig> EmptyMessageBuilder for ProverConstraintFolder<'_, SC> {}
360impl<SC: StarkGenericConfig> EmptyMessageBuilder for VerifierConstraintFolder<'_, SC> {}
361impl<
362        F: Field,
363        K: Field + From<F> + Add<F, Output = K> + Sub<F, Output = K> + Mul<F, Output = K>,
364        EF: Field + Mul<K, Output = EF>,
365    > EmptyMessageBuilder for ConstraintSumcheckFolder<'_, F, K, EF>
366{
367}
368impl<F: Field> EmptyMessageBuilder for SymbolicAirBuilder<F> {}
369
370#[cfg(debug_assertions)]
371#[cfg(not(doctest))]
372impl<F: Field> EmptyMessageBuilder for slop_uni_stark::DebugConstraintBuilder<'_, F> {}
sp1_hypercube/air/builder.rs

sp1_hypercube/air/
builder.rs
