Adding random values to IR

air-script-core/src/expression.rs

@@ -16,9 +16,9 @@ pub enum Expression {
     MatrixAccess(MatrixAccess),
     IndexedTraceAccess(IndexedTraceAccess),
     NamedTraceAccess(NamedTraceAccess),
-    /// 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),
+    /// Represents a random value provided by the verifier. The first inner value is the name of
+    /// the random values array and the second is the index of this random value in that array
+    Rand(Identifier, usize),
     Add(Box<Expression>, Box<Expression>),
     Sub(Box<Expression>, Box<Expression>),
     Mul(Box<Expression>, Box<Expression>),

air-script/tests/aux_trace/aux_trace.air

@@ -7,6 +7,9 @@ trace_columns:
 public_inputs:
     stack_inputs: [16]
 
+random_values:
+    rand: [2]
+
 boundary_constraints:
     enf a.first = 1
     enf b.first = 1

air-script/tests/main.rs

@@ -109,3 +109,20 @@ fn indexed_trace_access() {
     let expected = expect_file!["indexed_trace_access/indexed_trace_access.rs"];
     expected.assert_eq(&generated_air);
 }
+
+#[test]
+fn random_values() {
+    let generated_air = Test::new("tests/random_values/random_values_simple.air".to_string())
+        .transpile()
+        .unwrap();
+
+    let expected = expect_file!["random_values/random_values.rs"];
+    expected.assert_eq(&generated_air);
+
+    let generated_air = Test::new("tests/random_values/random_values_bindings.air".to_string())
+        .transpile()
+        .unwrap();
+
+    let expected = expect_file!["random_values/random_values.rs"];
+    expected.assert_eq(&generated_air);
+}

air-script/tests/random_values/random_values.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 RandomValuesAir {
+    context: AirContext<Felt>,
+    stack_inputs: [Felt; 16],
+}
+
+impl RandomValuesAir {
+    pub fn last_step(&self) -> usize {
+        self.trace_length() - self.context().num_transition_exemptions()
+    }
+}
+
+impl Air for RandomValuesAir {
+    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![];
+        let aux_degrees = vec![TransitionConstraintDegree::new(1)];
+        let num_main_assertions = 0;
+        let num_aux_assertions = 2;
+
+        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(0, 0, aux_rand_elements.get_segment_elements(0)[5] + aux_rand_elements.get_segment_elements(0)[3] + aux_rand_elements.get_segment_elements(0)[15]));
+        result.push(Assertion::single(0, self.last_step(), aux_rand_elements.get_segment_elements(0)[0] + aux_rand_elements.get_segment_elements(0)[15] + aux_rand_elements.get_segment_elements(0)[11]));
+        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();
+    }
+
+    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();
+        result[0] = next[0] - (aux_rand_elements.get_segment_elements(0)[15] - (aux_rand_elements.get_segment_elements(0)[0]) + aux_rand_elements.get_segment_elements(0)[3]);
+    }
+}

air-script/tests/random_values/random_values_bindings.air

@@ -0,0 +1,18 @@
+def RandomValuesAir
+
+trace_columns:
+    main: [a, b]
+    aux: [c, d]
+
+public_inputs:
+    stack_inputs: [16]
+
+random_values:
+    alphas: [x, y[14], z]
+
+boundary_constraints:
+    enf c.first = $alphas[5] + y[2] + z
+    enf c.last = x + $alphas[15] + y[10]
+
+integrity_constraints:
+    enf c' = z - $alphas[0] + y[2]

air-script/tests/random_values/random_values_simple.air

@@ -0,0 +1,18 @@
+def RandomValuesAir
+
+trace_columns:
+    main: [a, b]
+    aux: [c, d]
+
+public_inputs:
+    stack_inputs: [16]
+
+random_values:
+    rand: [16]
+
+boundary_constraints:
+    enf c.first = $rand[5] + $rand[3] + $rand[15]
+    enf c.last = $rand[0] + $rand[15] + $rand[11]
+
+integrity_constraints:
+    enf c' = $rand[15] - $rand[0] + $rand[3]

air-script/tests/variables/variables.air

@@ -13,6 +13,9 @@ public_inputs:
 periodic_columns:
     k0: [1, 1, 1, 1, 1, 1, 1, 0]
 
+random_values:
+    rand: [1]
+
 boundary_constraints:
     enf a.first = 0
     enf a.last = 1
@@ -36,4 +39,4 @@ integrity_constraints:
     enf s * (c - a * b) = o[0][0] - o[0][1] - o[1][0]
 
     # the auxiliary column contains the product of values of c offset by a random value.
-    enf p' = p * (c + $rand[0])
+    enf p' = p * (c + $rand[0])

ir/src/constraints/graph.rs

@@ -107,14 +107,19 @@ impl AlgebraicGraph {
                 variable_roots,
                 default_domain,
             ),
-            Expression::Rand(index) => {
+            Expression::Rand(name, index) => {
                 // The constraint target for random values defaults to the second (auxiliary) trace
                 // segment.
                 // TODO: make this more general, so random values from further trace segments can be
                 // used. This requires having a way to describe different sets of randomness in
                 // the AirScript syntax.
-                let node_index = self.insert_op(Operation::RandomValue(*index));
-                Ok(ExprDetails::new(node_index, AUX_SEGMENT, default_domain))
+                self.insert_random_access(
+                    symbol_table,
+                    name.name(),
+                    *index,
+                    AUX_SEGMENT,
+                    default_domain,
+                )
             }
             Expression::IndexedTraceAccess(column_access) => {
                 self.insert_indexed_trace_access(symbol_table, column_access)
@@ -273,6 +278,27 @@ impl AlgebraicGraph {
         Ok(ExprDetails::new(node_index, DEFAULT_SEGMENT, domain))
     }
 
+    /// Inserts random value with specified index into the graph and returns the resulting
+    /// expression details.
+    fn insert_random_value(
+        &mut self,
+        symbol_table: &SymbolTable,
+        index: usize,
+        trace_segment: u8,
+        domain: ConstraintDomain,
+    ) -> Result<ExprDetails, SemanticError> {
+        if index >= symbol_table.num_random_values() as usize {
+            return Err(SemanticError::IndexOutOfRange(format!(
+                "Random value index {} is greater than or equal to the total number of random values ({}).", 
+                index,
+                symbol_table.num_random_values()
+            )));
+        }
+
+        let node_index = self.insert_op(Operation::RandomValue(index));
+        Ok(ExprDetails::new(node_index, trace_segment, domain))
+    }
+
     /// Looks up the specified variable value in the variable roots and returns the expression
     /// details if it is found. Otherwise, inserts the variable expression into the graph, adds it
     /// to the variable roots, and returns the resulting expression details.
@@ -330,8 +356,8 @@ impl AlgebraicGraph {
         Ok(ExprDetails::new(node_index, trace_segment, domain))
     }
 
-    /// Adds a trace column, periodic column, named constant or a variable to the graph and returns
-    /// the [ExprDetails] of the inserted expression.
+    /// Adds a trace column, periodic column, random value, named constant or a variable to the
+    /// graph and returns the [ExprDetails] of the inserted expression.
     ///
     /// # Errors
     /// Returns an error if the identifier is not present in the symbol table or is not a supported
@@ -348,6 +374,9 @@ impl AlgebraicGraph {
             IdentifierType::Constant(ConstantType::Scalar(_)) => {
                 self.insert_constant(ConstantValue::Scalar(ident.to_string()), domain)
             }
+            IdentifierType::RandomValuesBinding(offset, _) => {
+                self.insert_random_value(symbol_table, *offset, AUX_SEGMENT, domain)
+            }
             IdentifierType::IntegrityVariable(integrity_variable) => {
                 if let VariableType::Scalar(variable_expr) = integrity_variable.value() {
                     self.insert_variable(
@@ -426,6 +455,12 @@ impl AlgebraicGraph {
             IdentifierType::PublicInput(_) => {
                 unimplemented!("TODO: add support for public inputs.")
             }
+            IdentifierType::RandomValuesBinding(offset, _) => self.insert_random_value(
+                symbol_table,
+                *offset + vector_access.idx(),
+                AUX_SEGMENT,
+                domain,
+            ),
             _ => Err(SemanticError::invalid_vector_access(
                 vector_access,
                 symbol_type,
@@ -472,6 +507,35 @@ impl AlgebraicGraph {
             )),
         }
     }
+
+    /// Inserts random value by index access into the graph and returns the resulting
+    /// expression details.
+    ///
+    /// # Errors
+    /// Returns an error if the identifier is not present in the symbol table or is not a supported
+    /// type.
+    ///
+    /// # Example
+    /// This function inserts values like `$alphas[3]`, having `name` as `alphas` and `index` as
+    /// `3`
+    fn insert_random_access(
+        &mut self,
+        symbol_table: &SymbolTable,
+        name: &str,
+        index: usize,
+        trace_segment: u8,
+        domain: ConstraintDomain,
+    ) -> Result<ExprDetails, SemanticError> {
+        let elem_type = symbol_table.get_type(name)?;
+        match elem_type {
+            IdentifierType::RandomValuesBinding(_, _) => {
+                self.insert_random_value(symbol_table, index, trace_segment, domain)
+            }
+            _ => Err(SemanticError::InvalidUsage(format!(
+                "Identifier {name} was declared as a {elem_type} not as a random values"
+            ))),
+        }
+    }
 }
 
 /// Reference to a node in a graph by its index in the nodes vector of the graph struct.

ir/src/lib.rs

@@ -89,6 +89,9 @@ impl AirIR {
                     // process & validate the periodic columns
                     symbol_table.insert_periodic_columns(columns)?;
                 }
+                ast::SourceSection::RandomValues(values) => {
+                    symbol_table.insert_random_values(values)?;
+                }
                 _ => {}
             }
         }