From 7fee3bd69cb80058fa9cfd2fe6daa3e2823e639f Mon Sep 17 00:00:00 2001 From: tohrnii <100405913+tohrnii@users.noreply.github.com> Date: Fri, 24 Feb 2023 04:04:43 +0000 Subject: [PATCH] test: add list comprehension integration test --- .../list_comprehension/list_comprehension.air | 35 +++++++ .../list_comprehension/list_comprehension.rs | 96 +++++++++++++++++++ air-script/tests/main.rs | 11 +++ 3 files changed, 142 insertions(+) create mode 100644 air-script/tests/list_comprehension/list_comprehension.air create mode 100644 air-script/tests/list_comprehension/list_comprehension.rs diff --git a/air-script/tests/list_comprehension/list_comprehension.air b/air-script/tests/list_comprehension/list_comprehension.air new file mode 100644 index 00000000..1bff44e8 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/air-script/tests/list_comprehension/list_comprehension.rs b/air-script/tests/list_comprehension/list_comprehension.rs new file mode 100644 index 00000000..5c29d5b7 --- /dev/null +++ b/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(&self, target: &mut W) { + target.write(self.stack_inputs.as_slice()); + } +} + +pub struct ListComprehensionAir { + context: AirContext, + 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 { + &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![] + } + + fn get_assertions(&self) -> Vec> { + let mut result = Vec::new(); + result + } + + fn get_aux_assertions>(&self, aux_rand_elements: &AuxTraceRandElements) -> Vec> { + let mut result = Vec::new(); + result.push(Assertion::single(6, 0, E::from(0_u64))); + result + } + + fn evaluate_transition>(&self, frame: &EvaluationFrame, 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(&self, main_frame: &EvaluationFrame, aux_frame: &EvaluationFrame, _periodic_values: &[F], aux_rand_elements: &AuxTraceRandElements, result: &mut [E]) + where F: FieldElement, + E: FieldElement + ExtensionOf, + { + 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]); + } +} \ No newline at end of file diff --git a/air-script/tests/main.rs b/air-script/tests/main.rs index 529e5a63..c6175242 100644 --- a/air-script/tests/main.rs +++ b/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); +}