Skip to content

Commit

Permalink
Merge pull request #102 from 0xPolygonMiden/tohrnii-col-grouping-ir
Browse files Browse the repository at this point in the history
Add trace column grouping to IR
  • Loading branch information
tohrnii authored Jan 6, 2023
2 parents b9dd537 + 2fdc6f1 commit 9ffc6ca
Show file tree
Hide file tree
Showing 13 changed files with 296 additions and 238 deletions.
10 changes: 10 additions & 0 deletions air-script/tests/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
15 changes: 15 additions & 0 deletions air-script/tests/trace_col_groups/trace_col_groups.air
Original file line number Diff line number Diff line change
@@ -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
89 changes: 89 additions & 0 deletions air-script/tests/trace_col_groups/trace_col_groups.rs
Original file line number Diff line number Diff line change
@@ -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();
}
}
142 changes: 0 additions & 142 deletions ir/src/boundary_constraints.rs

This file was deleted.

28 changes: 18 additions & 10 deletions ir/src/boundary_stmts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
11 changes: 10 additions & 1 deletion ir/src/error.rs
Original file line number Diff line number Diff line change
@@ -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 {
Expand Down Expand Up @@ -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 {}",
Expand Down
9 changes: 9 additions & 0 deletions ir/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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
Expand Down
Loading

0 comments on commit 9ffc6ca

Please sign in to comment.