openvm_circuit_primitives/
sub_air.rs

1use openvm_stark_backend::p3_air::AirBuilder;
2
3/// Trait with associated types intended to allow reuse of constraint logic
4/// inside other AIRs.
5///
6/// A `SubAir` is **not** an [Air](openvm_stark_backend::p3_air::Air) itself.
7/// A `SubAir` is a struct that holds the means to generate a particular set of constraints,
8/// meant to be reusable within other AIRs.
9///
10/// The trait is designed to be maximally flexible, but typical implementations will separate
11/// the `AirContext` into two parts: `Io` and `AuxCols`. The `Io` part will consist of
12/// expressions (built using `AB::Expr`) that the `SubAir` does not own, while the `AuxCols`
13/// are any internal columns that the `SubAir` requires to generate its constraints. The
14/// `AuxCols` are columns that the `SubAir` fully owns and should be internally determined by
15/// the `SubAir` from the `Io` part. These `AuxCols` are typically just slices of `AB::Var`.
16///
17/// This trait only owns the constraints, but it is expected that the [TraceSubRowGenerator] trait
18/// or some analogous functionality is also implemented so that the trace generation of the `AuxCols`
19/// of each row can be done purely in terms of the `Io` part.
20pub trait SubAir<AB: AirBuilder> {
21    /// Type to define the context, typically in terms of `AB::Expr` that are needed
22    /// to define the SubAir's constraints.
23    type AirContext<'a>
24    where
25        Self: 'a,
26        AB: 'a,
27        AB::Var: 'a,
28        AB::Expr: 'a;
29
30    fn eval<'a>(&'a self, builder: &'a mut AB, ctx: Self::AirContext<'a>)
31    where
32        AB::Var: 'a,
33        AB::Expr: 'a;
34}
35
36/// This is a helper for generation of the trace on a subset of the columns in a single row
37/// of the trace matrix.
38// [jpw] This could be part of SubAir, but I want to keep SubAir to be constraints only
39pub trait TraceSubRowGenerator<F> {
40    /// The minimal amount of information needed to generate the sub-row of the trace matrix.
41    /// This type has a lifetime so other context, such as references to other chips, can be provided.
42    type TraceContext<'a>
43    where
44        Self: 'a;
45    /// The type for the columns to mutate. Often this can be `&'a mut Cols<F>` if `Cols` is on the stack.
46    /// For structs that use the heap, this should be a struct that contains mutable slices.
47    type ColsMut<'a>
48    where
49        Self: 'a;
50
51    fn generate_subrow<'a>(&'a self, ctx: Self::TraceContext<'a>, sub_row: Self::ColsMut<'a>);
52}