Merge pull request #245 from AleoHQ/feature/context-annotation

Feature/context annotation

Cargo.lock

@@ -1122,6 +1122,7 @@ dependencies = [
  "leo-ast",
  "leo-gadgets",
  "leo-input",
+ "leo-package",
  "leo-state",
  "leo-typed",
  "log",

ast/src/annotations/annotation_arguments.rs

@@ -0,0 +1,27 @@
+use crate::{
+    ast::{span_into_string, Rule},
+    SpanDef,
+};
+
+use pest::Span;
+use pest_ast::FromPest;
+use serde::Serialize;
+
+#[derive(Clone, Debug, FromPest, PartialEq, Serialize)]
+#[pest_ast(rule(Rule::annotation_arguments))]
+pub struct AnnotationArguments<'ast> {
+    pub arguments: Vec<AnnotationArgument<'ast>>,
+    #[pest_ast(outer())]
+    #[serde(with = "SpanDef")]
+    pub span: Span<'ast>,
+}
+
+#[derive(Clone, Debug, FromPest, PartialEq, Serialize)]
+#[pest_ast(rule(Rule::annotation_argument))]
+pub struct AnnotationArgument<'ast> {
+    #[pest_ast(outer(with(span_into_string)))]
+    pub value: String,
+    #[pest_ast(outer())]
+    #[serde(with = "SpanDef")]
+    pub span: Span<'ast>,
+}

ast/src/annotations/annotation_name.rs

@@ -0,0 +1,19 @@
+use crate::{ast::Rule, SpanDef};
+
+use pest::Span;
+use pest_ast::FromPest;
+use serde::Serialize;
+
+#[derive(Clone, Debug, FromPest, PartialEq, Serialize)]
+#[pest_ast(rule(Rule::annotation_name))]
+pub enum AnnotationName<'ast> {
+    Context(Context<'ast>),
+}
+
+#[derive(Clone, Debug, FromPest, PartialEq, Serialize)]
+#[pest_ast(rule(Rule::context))]
+pub struct Context<'ast> {
+    #[pest_ast(outer())]
+    #[serde(with = "SpanDef")]
+    pub span: Span<'ast>,
+}

ast/src/annotations/annotation_symbol.rs

@@ -0,0 +1,13 @@
+use crate::{ast::Rule, SpanDef};
+
+use pest::Span;
+use pest_ast::FromPest;
+use serde::Serialize;
+
+#[derive(Clone, Debug, FromPest, PartialEq, Serialize)]
+#[pest_ast(rule(Rule::annotation_symbol))]
+pub struct AnnotationSymbol<'ast> {
+    #[pest_ast(outer())]
+    #[serde(with = "SpanDef")]
+    pub span: Span<'ast>,
+}

ast/src/annotations/annotations.rs

@@ -0,0 +1,20 @@
+use crate::{
+    annotations::{AnnotationArguments, AnnotationName, AnnotationSymbol},
+    ast::Rule,
+    SpanDef,
+};
+
+use pest::Span;
+use pest_ast::FromPest;
+use serde::Serialize;
+
+#[derive(Clone, Debug, FromPest, PartialEq, Serialize)]
+#[pest_ast(rule(Rule::annotation))]
+pub struct Annotation<'ast> {
+    pub symbol: AnnotationSymbol<'ast>,
+    pub name: AnnotationName<'ast>,
+    pub arguments: AnnotationArguments<'ast>,
+    #[pest_ast(outer())]
+    #[serde(with = "SpanDef")]
+    pub span: Span<'ast>,
+}

ast/src/annotations/mod.rs

@@ -0,0 +1,11 @@
+pub mod annotations;
+pub use annotations::*;
+
+pub mod annotation_symbol;
+pub use annotation_symbol::*;
+
+pub mod annotation_name;
+pub use annotation_name::*;
+
+pub mod annotation_arguments;
+pub use annotation_arguments::*;

ast/src/definitions/annotated_definition.rs

@@ -0,0 +1,15 @@
+use crate::{annotations::Annotation, ast::Rule, definitions::Definition, SpanDef};
+
+use pest::Span;
+use pest_ast::FromPest;
+use serde::Serialize;
+
+#[derive(Clone, Debug, FromPest, PartialEq, Serialize)]
+#[pest_ast(rule(Rule::definition_annotated))]
+pub struct AnnotatedDefinition<'ast> {
+    pub annotation: Annotation<'ast>,
+    pub definition: Box<Definition<'ast>>,
+    #[pest_ast(outer())]
+    #[serde(with = "SpanDef")]
+    pub span: Span<'ast>,
+}

ast/src/definitions/definition.rs

@@ -1,6 +1,7 @@
 use crate::{
     ast::Rule,
     circuits::Circuit,
+    definitions::AnnotatedDefinition,
     functions::{Function, TestFunction},
     imports::Import,
 };
@@ -11,6 +12,7 @@ use serde::Serialize;
 #[derive(Clone, Debug, FromPest, PartialEq, Serialize)]
 #[pest_ast(rule(Rule::definition))]
 pub enum Definition<'ast> {
+    Annotated(AnnotatedDefinition<'ast>),
     Import(Import<'ast>),
     Circuit(Circuit<'ast>),
     Function(Function<'ast>),

ast/src/definitions/mod.rs

@@ -1,2 +1,5 @@
+pub mod annotated_definition;
+pub use annotated_definition::*;
+
 pub mod definition;
 pub use definition::*;

ast/src/leo.pest

@@ -8,12 +8,16 @@ file = { SOI ~ NEWLINE* ~ definition* ~ NEWLINE* ~ EOI }
 
 // Declared in definitions/definition.rs
 definition = {
-    import
+    definition_annotated
+    | import
     | circuit
     | function
     | test_function
 }
 
+// Declared in definitions/annotated_definition.rs
+definition_annotated = { annotation ~ NEWLINE* ~ definition}
+
 // Declared in common/identifier.rs
 identifier = @{ ((!protected_name ~ ASCII_ALPHA) | (protected_name ~ (ASCII_ALPHANUMERIC | "_"))) ~ (ASCII_ALPHANUMERIC | "_")* }
 protected_name = {
@@ -468,3 +472,23 @@ debug = {"debug"}
 // Declared in macros/error.rs
 error = {"error"}
 
+/// Annotations
+
+// Declared in annotations/annotation.rs
+annotation = ${annotation_symbol ~ annotation_name ~ annotation_arguments}
+
+// Declared in annotations/annotation_symbol.rs
+annotation_symbol = ${"@"}
+
+// Declared in annotations/annotation_name.rs
+annotation_name = {
+    context
+}
+
+// Declared in annotations/context.rs
+context = {"context"}
+
+// Declared in annotations/annotation_argument.rs
+annotation_arguments = !{"(" ~ NEWLINE* ~ annotation_argument ~ ("," ~ NEWLINE* ~ annotation_argument)* ~ ","? ~  NEWLINE* ~ ")"}
+
+annotation_argument = @{ (ASCII_ALPHANUMERIC | "_")+ }

ast/src/lib.rs

@@ -8,6 +8,7 @@ extern crate thiserror;
 mod ast;
 
 pub mod access;
+pub mod annotations;
 pub mod circuits;
 pub mod common;
 pub mod definitions;

compiler/Cargo.toml

@@ -8,6 +8,7 @@ edition = "2018"
 leo-ast = { path = "../ast", version = "0.1.0" }
 leo-gadgets = { path = "../gadgets", version = "0.1.0" }
 leo-input = { path = "../input", version = "0.1.0" }
+leo-package = { path = "../package", version = "0.1.0"}
 leo-typed = { path = "../typed", version = "0.1.0" }
 leo-state = { path = "../state", version = "0.1.0" }
 

compiler/src/compiler.rs

@@ -10,6 +10,7 @@ use crate::{
 };
 use leo_ast::LeoAst;
 use leo_input::LeoInputParser;
+use leo_package::inputs::InputPairs;
 use leo_state::verify_local_data_commitment;
 use leo_typed::{Input, LeoTypedAst, MainInput, Program};
 
@@ -163,8 +164,13 @@ impl<F: Field + PrimeField, G: GroupType<F>> Compiler<F, G> {
     }
 
     /// Synthesizes the circuit for test functions with program input.
-    pub fn compile_test_constraints(self) -> Result<(), CompilerError> {
-        generate_test_constraints::<F, G>(self.program, self.program_input, &self.imported_programs)
+    pub fn compile_test_constraints(self, input_pairs: InputPairs) -> Result<(), CompilerError> {
+        generate_test_constraints::<F, G>(
+            self.program,
+            input_pairs,
+            &self.imported_programs,
+            &self.output_directory,
+        )
     }
 
     /// Calls the internal generate_constraints method with arguments
@@ -216,6 +222,9 @@ impl<F: Field + PrimeField, G: GroupType<F>> ConstraintSynthesizer<F> for Compil
 
         // Write results to file
         let output_file = OutputFile::new(&package_name);
+
+        log::info!("Writing to output registers...");
+
         output_file.write(&output_directory, result.bytes()).unwrap();
 
         Ok(())

compiler/src/constraints/constraints.rs

@@ -8,13 +8,17 @@ use crate::{
     GroupType,
     ImportParser,
     OutputBytes,
+    OutputFile,
 };
 use leo_typed::{Input, Program};
 
+use leo_input::LeoInputParser;
+use leo_package::inputs::InputPairs;
 use snarkos_models::{
     curves::{Field, PrimeField},
     gadgets::r1cs::{ConstraintSystem, TestConstraintSystem},
 };
+use std::path::PathBuf;
 
 pub fn generate_constraints<F: Field + PrimeField, G: GroupType<F>, CS: ConstraintSystem<F>>(
     cs: &mut CS,
@@ -43,27 +47,61 @@ pub fn generate_constraints<F: Field + PrimeField, G: GroupType<F>, CS: Constrai
 
 pub fn generate_test_constraints<F: Field + PrimeField, G: GroupType<F>>(
     program: Program,
-    input: Input,
+    input: InputPairs,
     imported_programs: &ImportParser,
+    output_directory: &PathBuf,
 ) -> Result<(), CompilerError> {
     let mut resolved_program = ConstrainedProgram::<F, G>::new();
     let program_name = program.get_name();
 
     let tests = program.tests.clone();
 
+    // Store definitions
     resolved_program.store_definitions(program, imported_programs)?;
 
+    // Get default input
+    let default = input.pairs.get(&program_name);
+
     log::info!("Running {} tests", tests.len());
 
-    for (test_name, test_function) in tests.into_iter() {
+    for (test_name, test) in tests.into_iter() {
         let cs = &mut TestConstraintSystem::<F>::new();
         let full_test_name = format!("{}::{}", program_name.clone(), test_name.to_string());
+        let mut output_file_name = program_name.clone();
 
-        let result = resolved_program.enforce_main_function(
+        // get input file name from annotation or use test_name
+        let input_pair = match test.input_file {
+            Some(file_id) => {
+                let file_name = file_id.name;
+
+                output_file_name = file_name.clone();
+
+                match input.pairs.get(&file_name) {
+                    Some(pair) => pair.to_owned(),
+                    None => return Err(CompilerError::InvalidTestContext(file_name)),
+                }
+            }
+            None => default.ok_or(CompilerError::NoTestInput)?,
+        };
+
+        // parse input files to abstract syntax trees
+        let input_file = &input_pair.input_file;
+        let state_file = &input_pair.state_file;
+
+        let input_ast = LeoInputParser::parse_file(input_file)?;
+        let state_ast = LeoInputParser::parse_file(state_file)?;
+
+        // parse input files into input struct
+        let mut input = Input::new();
+        input.parse_input(input_ast)?;
+        input.parse_state(state_ast)?;
+
+        // run test function on new program with input
+        let result = resolved_program.clone().enforce_main_function(
             cs,
             program_name.clone(),
-            test_function.0,
-            input.clone(), // pass program input into every test
+            test.function,
+            input, // pass program input into every test
         );
 
         if result.is_ok() {
@@ -72,6 +110,14 @@ pub fn generate_test_constraints<F: Field + PrimeField, G: GroupType<F>>(
                 full_test_name,
                 cs.is_satisfied()
             );
+
+            // write result to file
+            let output = result?;
+            let output_file = OutputFile::new(&output_file_name);
+
+            log::info!("\tWriting output to registers in `{}.out` ...", output_file_name);
+
+            output_file.write(output_directory, output.bytes()).unwrap();
         } else {
             log::error!("test {} errored: {}", full_test_name, result.unwrap_err());
         }

compiler/src/errors/compiler.rs

@@ -14,6 +14,9 @@ pub enum CompilerError {
     #[error("{}", _0)]
     InputParserError(#[from] InputParserError),
 
+    #[error("Cannot find input files with context name `{}`", _0)]
+    InvalidTestContext(String),
+
     #[error("{}", _0)]
     FunctionError(#[from] FunctionError),
 
@@ -29,6 +32,9 @@ pub enum CompilerError {
     #[error("`main` must be a function")]
     NoMainFunction,
 
+    #[error("Failed to find input files for the current test")]
+    NoTestInput,
+
     #[error("{}", _0)]
     OutputError(#[from] OutputFileError),
 

compiler/src/output/output_file.rs

@@ -40,7 +40,6 @@ impl OutputFile {
         // create output file
         let path = self.setup_file_path(path);
         let mut file = File::create(&path)?;
-        log::info!("Writing to output registers...");
 
         Ok(file.write_all(bytes)?)
     }

compiler/src/program/program.rs

@@ -6,6 +6,7 @@ use snarkos_models::curves::{Field, PrimeField};
 
 use std::collections::HashMap;
 
+#[derive(Clone)]
 pub struct ConstrainedProgram<F: Field + PrimeField, G: GroupType<F>> {
     pub identifiers: HashMap<String, ConstrainedValue<F, G>>,
 }

leo/commands/publish.rs

@@ -1,7 +1,7 @@
 use crate::{
     cli::*,
     cli_types::*,
-    commands::{BuildCommand, LoginCommand},
+    commands::LoginCommand,
     errors::{
         commands::PublishError::{ConnectionUnavalaible, PackageNotPublished},
         CLIError,

leo/commands/test.rs

@@ -6,7 +6,7 @@ use crate::{
 use leo_compiler::{compiler::Compiler, group::targets::edwards_bls12::EdwardsGroupType};
 use leo_package::{
     inputs::*,
-    outputs::OUTPUTS_DIRECTORY_NAME,
+    outputs::{OutputsDirectory, OUTPUTS_DIRECTORY_NAME},
     root::Manifest,
     source::{MainFile, MAIN_FILE_NAME, SOURCE_DIRECTORY_NAME},
 };
@@ -63,27 +63,23 @@ impl CLI for TestCommand {
         let mut output_directory = package_path.clone();
         output_directory.push(OUTPUTS_DIRECTORY_NAME);
 
-        // Load the input file at `package_name`
-        let input_string = InputFile::new(&package_name).read_from(&path)?;
+        // Create the output directory
+        OutputsDirectory::create(&package_path)?;
 
-        // Load the state file at `package_name.in`
-        let state_string = StateFile::new(&package_name).read_from(&path)?;
-
-        // Compute the current program checksum
-        let program = Compiler::<Fq, EdwardsGroupType>::parse_program_with_input(
+        // Parse the current main program file
+        let program = Compiler::<Fq, EdwardsGroupType>::parse_program_without_input(
             package_name.clone(),
             main_file_path.clone(),
             output_directory,
-            &input_string,
-            &state_string,
         )?;
 
-        // Generate the program on the constraint system and verify correctness
-        {
-            let temporary_program = program.clone();
-            let output = temporary_program.compile_test_constraints()?;
-            log::debug!("Compiled constraints - {:#?}", output);
-        }
+        // Parse all inputs as input pairs
+        let pairs = InputPairs::try_from(&package_path)?;
+
+        // Run tests
+        let temporary_program = program.clone();
+        let output = temporary_program.compile_test_constraints(pairs)?;
+        log::debug!("Compiled constraints - {:#?}", output);
 
         Ok(())
     }

package/src/errors/inputs/directory.rs

@@ -1,3 +1,5 @@
+use crate::{InputFileError, StateFileError};
+
 use std::{ffi::OsString, fs::FileType, io};
 
 #[derive(Debug, Error)]
@@ -11,15 +13,24 @@ pub enum InputsDirectoryError {
     #[error("file {:?} extension getting", _0)]
     GettingFileExtension(OsString),
 
+    #[error("file {:?} name getting", _0)]
+    GettingFileName(OsString),
+
     #[error("file {:?} type getting: {}", _0, _1)]
     GettingFileType(OsString, io::Error),
 
+    #[error("{}", _0)]
+    InputFileError(#[from] InputFileError),
+
     #[error("invalid file {:?} extension: {:?}", _0, _1)]
-    InvalidFileExtension(OsString, OsString),
+    InvalidFileExtension(String, OsString),
 
     #[error("invalid file {:?} type: {:?}", _0, _1)]
     InvalidFileType(OsString, FileType),
 
     #[error("reading: {}", _0)]
     Reading(io::Error),
+
+    #[error("{}", _0)]
+    StateFileError(#[from] StateFileError),
 }

package/src/inputs/directory.rs

@@ -1,6 +1,6 @@
-use crate::{errors::InputsDirectoryError, inputs::INPUT_FILE_EXTENSION};
+use crate::errors::InputsDirectoryError;
 
-use std::{fs, path::PathBuf};
+use std::{fs, fs::ReadDir, path::PathBuf};
 
 pub static INPUTS_DIRECTORY_NAME: &str = "inputs/";
 
@@ -17,42 +17,42 @@ impl InputsDirectory {
         fs::create_dir_all(&path).map_err(InputsDirectoryError::Creating)
     }
 
-    /// Returns a list of files in the source directory.
+    /// Returns a list of files in the input directory.
     pub fn files(path: &PathBuf) -> Result<Vec<PathBuf>, InputsDirectoryError> {
         let mut path = path.to_owned();
         path.push(PathBuf::from(INPUTS_DIRECTORY_NAME));
         let directory = fs::read_dir(&path).map_err(InputsDirectoryError::Reading)?;
 
         let mut file_paths = Vec::new();
-        for file_entry in directory.into_iter() {
-            let file_entry = file_entry.map_err(InputsDirectoryError::GettingFileEntry)?;
-            let file_path = file_entry.path();
-
-            // Verify that the entry is structured as a valid file
-            let file_type = file_entry
-                .file_type()
-                .map_err(|error| InputsDirectoryError::GettingFileType(file_path.as_os_str().to_owned(), error))?;
-            if !file_type.is_file() {
-                return Err(InputsDirectoryError::InvalidFileType(
-                    file_path.as_os_str().to_owned(),
-                    file_type,
-                ));
-            }
-
-            // Verify that the file has the default file extension
-            let file_extension = file_path
-                .extension()
-                .ok_or_else(|| InputsDirectoryError::GettingFileExtension(file_path.as_os_str().to_owned()))?;
-            if file_extension != INPUT_FILE_EXTENSION {
-                return Err(InputsDirectoryError::InvalidFileExtension(
-                    file_path.as_os_str().to_owned(),
-                    file_extension.to_owned(),
-                ));
-            }
-
-            file_paths.push(file_path);
-        }
+        parse_file_paths(directory, &mut file_paths)?;
 
         Ok(file_paths)
     }
 }
+
+fn parse_file_paths(directory: ReadDir, file_paths: &mut Vec<PathBuf>) -> Result<(), InputsDirectoryError> {
+    for file_entry in directory.into_iter() {
+        let file_entry = file_entry.map_err(InputsDirectoryError::GettingFileEntry)?;
+        let file_path = file_entry.path();
+
+        // Verify that the entry is structured as a valid file or directory
+        let file_type = file_entry
+            .file_type()
+            .map_err(|error| InputsDirectoryError::GettingFileType(file_path.as_os_str().to_owned(), error))?;
+        if file_type.is_dir() {
+            let directory = fs::read_dir(&file_path).map_err(InputsDirectoryError::Reading)?;
+
+            parse_file_paths(directory, file_paths)?;
+            continue;
+        } else if !file_type.is_file() {
+            return Err(InputsDirectoryError::InvalidFileType(
+                file_path.as_os_str().to_owned(),
+                file_type,
+            ));
+        }
+
+        file_paths.push(file_path);
+    }
+
+    Ok(())
+}

package/src/inputs/mod.rs

@@ -4,5 +4,8 @@ pub use directory::*;
 pub mod input;
 pub use input::*;
 
+pub mod pairs;
+pub use pairs::*;
+
 pub mod state;
 pub use state::*;

package/src/inputs/pairs.rs

@@ -0,0 +1,79 @@
+use crate::{
+    inputs::{InputFile, InputsDirectory, StateFile, INPUT_FILE_EXTENSION, STATE_FILE_EXTENSION},
+    InputsDirectoryError,
+};
+
+use std::{collections::HashMap, convert::TryFrom, path::PathBuf};
+
+pub struct InputPairs {
+    /// Maps file names to input file pairs
+    pub pairs: HashMap<String, InputPair>,
+}
+
+pub struct InputPair {
+    pub input_file: String,
+    pub state_file: String,
+}
+
+impl InputPairs {
+    pub fn new() -> Self {
+        Self { pairs: HashMap::new() }
+    }
+}
+
+impl TryFrom<&PathBuf> for InputPairs {
+    type Error = InputsDirectoryError;
+
+    fn try_from(directory: &PathBuf) -> Result<Self, Self::Error> {
+        let files = InputsDirectory::files(directory)?;
+
+        let mut pairs = HashMap::<String, InputPair>::new();
+
+        for file in files {
+            let file_extension = file
+                .extension()
+                .ok_or_else(|| InputsDirectoryError::GettingFileExtension(file.as_os_str().to_owned()))?;
+
+            let file_name = file
+                .file_stem()
+                .ok_or(InputsDirectoryError::GettingFileName(file.as_os_str().to_owned()))?
+                .to_str()
+                .ok_or(InputsDirectoryError::GettingFileName(file.as_os_str().to_owned()))?;
+
+            if file_extension == INPUT_FILE_EXTENSION.trim_start_matches(".") {
+                let input_file = InputFile::new(file_name).read_from(&file)?;
+
+                if pairs.contains_key(file_name) {
+                    let pair = pairs.get_mut(file_name).unwrap();
+                    pair.input_file = input_file;
+                } else {
+                    let pair = InputPair {
+                        input_file,
+                        state_file: "".to_owned(),
+                    };
+                    pairs.insert(file_name.to_owned(), pair);
+                }
+            } else if file_extension == STATE_FILE_EXTENSION.trim_start_matches(".") {
+                let state_file = StateFile::new(file_name).read_from(&file)?;
+
+                if pairs.contains_key(file_name) {
+                    let pair = pairs.get_mut(file_name).unwrap();
+                    pair.state_file = state_file;
+                } else {
+                    let pair = InputPair {
+                        input_file: "".to_owned(),
+                        state_file,
+                    };
+                    pairs.insert(file_name.to_owned(), pair);
+                }
+            } else {
+                return Err(InputsDirectoryError::InvalidFileExtension(
+                    file_name.to_owned(),
+                    file_extension.to_owned(),
+                ));
+            }
+        }
+
+        Ok(InputPairs { pairs })
+    }
+}

typed/src/annotation.rs

@@ -0,0 +1,58 @@
+use crate::{Circuit, Function, Identifier, Import, InputVariable, TestFunction};
+use leo_ast::{
+    annotations::{Annotation, AnnotationArguments, AnnotationName},
+    definitions::{AnnotatedDefinition, Definition},
+};
+
+use std::collections::HashMap;
+
+pub fn load_annotation(
+    annotated_definition: AnnotatedDefinition,
+    _imports: &mut Vec<Import>,
+    _circuits: &mut HashMap<Identifier, Circuit>,
+    _functions: &mut HashMap<Identifier, Function>,
+    tests: &mut HashMap<Identifier, TestFunction>,
+    _expected: &mut Vec<InputVariable>,
+) {
+    let ast_annotation = annotated_definition.annotation;
+    let ast_definition = *annotated_definition.definition;
+
+    match ast_definition {
+        Definition::Import(_) => unimplemented!("annotated imports are not supported yet"),
+        Definition::Circuit(_) => unimplemented!("annotated circuits are not supported yet"),
+        Definition::Function(_) => unimplemented!("annotated functions are not supported yet"),
+        Definition::TestFunction(ast_test) => {
+            let test = TestFunction::from(ast_test);
+            load_annotated_test(test, ast_annotation, tests)
+        }
+        Definition::Annotated(_) => unimplemented!("nested annotations are not supported yet"),
+    }
+}
+
+pub fn load_annotated_test(test: TestFunction, annotation: Annotation, tests: &mut HashMap<Identifier, TestFunction>) {
+    let name = annotation.name;
+    let ast_arguments = annotation.arguments;
+
+    match name {
+        AnnotationName::Context(_) => load_annotated_test_context(test, ast_arguments, tests),
+    }
+}
+
+pub fn load_annotated_test_context(
+    mut test: TestFunction,
+    ast_arguments: AnnotationArguments,
+    tests: &mut HashMap<Identifier, TestFunction>,
+) {
+    let arguments = ast_arguments.arguments;
+
+    if arguments.len() != 1 {
+        panic!("text context annotation must have one argument identifier")
+    }
+
+    let ast_input_file = arguments[0].to_owned();
+    let input_file = Identifier::from(ast_input_file);
+
+    test.input_file = Some(input_file);
+
+    tests.insert(test.function.identifier.clone(), test);
+}

typed/src/common/identifier.rs

@@ -1,5 +1,9 @@
 use crate::Span;
-use leo_ast::{common::Identifier as AstIdentifier, imports::PackageName as AstPackageName};
+use leo_ast::{
+    annotations::AnnotationArgument,
+    common::Identifier as AstIdentifier,
+    imports::PackageName as AstPackageName,
+};
 use leo_input::common::Identifier as InputAstIdentifier;
 
 use serde::{
@@ -55,6 +59,15 @@ impl<'ast> From<InputAstIdentifier<'ast>> for Identifier {
     }
 }
 
+impl<'ast> From<AnnotationArgument<'ast>> for Identifier {
+    fn from(argument: AnnotationArgument<'ast>) -> Self {
+        Self {
+            name: argument.value,
+            span: Span::from(argument.span),
+        }
+    }
+}
+
 impl fmt::Display for Identifier {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(f, "{}", self.name)

typed/src/functions/test_function.rs

@@ -1,13 +1,19 @@
-use crate::Function;
+use crate::{Function, Identifier};
 use leo_ast::functions::TestFunction as AstTestFunction;
 
 use serde::{Deserialize, Serialize};
 
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
-pub struct TestFunction(pub Function);
+pub struct TestFunction {
+    pub function: Function,
+    pub input_file: Option<Identifier>,
+}
 
 impl<'ast> From<AstTestFunction<'ast>> for TestFunction {
     fn from(test: AstTestFunction) -> Self {
-        TestFunction(Function::from(test.function))
+        TestFunction {
+            function: Function::from(test.function),
+            input_file: None, // pass custom input file with `@context` annotation
+        }
     }
 }

typed/src/input/input.rs

@@ -6,6 +6,7 @@ use leo_input::{
 
 #[derive(Clone, PartialEq, Eq)]
 pub struct Input {
+    name: String,
     program_input: ProgramInput,
     program_state: ProgramState,
 }
@@ -13,6 +14,7 @@ pub struct Input {
 impl Input {
     pub fn new() -> Self {
         Self {
+            name: "default".to_owned(),
             program_input: ProgramInput::new(),
             program_state: ProgramState::new(),
         }
@@ -25,6 +27,7 @@ impl Input {
         let state = self.program_state.empty();
 
         Self {
+            name: self.name.clone(),
             program_input: input,
             program_state: state,
         }

typed/src/lib.rs

@@ -1,6 +1,9 @@
 //! A typed syntax tree is represented as a `Program` and consists of import, circuit, and function definitions.
 //! Each defined type consists of typed statements and expressions.
 
+pub mod annotation;
+pub use self::annotation::*;
+
 pub mod circuits;
 pub use self::circuits::*;
 

typed/src/program.rs

@@ -1,7 +1,7 @@
 //! A typed Leo program consists of import, circuit, and function definitions.
 //! Each defined type consists of typed statements and expressions.
 
-use crate::{Circuit, Function, Identifier, Import, InputVariable, TestFunction};
+use crate::{load_annotation, Circuit, Function, Identifier, Import, InputVariable, TestFunction};
 use leo_ast::{definitions::Definition, files::File};
 
 use serde::{Deserialize, Serialize};
@@ -47,7 +47,17 @@ impl<'ast> Program {
                 }
                 Definition::TestFunction(test_def) => {
                     let test = TestFunction::from(test_def);
-                    tests.insert(test.0.identifier.clone(), test);
+                    tests.insert(test.function.identifier.clone(), test);
+                }
+                Definition::Annotated(annotated_definition) => {
+                    load_annotation(
+                        annotated_definition,
+                        &mut imports,
+                        &mut circuits,
+                        &mut functions,
+                        &mut tests,
+                        &mut expected_input,
+                    );
                 }
             });