add test type to pest

compiler/src/ast.rs

@@ -1225,6 +1225,16 @@ pub struct Import<'ast> {
     pub span: Span<'ast>,
 }
 
+// Tests
+
+#[derive(Clone, Debug, FromPest, PartialEq)]
+#[pest_ast(rule(Rule::test))]
+pub struct Test<'ast> {
+    pub function: Function<'ast>,
+    #[pest_ast(outer())]
+    pub span: Span<'ast>,
+}
+
 // File
 
 #[derive(Clone, Debug, FromPest, PartialEq)]
@@ -1233,6 +1243,7 @@ pub struct File<'ast> {
     pub imports: Vec<Import<'ast>>,
     pub circuits: Vec<Circuit<'ast>>,
     pub functions: Vec<Function<'ast>>,
+    pub tests: Vec<Test<'ast>>,
     pub eoi: EOI,
     #[pest_ast(outer())]
     pub span: Span<'ast>,

compiler/src/leo.pest

@@ -226,6 +226,10 @@ import_symbol_tuple = _{ import_symbol ~ ("," ~ NEWLINE* ~ import_symbol)* }
 
 import = { "from" ~ "\"" ~ import_source ~ "\"" ~ "import" ~ ("*" | ("{" ~ NEWLINE* ~ import_symbol_tuple ~ NEWLINE* ~ "}") | import_symbol) ~ LINE_END}
 
+/// Tests
+
+test = {"test" ~ function_definition}
+
 /// Program File
 
-file = { SOI ~ NEWLINE* ~ import* ~ NEWLINE* ~ circuit_definition* ~ NEWLINE* ~ function_definition* ~ NEWLINE* ~ EOI }
+file = { SOI ~ NEWLINE* ~ import* ~ NEWLINE* ~ circuit_definition* ~ NEWLINE* ~ function_definition* ~ NEWLINE* ~ test* ~ NEWLINE* ~ EOI }

compiler/src/types.rs

@@ -240,6 +240,11 @@ impl Function {
     }
 }
 
+/// Tests
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Test(pub Function);
+
 /// A simple program with statement expressions, program arguments and program returns.
 #[derive(Debug, Clone)]
 pub struct Program {
@@ -248,6 +253,7 @@ pub struct Program {
     pub imports: Vec<Import>,
     pub circuits: HashMap<Identifier, Circuit>,
     pub functions: HashMap<Identifier, Function>,
+    pub tests: HashMap<Identifier, Test>,
 }
 
 impl<'ast> Program {
@@ -258,6 +264,7 @@ impl<'ast> Program {
             imports: vec![],
             circuits: HashMap::new(),
             functions: HashMap::new(),
+            tests: HashMap::new(),
         }
     }
 

compiler/src/types_from.rs

@@ -788,6 +788,13 @@ impl<'ast> From<ast::Import<'ast>> for Import {
     }
 }
 
+/// pest ast -> Test
+impl<'ast> From<ast::Test<'ast>> for types::Test {
+    fn from(test: ast::Test) -> Self {
+        types::Test(types::Function::from(test.function))
+    }
+}
+
 /// pest ast -> types::Program
 
 impl<'ast> types::Program {
@@ -801,6 +808,7 @@ impl<'ast> types::Program {
 
         let mut circuits = HashMap::new();
         let mut functions = HashMap::new();
+        let mut tests = HashMap::new();
         let mut num_parameters = 0usize;
 
         file.circuits.into_iter().for_each(|circuit| {
@@ -815,6 +823,12 @@ impl<'ast> types::Program {
                 types::Function::from(function_def),
             );
         });
+        file.tests.into_iter().for_each(|test_def| {
+            tests.insert(
+                types::Identifier::from(test_def.function.function_name.clone()),
+                types::Test::from(test_def),
+            );
+        });
 
         if let Some(main_function) = functions.get(&types::Identifier::new("main".into())) {
             num_parameters = main_function.inputs.len();
@@ -826,6 +840,7 @@ impl<'ast> types::Program {
             imports,
             circuits,
             functions,
+            tests,
         }
     }
 }