Merge pull request #102 from 0xPolygonMiden/tohrnii-col-grouping-ir

Add trace column grouping to IR

air-script/tests/main.rs

@@ -85,3 +85,13 @@ fn variables() {
     let expected = expect_file!["variables/variables.rs"];
     expected.assert_eq(&generated_air);
 }
+
+#[test]
+fn trace_col_groups() {
+    let generated_air = Test::new("tests/trace_col_groups/trace_col_groups.air".to_string())
+        .transpile()
+        .unwrap();
+
+    let expected = expect_file!["trace_col_groups/trace_col_groups.rs"];
+    expected.assert_eq(&generated_air);
+}

air-script/tests/trace_col_groups/trace_col_groups.air

@@ -0,0 +1,15 @@
+def TraceColGroupAir
+
+trace_columns:
+    main: [clk, fmp[2], ctx]
+    aux: [a, b, c[3]]
+
+public_inputs:
+    stack_inputs: [16]
+
+transition_constraints:
+    enf fmp[1]' = fmp[1] + 1
+    enf fmp[0]' = fmp[0] - 1
+
+boundary_constraints:
+    enf c[2].first = 0

air-script/tests/trace_col_groups/trace_col_groups.rs

@@ -0,0 +1,89 @@
+use winter_air::{Air, AirContext, Assertion, AuxTraceRandElements, EvaluationFrame, ProofOptions as WinterProofOptions, TransitionConstraintDegree, TraceInfo};
+use winter_math::fields::f64::BaseElement as Felt;
+use winter_math::{ExtensionOf, FieldElement};
+use winter_utils::collections::Vec;
+use winter_utils::{ByteWriter, Serializable};
+
+pub struct PublicInputs {
+    stack_inputs: [Felt; 16],
+}
+
+impl PublicInputs {
+    pub fn new(stack_inputs: [Felt; 16]) -> Self {
+        Self { stack_inputs }
+    }
+}
+
+impl Serializable for PublicInputs {
+    fn write_into<W: ByteWriter>(&self, target: &mut W) {
+        target.write(self.stack_inputs.as_slice());
+    }
+}
+
+pub struct TraceColGroupAir {
+    context: AirContext<Felt>,
+    stack_inputs: [Felt; 16],
+}
+
+impl TraceColGroupAir {
+    pub fn last_step(&self) -> usize {
+        self.trace_length() - self.context().num_transition_exemptions()
+    }
+}
+
+impl Air for TraceColGroupAir {
+    type BaseField = Felt;
+    type PublicInputs = PublicInputs;
+
+    fn context(&self) -> &AirContext<Felt> {
+        &self.context
+    }
+
+    fn new(trace_info: TraceInfo, public_inputs: PublicInputs, options: WinterProofOptions) -> Self {
+        let main_degrees = vec![TransitionConstraintDegree::new(1), TransitionConstraintDegree::new(1)];
+        let aux_degrees = vec![];
+        let num_main_assertions = 0;
+        let num_aux_assertions = 1;
+
+        let context = AirContext::new_multi_segment(
+            trace_info,
+            main_degrees,
+            aux_degrees,
+            num_main_assertions,
+            num_aux_assertions,
+            options,
+        )
+        .set_num_transition_exemptions(2);
+        Self { context, stack_inputs: public_inputs.stack_inputs }
+    }
+
+    fn get_periodic_column_values(&self) -> Vec<Vec<Felt>> {
+        vec![]
+    }
+
+    fn get_assertions(&self) -> Vec<Assertion<Felt>> {
+        let mut result = Vec::new();
+        result
+    }
+
+    fn get_aux_assertions<E: FieldElement<BaseField = Felt>>(&self, aux_rand_elements: &AuxTraceRandElements<E>) -> Vec<Assertion<E>> {
+        let mut result = Vec::new();
+        result.push(Assertion::single(4, 0, E::from(0_u64)));
+        result
+    }
+
+    fn evaluate_transition<E: FieldElement<BaseField = Felt>>(&self, frame: &EvaluationFrame<E>, periodic_values: &[E], result: &mut [E]) {
+        let current = frame.current();
+        let next = frame.next();
+        result[0] = next[2] - (current[2] + E::from(1_u64));
+        result[1] = next[1] - (current[1] - (E::from(1_u64)));
+    }
+
+    fn evaluate_aux_transition<F, E>(&self, main_frame: &EvaluationFrame<F>, aux_frame: &EvaluationFrame<E>, _periodic_values: &[F], aux_rand_elements: &AuxTraceRandElements<E>, result: &mut [E])
+    where F: FieldElement<BaseField = Felt>,
+          E: FieldElement<BaseField = Felt> + ExtensionOf<F>,
+    {
+        let current = aux_frame.current();
+        let next = aux_frame.next();
+    }
+}

ir/src/boundary_stmts.rs

@@ -133,16 +133,24 @@ impl BoundaryStmts {
                 // add the constraint to the specified boundary for the specified trace
                 let col_type = symbol_table.get_type(constraint.column().name())?;
                 let result = match col_type {
-                    IdentifierType::TraceColumn(column) => match constraint.boundary() {
-                        ast::Boundary::First => self.boundary_constraints
-                            [column.trace_segment() as usize]
-                            .first_mut()
-                            .insert(column.col_idx(), expr),
-                        ast::Boundary::Last => self.boundary_constraints
-                            [column.trace_segment() as usize]
-                            .last_mut()
-                            .insert(column.col_idx(), expr),
-                    },
+                    IdentifierType::TraceColumns(columns) => {
+                        if constraint.column().idx() >= columns.size() {
+                            return Err(SemanticError::trace_access_out_of_bounds(
+                                constraint.column(),
+                                columns.size(),
+                            ));
+                        }
+                        match constraint.boundary() {
+                            ast::Boundary::First => self.boundary_constraints
+                                [columns.trace_segment() as usize]
+                                .first_mut()
+                                .insert(columns.offset() + constraint.column().idx(), expr),
+                            ast::Boundary::Last => self.boundary_constraints
+                                [columns.trace_segment() as usize]
+                                .last_mut()
+                                .insert(columns.offset() + constraint.column().idx(), expr),
+                        }
+                    }
                     _ => {
                         return Err(SemanticError::InvalidUsage(format!(
                             "Identifier {} was declared as a {}, not as a trace column",

ir/src/error.rs

@@ -1,5 +1,5 @@
 use crate::symbol_table::IdentifierType;
-use parser::ast::{MatrixAccess, VectorAccess};
+use parser::ast::{MatrixAccess, TraceAccess, VectorAccess};
 
 #[derive(Debug)]
 pub enum SemanticError {
@@ -49,6 +49,15 @@ impl SemanticError {
         ))
     }
 
+    pub(super) fn trace_access_out_of_bounds(access: &TraceAccess, size: usize) -> Self {
+        Self::IndexOutOfRange(format!(
+            "Out-of-range index {} in trace {} of length {}",
+            access.idx(),
+            access.name(),
+            size
+        ))
+    }
+
     pub(super) fn public_inputs_out_of_bounds(access: &VectorAccess, size: usize) -> Self {
         SemanticError::IndexOutOfRange(format!(
             "Out-of-range index {} in public input {} of length {}",

ir/src/lib.rs

@@ -14,6 +14,8 @@ pub mod transition_stmts;
 use transition_stmts::{AlgebraicGraph, TransitionStmts, VariableValue, MIN_CYCLE_LENGTH};
 pub use transition_stmts::{NodeIndex, TransitionConstraintDegree};
 
+mod trace_columns;
+
 mod error;
 use error::SemanticError;
 
@@ -23,13 +25,20 @@ use helpers::SourceValidator;
 #[cfg(test)]
 mod tests;
 
+// ==== ALIASES ===================================================================================
 pub type TraceSegment = u8;
 pub type Constants = Vec<Constant>;
 pub type PublicInputs = Vec<(String, usize)>;
 pub type PeriodicColumns = Vec<Vec<u64>>;
 pub type BoundaryConstraintsMap = BTreeMap<usize, BoundaryExpr>;
 pub type VariableRoots = BTreeMap<VariableValue, (TraceSegment, NodeIndex)>;
 
+// ==== CONSTANTS =================================================================================
+const CURRENT_ROW: usize = 0;
+const NEXT_ROW: usize = 1;
+
+// ==== AIR IR ====================================================================================
+
 /// Internal representation of an AIR.
 ///
 /// TODO: docs