Add support for trace col grouping

ir/src/boundary_constraints.rs

@@ -79,7 +79,7 @@ impl BoundaryConstraints {
         validate_expression(symbol_table, &expr)?;
 
         // add the constraint to the specified boundary for the specified trace
-        let col_type = symbol_table.get_type(constraint.column())?;
+        let col_type = symbol_table.get_type(constraint.column().name())?;
         let result = match col_type {
             IdentifierType::TraceColumn(column) => match column.trace_segment() {
                 0 => match constraint.boundary() {
@@ -95,7 +95,7 @@ impl BoundaryConstraints {
             _ => {
                 return Err(SemanticError::InvalidUsage(format!(
                     "Identifier {} was declared as a {}, not as a trace column",
-                    constraint.column(),
+                    constraint.column().name(),
                     col_type
                 )));
             }
@@ -106,7 +106,7 @@ impl BoundaryConstraints {
             return Err(SemanticError::TooManyConstraints(format!(
                 "A boundary constraint was already defined for {} '{}' at the {}",
                 col_type,
-                constraint.column(),
+                constraint.column().name(),
                 constraint.boundary()
             )));
         }

ir/src/lib.rs

@@ -68,7 +68,7 @@ impl AirIR {
                 ast::SourceSection::Constant(constant) => {
                     symbol_table.insert_constant(constant)?;
                 }
-                ast::SourceSection::TraceCols(columns) => {
+                ast::SourceSection::Trace(columns) => {
                     // process & validate the main trace columns
                     symbol_table.insert_trace_columns(0, &columns.main_cols)?;
                     // process & validate the auxiliary trace columns

ir/src/symbol_table.rs

@@ -4,7 +4,7 @@ use super::{
 };
 use parser::ast::{
     constants::{Constant, ConstantType},
-    Identifier, MatrixAccess, PeriodicColumn, PublicInput, VectorAccess,
+    Identifier, MatrixAccess, PeriodicColumn, PublicInput, TraceCols, VectorAccess,
 };
 use std::fmt::Display;
 
@@ -130,13 +130,15 @@ impl SymbolTable {
     pub(super) fn insert_trace_columns(
         &mut self,
         trace_segment: TraceSegment,
-        columns: &[Identifier],
+        trace: &[TraceCols],
     ) -> Result<(), SemanticError> {
         self.set_num_trace_segments(trace_segment);
 
-        for (idx, Identifier(name)) in columns.iter().enumerate() {
-            let trace_column = TraceColumn::new(trace_segment, idx);
-            self.insert_symbol(name, IdentifierType::TraceColumn(trace_column))?;
+        let mut col_idx = 0;
+        for trace_cols in trace {
+            let trace_column = TraceColumn::new(trace_segment, col_idx);
+            self.insert_symbol(trace_cols.name(), IdentifierType::TraceColumn(trace_column))?;
+            col_idx += trace_cols.size() as usize;
         }
 
         Ok(())

ir/src/transition_constraints/graph.rs

@@ -106,7 +106,9 @@ impl AlgebraicGraph {
             TransitionExpr::MatrixAccess(matrix_access) => {
                 self.insert_matrix_access(symbol_table, &matrix_access)
             }
-            TransitionExpr::Next(Identifier(ident)) => self.insert_next(symbol_table, &ident),
+            TransitionExpr::Next(trace_access) => {
+                self.insert_next(symbol_table, trace_access.name())
+            }
             TransitionExpr::Rand(index) => {
                 // constraint target for random values defaults to the second trace segment.
                 // TODO: make this more general, so random values from further trace segments can be

parser/src/ast/boundary_constraints.rs

@@ -1,4 +1,4 @@
-use super::{Identifier, MatrixAccess, VectorAccess};
+use super::{Identifier, MatrixAccess, TraceAccess, VectorAccess};
 use std::fmt::Display;
 
 // BOUNDARY CONSTRAINTS
@@ -13,22 +13,22 @@ pub enum BoundaryStmt {
 /// Stores the expression corresponding to the boundary constraint.
 #[derive(Debug, PartialEq)]
 pub struct BoundaryConstraint {
-    column: Identifier,
+    column: TraceAccess,
     boundary: Boundary,
     value: BoundaryExpr,
 }
 
 impl BoundaryConstraint {
-    pub fn new(column: Identifier, boundary: Boundary, value: BoundaryExpr) -> Self {
+    pub fn new(column: TraceAccess, boundary: Boundary, value: BoundaryExpr) -> Self {
         Self {
             column,
             boundary,
             value,
         }
     }
 
-    pub fn column(&self) -> &str {
-        &self.column.0
+    pub fn column(&self) -> &TraceAccess {
+        &self.column
     }
 
     pub fn boundary(&self) -> Boundary {

parser/src/ast/mod.rs

@@ -40,7 +40,7 @@ pub struct Source(pub Vec<SourceSection>);
 pub enum SourceSection {
     AirDef(Identifier),
     Constant(Constant),
-    TraceCols(TraceCols),
+    Trace(Trace),
     PublicInputs(Vec<PublicInput>),
     PeriodicColumns(Vec<PeriodicColumn>),
     BoundaryConstraints(Vec<BoundaryStmt>),
@@ -50,11 +50,56 @@ pub enum SourceSection {
 // TRACE
 // ================================================================================================
 
-/// [TraceCols] contains the main and auxiliary trace columns of the execution trace.
+/// [Trace] contains the main and auxiliary trace segments of the execution trace.
+#[derive(Debug, Eq, PartialEq)]
+pub struct Trace {
+    pub main_cols: Vec<TraceCols>,
+    pub aux_cols: Vec<TraceCols>,
+}
+
+/// [TraceCols] is used to represent a single or a group of columns in the execution trace. For
+/// single columns, the size is 1. For groups, the size is the number of columns in the group.
 #[derive(Debug, Eq, PartialEq)]
 pub struct TraceCols {
-    pub main_cols: Vec<Identifier>,
-    pub aux_cols: Vec<Identifier>,
+    name: Identifier,
+    size: u64,
+}
+
+impl TraceCols {
+    pub fn new(name: Identifier, size: u64) -> Self {
+        Self { name, size }
+    }
+
+    pub fn name(&self) -> &str {
+        self.name.name()
+    }
+
+    pub fn size(&self) -> u64 {
+        self.size
+    }
+}
+
+/// [TraceAccess] is used to indicate a column in the trace by specifying its index within a set of
+/// [TraceCols] with the given identifier. If the identifier refers to a single column ([TraceCols]
+/// with size 1), then the index is always zero.
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct TraceAccess {
+    name: Identifier,
+    idx: usize,
+}
+
+impl TraceAccess {
+    pub fn new(name: Identifier, idx: usize) -> Self {
+        Self { name, idx }
+    }
+
+    pub fn name(&self) -> &str {
+        self.name.name()
+    }
+
+    pub fn idx(&self) -> usize {
+        self.idx
+    }
 }
 
 // SHARED ATOMIC TYPES

parser/src/ast/transition_constraints.rs

@@ -1,4 +1,4 @@
-use super::{Identifier, MatrixAccess, VectorAccess};
+use super::{Identifier, MatrixAccess, TraceAccess, VectorAccess};
 
 // TRANSITION CONSTRAINTS
 // ================================================================================================
@@ -40,7 +40,7 @@ pub enum TransitionExpr {
     /// Represents an element inside a constant or variable matrix. [MatrixAccess] contains the
     /// name of the matrix and indices of the element to access.
     MatrixAccess(MatrixAccess),
-    Next(Identifier),
+    Next(TraceAccess),
     /// Represents a random value provided by the verifier. The inner value is the index of this
     /// random value in the array of all random values.
     Rand(usize),

parser/src/parser/grammar.lalrpop

@@ -5,8 +5,8 @@ use crate::{
         transition_constraints::{TransitionConstraint, TransitionExpr, TransitionStmt,
         TransitionVariable, TransitionVariableType},
         constants::{Constant, ConstantType},
-        Identifier, Source, SourceSection, TraceCols, PublicInput, PeriodicColumn, VectorAccess,
-        MatrixAccess
+        Identifier, Source, SourceSection, Trace, TraceAccess, TraceCols, PublicInput,
+        PeriodicColumn, VectorAccess, MatrixAccess
     }, error::{Error, ParseError::{InvalidInt, InvalidTraceCols, MissingMainTraceCols,
     InvalidConst, MissingBoundaryConstraint, MissingTransitionConstraint}}, lexer::Token
 };
@@ -25,7 +25,7 @@ pub Source: Source = {
 SourceSection: SourceSection = {
     AirDef => SourceSection::AirDef(<>),
     Constant => SourceSection::Constant(<>),
-    TraceCols => SourceSection::TraceCols(<>),
+    Trace => SourceSection::Trace(<>),
     PublicInputs => SourceSection::PublicInputs(<>),
     PeriodicColumns => SourceSection::PeriodicColumns(<>),
     BoundaryConstraints => SourceSection::BoundaryConstraints(<>),
@@ -42,11 +42,11 @@ AirDef: Identifier = {
 // TRACE COLUMNS
 // ================================================================================================
 
-TraceCols: TraceCols = {
+Trace: Trace = {
     "trace_columns" ":" <main_cols: MainCols?> <aux_cols: AuxCols?> =>? match (main_cols, aux_cols)
     {
-        (Some(main_cols), Some(aux_cols)) => Ok(TraceCols { main_cols, aux_cols }),
-        (Some(main_cols), None) => Ok(TraceCols { main_cols, aux_cols: vec![] }),
+        (Some(main_cols), Some(aux_cols)) => Ok(Trace { main_cols, aux_cols }),
+        (Some(main_cols), None) => Ok(Trace { main_cols, aux_cols: vec![] }),
         (None, Some(_aux_cols)) => Err(ParseError::User {
             error: Error::ParseError(
                 MissingMainTraceCols("Declaration of main trace columns is required".to_string())
@@ -58,12 +58,17 @@ TraceCols: TraceCols = {
     } 
 }
 
-MainCols: Vec<Identifier> = {
-    "main" ":" "[" <main_cols: CommaElems<Identifier>> "]" => main_cols,
+MainCols: Vec<TraceCols> = {
+    "main" ":" "[" <main_cols: CommaElems<TraceCols>> "]" => main_cols,
+}
+
+AuxCols: Vec<TraceCols> = {
+    "aux" ":" "[" <aux_cols: CommaElems<TraceCols>> "]" => aux_cols,
 }
 
-AuxCols: Vec<Identifier> = {
-    "aux" ":" "[" <aux_cols: CommaElems<Identifier>> "]" => aux_cols,
+TraceCols: TraceCols = {
+    <name: Identifier> => TraceCols::new(name, 1),
+    <name: Identifier> <size: Size> => TraceCols::new(name, size)
 }
 
 // CONSTANTS
@@ -139,7 +144,7 @@ BoundaryConstraints: Vec<BoundaryStmt> = {
 }
 
 BoundaryStmt: BoundaryStmt = {
-    "enf" <column: Identifier> "." <boundary: Boundary> "=" <value: BoundaryExpr> =>
+    "enf" <column: TraceAccess> "." <boundary: Boundary> "=" <value: BoundaryExpr> =>
         BoundaryStmt::Constraint(BoundaryConstraint::new(column, boundary, value)),
     "let" <name: Identifier> "=" <boundary_variable_type: BoundaryVariableType> =>
         BoundaryStmt::Variable(BoundaryVariable::new(name, boundary_variable_type)),
@@ -248,7 +253,7 @@ TransitionAtom: TransitionExpr = {
     <ident: Identifier> => TransitionExpr::Elem(ident),
     <vector_access: VectorAccess> => TransitionExpr::VectorAccess(vector_access),
     <matrix_access: MatrixAccess> => TransitionExpr::MatrixAccess(matrix_access),
-    <ident: Identifier> "'" => TransitionExpr::Next(ident)
+    <trace_access: TraceAccess> "'" => TransitionExpr::Next(trace_access)
 }
 
 // ATOMS
@@ -287,6 +292,11 @@ MatrixAccess: MatrixAccess = {
         MatrixAccess::new(ident, row, col)
 }
 
+TraceAccess: TraceAccess = {
+    <ident: Identifier> => TraceAccess::new(ident, 0),
+    <ident: Identifier> <idx: Index> => TraceAccess::new(ident, idx)
+}
+
 Identifier: Identifier = {
     <n:identifier> => Identifier(n.to_string())
 }

parser/src/parser/tests/arithmetic_ops.rs

@@ -2,6 +2,7 @@ use super::{
     build_parse_test, Identifier, Source, SourceSection::*, TransitionConstraint,
     TransitionExpr::*, TransitionStmt::*,
 };
+use crate::ast::TraceAccess;
 
 // EXPRESSIONS
 // ================================================================================================
@@ -15,7 +16,7 @@ fn single_addition() {
     let expected = Source(vec![TransitionConstraints(vec![Constraint(
         TransitionConstraint::new(
             Add(
-                Box::new(Next(Identifier("clk".to_string()))),
+                Box::new(Next(TraceAccess::new(Identifier("clk".to_string()), 0))),
                 Box::new(Elem(Identifier("clk".to_string()))),
             ),
             Const(0),
@@ -34,7 +35,7 @@ fn multi_addition() {
         TransitionConstraint::new(
             Add(
                 Box::new(Add(
-                    Box::new(Next(Identifier("clk".to_string()))),
+                    Box::new(Next(TraceAccess::new(Identifier("clk".to_string()), 0))),
                     Box::new(Elem(Identifier("clk".to_string()))),
                 )),
                 Box::new(Const(2)),
@@ -54,7 +55,7 @@ fn single_subtraction() {
     let expected = Source(vec![TransitionConstraints(vec![Constraint(
         TransitionConstraint::new(
             Sub(
-                Box::new(Next(Identifier("clk".to_string()))),
+                Box::new(Next(TraceAccess::new(Identifier("clk".to_string()), 0))),
                 Box::new(Elem(Identifier("clk".to_string()))),
             ),
             Const(0),
@@ -73,7 +74,7 @@ fn multi_subtraction() {
         TransitionConstraint::new(
             Sub(
                 Box::new(Sub(
-                    Box::new(Next(Identifier("clk".to_string()))),
+                    Box::new(Next(TraceAccess::new(Identifier("clk".to_string()), 0))),
                     Box::new(Elem(Identifier("clk".to_string()))),
                 )),
                 Box::new(Const(1)),
@@ -93,7 +94,7 @@ fn single_multiplication() {
     let expected = Source(vec![TransitionConstraints(vec![Constraint(
         TransitionConstraint::new(
             Mul(
-                Box::new(Next(Identifier("clk".to_string()))),
+                Box::new(Next(TraceAccess::new(Identifier("clk".to_string()), 0))),
                 Box::new(Elem(Identifier("clk".to_string()))),
             ),
             Const(0),
@@ -112,7 +113,7 @@ fn multi_multiplication() {
         TransitionConstraint::new(
             Mul(
                 Box::new(Mul(
-                    Box::new(Next(Identifier("clk".to_string()))),
+                    Box::new(Next(TraceAccess::new(Identifier("clk".to_string()), 0))),
                     Box::new(Elem(Identifier("clk".to_string()))),
                 )),
                 Box::new(Const(2)),
@@ -145,7 +146,7 @@ fn ops_with_parens() {
         TransitionConstraint::new(
             Mul(
                 Box::new(Add(
-                    Box::new(Next(Identifier("clk".to_string()))),
+                    Box::new(Next(TraceAccess::new(Identifier("clk".to_string()), 0))),
                     Box::new(Elem(Identifier("clk".to_string()))),
                 )),
                 Box::new(Const(2)),
@@ -164,7 +165,10 @@ fn single_exponentiation() {
         enf clk'^2 = 1";
     let expected = Source(vec![TransitionConstraints(vec![Constraint(
         TransitionConstraint::new(
-            Exp(Box::new(Next(Identifier("clk".to_string()))), 2),
+            Exp(
+                Box::new(Next(TraceAccess::new(Identifier("clk".to_string()), 0))),
+                2,
+            ),
             Const(1),
         ),
     )])]);
@@ -209,7 +213,7 @@ fn multi_arithmetic_ops_same_precedence() {
             Add(
                 Box::new(Sub(
                     Box::new(Sub(
-                        Box::new(Next(Identifier("clk".to_string()))),
+                        Box::new(Next(TraceAccess::new(Identifier("clk".to_string()), 0))),
                         Box::new(Elem(Identifier("clk".to_string()))),
                     )),
                     Box::new(Const(2)),
@@ -237,7 +241,10 @@ fn multi_arithmetic_ops_different_precedence() {
         TransitionConstraint::new(
             Sub(
                 Box::new(Sub(
-                    Box::new(Exp(Box::new(Next(Identifier("clk".to_string()))), 2)),
+                    Box::new(Exp(
+                        Box::new(Next(TraceAccess::new(Identifier("clk".to_string()), 0))),
+                        2,
+                    )),
                     Box::new(Mul(
                         Box::new(Elem(Identifier("clk".to_string()))),
                         Box::new(Const(2)),
@@ -266,7 +273,7 @@ fn multi_arithmetic_ops_different_precedence_w_parens() {
     let expected = Source(vec![TransitionConstraints(vec![Constraint(
         TransitionConstraint::new(
             Sub(
-                Box::new(Next(Identifier("clk".to_string()))),
+                Box::new(Next(TraceAccess::new(Identifier("clk".to_string()), 0))),
                 Box::new(Mul(
                     Box::new(Exp(Box::new(Elem(Identifier("clk".to_string()))), 2)),
                     Box::new(Sub(Box::new(Const(2)), Box::new(Const(1)))),