test: add list comprehension integration test

air-script/tests/list_comprehension/list_comprehension.air

@@ -0,0 +1,35 @@
+def ListComprehensionAir
+
+trace_columns:
+    main: [clk, fmp[2], ctx]
+    aux: [a, b[3], c[4], d[4]]
+
+public_inputs:
+    stack_inputs: [16]
+
+boundary_constraints:
+    enf c[2].first = 0
+
+integrity_constraints:
+    let x = [fmp for fmp in fmp]
+    enf clk = x[1]
+
+    let enumerate = [2^i * c for (i, c) in (0..4, c)]
+    enf a = clk * enumerate[3]
+
+    let diff_ident_iterables = [x' - y' for (x, y) in (c, d)]
+    enf a = clk * diff_ident_iterables[0]
+
+    let diff_slice_iterables = [x - y for (x, y) in (c[0..2], d[1..3])]
+    enf b[1] = clk * diff_slice_iterables[1]
+
+    let x0 = [w + x - y - z for (w, x, y, z) in (0..3, b, c[0..3], d[0..3])]
+    enf a = x0[0] + x0[1] + x0[2]
+
+    let x1 = sum([c * d for (c, d) in (c, d)])
+    let y1 = prod([c + d for (c, d) in (c, d)])    
+    enf c[3] = x1 + y1
+
+    let x2 = sum([c * d for (c, d) in (c, d)])
+    let y2 = [m + x2 for m in fmp]
+    enf a * clk = y2[0] + x2

air-script/tests/list_comprehension/list_comprehension.rs

@@ -0,0 +1,96 @@
+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 ListComprehensionAir {
+    context: AirContext<Felt>,
+    stack_inputs: [Felt; 16],
+}
+
+impl ListComprehensionAir {
+    pub fn last_step(&self) -> usize {
+        self.trace_length() - self.context().num_transition_exemptions()
+    }
+}
+
+impl Air for ListComprehensionAir {
+    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)];
+        let aux_degrees = vec![TransitionConstraintDegree::new(2), TransitionConstraintDegree::new(2), TransitionConstraintDegree::new(2), TransitionConstraintDegree::new(1), TransitionConstraintDegree::new(4), TransitionConstraintDegree::new(2)];
+        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(6, 0, E::from(0_u64)));
+        result
+    }
+
+    fn evaluate_transition<E: FieldElement<BaseField = Felt>>(&self, frame: &EvaluationFrame<E>, periodic_values: &[E], result: &mut [E]) {
+        let main_current = frame.current();
+        let main_next = frame.next();
+        result[0] = main_current[0] - main_current[2];
+    }
+
+    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 main_current = main_frame.current();
+        let main_next = main_frame.next();
+        let aux_current = aux_frame.current();
+        let aux_next = aux_frame.next();
+        result[0] = aux_current[0] - main_current[0] * E::from(2_u64).exp(E::PositiveInteger::from(3_u64)) * aux_current[7];
+        result[1] = aux_current[0] - main_current[0] * (aux_next[4] - aux_next[8]);
+        result[2] = aux_current[2] - main_current[0] * (aux_current[5] - aux_current[10]);
+        result[3] = aux_current[0] - (E::from(0_u64) + aux_current[1] - aux_current[4] - aux_current[8] + E::from(1_u64) + aux_current[2] - aux_current[5] - aux_current[9] + E::from(2_u64) + aux_current[3] - aux_current[6] - aux_current[10]);
+        result[4] = aux_current[7] - (aux_current[4] * aux_current[8] + aux_current[5] * aux_current[9] + aux_current[6] * aux_current[10] + aux_current[7] * aux_current[11] + (aux_current[4] + aux_current[8]) * (aux_current[5] + aux_current[9]) * (aux_current[6] + aux_current[10]) * (aux_current[7] + aux_current[11]));
+        result[5] = aux_current[0] * main_current[0] - (main_current[1] + aux_current[4] * aux_current[8] + aux_current[5] * aux_current[9] + aux_current[6] * aux_current[10] + aux_current[7] * aux_current[11] + aux_current[4] * aux_current[8] + aux_current[5] * aux_current[9] + aux_current[6] * aux_current[10] + aux_current[7] * aux_current[11]);
+    }
+}

air-script/tests/main.rs

@@ -123,3 +123,14 @@ fn random_values() {
     let expected = expect_file!["random_values/random_values.rs"];
     expected.assert_eq(&generated_air);
 }
+
+#[test]
+fn list_comprehension() {
+    // TODO: Improve this test to include more complicated expressions
+    let generated_air = Test::new("tests/list_comprehension/list_comprehension.air".to_string())
+        .transpile()
+        .unwrap();
+
+    let expected = expect_file!["list_comprehension/list_comprehension.rs"];
+    expected.assert_eq(&generated_air);
+}