added tests for all examples

gladiopeace · Oct 6, 2021 · dd86abc · dd86abc
1 parent 24f1bad
commit dd86abc
Show file tree

Hide file tree

Showing 20 changed files with 247 additions and 176 deletions.
diff --git a/air/Cargo.toml b/air/Cargo.toml
@@ -12,6 +12,7 @@ edition = "2018"
 
 [lib]
 bench = false
+doctest = false
 
 [features]
 default = ["std"]

diff --git a/assembly/Cargo.toml b/assembly/Cargo.toml
@@ -12,6 +12,7 @@ edition = "2018"
 
 [lib]
 bench = false
+doctest = false
 
 [features]
 default = ["std"]

diff --git a/core/Cargo.toml b/core/Cargo.toml
@@ -12,6 +12,7 @@ edition = "2018"
 
 [lib]
 bench = false
+doctest = false
 
 [features]
 default = ["std"]

diff --git a/core/src/program/inputs.rs b/core/src/program/inputs.rs
@@ -11,11 +11,7 @@ pub struct ProgramInputs {
 
 impl ProgramInputs {
     /// Returns `ProgramInputs` initialized with the provided public and secret inputs.
-    pub fn new(
-        public: &[BaseElement],
-        secret_a: &[BaseElement],
-        secret_b: &[BaseElement],
-    ) -> ProgramInputs {
+    pub fn new(public: &[u128], secret_a: &[u128], secret_b: &[u128]) -> ProgramInputs {
         assert!(
             public.len() <= MAX_PUBLIC_INPUTS,
             "expected no more than {} public inputs, but received {}",
@@ -26,8 +22,11 @@ impl ProgramInputs {
             "number of primary secret inputs cannot be smaller than the number of secondary secret inputs");
 
         ProgramInputs {
-            public: public.to_vec(),
-            secret: [secret_a.to_vec(), secret_b.to_vec()],
+            public: public.iter().map(|&v| BaseElement::new(v)).collect(),
+            secret: [
+                secret_a.iter().map(|&v| BaseElement::new(v)).collect(),
+                secret_b.iter().map(|&v| BaseElement::new(v)).collect(),
+            ],
         }
     }
 
@@ -41,18 +40,18 @@ impl ProgramInputs {
 
     /// Returns `ProgramInputs` initialized with the provided public inputs and secret
     /// input tapes set to empty vectors.
-    pub fn from_public(public: &[BaseElement]) -> ProgramInputs {
+    pub fn from_public(public: &[u128]) -> ProgramInputs {
         ProgramInputs {
-            public: public.to_vec(),
+            public: public.iter().map(|&v| BaseElement::new(v)).collect(),
             secret: [vec![], vec![]],
         }
     }
 
-    pub fn get_public_inputs(&self) -> &[BaseElement] {
+    pub fn public_inputs(&self) -> &[BaseElement] {
         &self.public
     }
 
-    pub fn get_secret_inputs(&self) -> &[Vec<BaseElement>; 2] {
+    pub fn secret_inputs(&self) -> &[Vec<BaseElement>; 2] {
         &self.secret
     }
 }
diff --git a/distaff/Cargo.toml b/distaff/Cargo.toml
@@ -12,6 +12,7 @@ license = "MIT"
 name = "distaff"
 path = "src/lib.rs"
 bench = false
+doctest = false
 
 [features]
 concurrent = ["prover/concurrent", "std"]

diff --git a/distaff/src/lib.rs b/distaff/src/lib.rs
@@ -71,7 +71,7 @@ pub fn execute(
 
     // generate STARK proof
     let inputs = inputs
-        .get_public_inputs()
+        .public_inputs()
         .iter()
         .map(|&v| v.as_int())
         .collect::<Vec<_>>();

diff --git a/distaff/src/tests.rs b/distaff/src/tests.rs
@@ -7,7 +7,7 @@ use air::ToElements;
 #[test]
 fn execute_span() {
     let program = assembly::compile("begin add push.5 mul push.7 end").unwrap();
-    let inputs = ProgramInputs::from_public(&[BaseElement::new(1), BaseElement::new(2)]);
+    let inputs = ProgramInputs::from_public(&[1, 2]);
 
     let trace = processor::execute(&program, &inputs);
     let trace_length = trace.length();
@@ -29,7 +29,7 @@ fn execute_span() {
 #[test]
 fn execute_block() {
     let program = assembly::compile("begin add block push.5 mul push.7 end end").unwrap();
-    let inputs = ProgramInputs::from_public(&[BaseElement::new(1), BaseElement::new(2)]);
+    let inputs = ProgramInputs::from_public(&[1, 2]);
 
     let trace = processor::execute(&program, &inputs);
     let trace_length = trace.length();
@@ -56,11 +56,7 @@ fn execute_if_else() {
             .unwrap();
 
     // execute true branch
-    let inputs = ProgramInputs::new(
-        &[BaseElement::new(5), BaseElement::new(3)],
-        &[BaseElement::new(1)],
-        &[],
-    );
+    let inputs = ProgramInputs::new(&[5, 3], &[1], &[]);
     let trace = processor::execute(&program, &inputs);
     let trace_length = trace.length();
     let trace_width = trace.width();
@@ -79,11 +75,7 @@ fn execute_if_else() {
     assert_eq!([24, 0, 0, 0, 0, 0, 0, 0].to_elements(), state.user_stack());
 
     // execute false branch
-    let inputs = ProgramInputs::new(
-        &[BaseElement::new(5), BaseElement::new(3)],
-        &[BaseElement::new(0)],
-        &[],
-    );
+    let inputs = ProgramInputs::new(&[5, 3], &[0], &[]);
     let trace = processor::execute(&program, &inputs);
     let trace_length = trace.length();
     let trace_width = trace.width();
@@ -107,11 +99,7 @@ fn execute_loop() {
     let program = assembly::compile("begin mul read while.true dup mul read end end").unwrap();
 
     // don't enter the loop
-    let inputs = ProgramInputs::new(
-        &[BaseElement::new(5), BaseElement::new(3)],
-        &[BaseElement::new(0)],
-        &[],
-    );
+    let inputs = ProgramInputs::new(&[5, 3], &[0], &[]);
     let trace = processor::execute(&program, &inputs);
 
     assert_eq!(64, trace.length());
@@ -128,11 +116,7 @@ fn execute_loop() {
     assert_eq!([15, 0, 0, 0, 0, 0, 0, 0].to_elements(), state.user_stack());
 
     // execute one iteration
-    let inputs = ProgramInputs::new(
-        &[BaseElement::new(5), BaseElement::new(3)],
-        &[BaseElement::new(1), BaseElement::new(0)],
-        &[],
-    );
+    let inputs = ProgramInputs::new(&[5, 3], &[1, 0], &[]);
     let trace = processor::execute(&program, &inputs);
 
     assert_eq!(128, trace.length());
@@ -149,18 +133,7 @@ fn execute_loop() {
     assert_eq!([225, 0, 0, 0, 0, 0, 0, 0].to_elements(), state.user_stack());
 
     // execute five iteration
-    let inputs = ProgramInputs::new(
-        &[BaseElement::new(5), BaseElement::new(3)],
-        &[
-            BaseElement::new(1),
-            BaseElement::new(1),
-            BaseElement::new(1),
-            BaseElement::new(1),
-            BaseElement::new(1),
-            BaseElement::new(0),
-        ],
-        &[],
-    );
+    let inputs = ProgramInputs::new(&[5, 3], &[1, 1, 1, 1, 1, 0], &[]);
     let trace = processor::execute(&program, &inputs);
 
     assert_eq!(256, trace.length());

diff --git a/examples/Cargo.toml b/examples/Cargo.toml
@@ -10,13 +10,13 @@ edition = "2018"
 
 [lib]
 bench = false
-doc = false
+doctest = false
 
 [[bin]]
 name = "distaff"
 path = "src/main.rs"
 bench = false
-doc = false
+doctest = false
 
 [dependencies]
 distaff = { path = "../distaff" }

diff --git a/examples/src/collatz.rs b/examples/src/collatz.rs
@@ -1,12 +1,16 @@
 use crate::Example;
 use distaff::{assembly, BaseElement, FieldElement, ProgramInputs, StarkField};
+use log::debug;
+
+// EXAMPLE BUILDER
+// ================================================================================================
 
 pub fn get_example(start_value: usize) -> Example {
     // convert starting value of the sequence into a field element
     let start_value = BaseElement::new(start_value as u128);
 
     // determine the expected result
-    let expected_result = compute_collatz_steps(start_value);
+    let expected_result = compute_collatz_steps(start_value).as_int();
 
     // construct the program which executes an unbounded loop to compute a Collatz sequence
     // which starts with the provided value; the output of the program is the number of steps
@@ -29,22 +33,17 @@ pub fn get_example(start_value: usize) -> Example {
     )
     .unwrap();
 
-    println!(
+    debug!(
         "Generated a program to compute Collatz sequence; expected result: {}",
         expected_result
     );
 
-    // put the starting value as the only secret input for tape A
-    let inputs = ProgramInputs::new(&[], &[start_value], &[]);
-
-    // a single element from the top of the stack will be the output
-    let num_outputs = 1;
-
     Example {
         program,
-        inputs,
+        inputs: ProgramInputs::new(&[], &[start_value.as_int()], &[]),
+        pub_inputs: vec![],
         expected_result: vec![expected_result],
-        num_outputs,
+        num_outputs: 1,
     }
 }
 
@@ -62,3 +61,18 @@ fn compute_collatz_steps(mut value: BaseElement) -> BaseElement {
 
     BaseElement::new(i)
 }
+
+// EXAMPLE TESTER
+// ================================================================================================
+
+#[test]
+fn test_collatz_example() {
+    let example = get_example(5);
+    super::test_example(example, false);
+}
+
+#[test]
+fn test_collatz_example_fail() {
+    let example = get_example(5);
+    super::test_example(example, true);
+}
diff --git a/examples/src/comparison.rs b/examples/src/comparison.rs
@@ -1,16 +1,14 @@
 use crate::Example;
-use distaff::{assembly, BaseElement, ProgramInputs, StarkField};
+use distaff::{assembly, ProgramInputs};
+use log::debug;
 
-pub fn get_example(value: usize) -> Example {
-    // convert value to a field element
-    let value = BaseElement::new(value as u128);
+// EXAMPLE BUILDER
+// ================================================================================================
 
+pub fn get_example(value: usize) -> Example {
     // determine the expected result
-    let expected_result: BaseElement = if value.as_int() < 9 {
-        value * BaseElement::new(9)
-    } else {
-        value + BaseElement::new(9)
-    };
+    let value = value as u128;
+    let expected_result = if value < 9 { value * 9 } else { value + 9 };
 
     // construct the program which checks if the value provided via secret inputs is
     // less than 9; if it is, the value is multiplied by 9, otherwise, 9 is added
@@ -33,24 +31,31 @@ pub fn get_example(value: usize) -> Example {
     )
     .unwrap();
 
-    println!(
+    debug!(
         "Generated a program to test comparisons; expected result: {}",
         expected_result
     );
 
-    // put the flag as the only secret input for tape A
-    let inputs = ProgramInputs::new(&[], &[value], &[]);
-
-    // a single element from the top of the stack will be the output
-    let num_outputs = 2;
-
     Example {
         program,
-        inputs,
-        expected_result: vec![
-            BaseElement::new(expected_result.as_int() & 1),
-            expected_result,
-        ],
-        num_outputs,
+        inputs: ProgramInputs::new(&[], &[value], &[]),
+        pub_inputs: vec![],
+        expected_result: vec![expected_result & 1, expected_result],
+        num_outputs: 2,
     }
 }
+
+// EXAMPLE TESTER
+// ================================================================================================
+
+#[test]
+fn test_comparison_example() {
+    let example = get_example(10);
+    super::test_example(example, false);
+}
+
+#[test]
+fn test_comparison_example_fail() {
+    let example = get_example(10);
+    super::test_example(example, true);
+}
diff --git a/examples/src/conditional.rs b/examples/src/conditional.rs
@@ -1,14 +1,18 @@
 use crate::Example;
-use distaff::{assembly, BaseElement, ProgramInputs, StarkField};
+use distaff::{assembly, ProgramInputs};
+use log::debug;
+
+// EXAMPLE BUILDER
+// ================================================================================================
 
 pub fn get_example(flag: usize) -> Example {
     // convert flag to a field element
-    let flag = BaseElement::new(flag as u128);
+    let flag = flag as u128;
 
     // determine the expected result
-    let expected_result = match flag.as_int() {
-        0 => BaseElement::new(15),
-        1 => BaseElement::new(8),
+    let expected_result = match flag {
+        0 => 15u128,
+        1 => 8u128,
         _ => panic!("flag must be a binary value"),
     };
 
@@ -29,21 +33,31 @@ pub fn get_example(flag: usize) -> Example {
     )
     .unwrap();
 
-    println!(
+    debug!(
         "Generated a program to test conditional execution; expected result: {}",
         expected_result
     );
 
-    // put the flag as the only secret input for tape A
-    let inputs = ProgramInputs::new(&[], &[flag], &[]);
-
-    // a single element from the top of the stack will be the output
-    let num_outputs = 1;
-
     Example {
         program,
-        inputs,
+        inputs: ProgramInputs::new(&[], &[flag], &[]),
+        pub_inputs: vec![],
         expected_result: vec![expected_result],
-        num_outputs,
+        num_outputs: 1,
     }
 }
+
+// EXAMPLE TESTER
+// ================================================================================================
+
+#[test]
+fn test_conditional_example() {
+    let example = get_example(1);
+    super::test_example(example, false);
+}
+
+#[test]
+fn test_conditional_example_fail() {
+    let example = get_example(1);
+    super::test_example(example, true);
+}