Validate annotations; refactor FunctionInput in passes

compiler/parser/src/parser/file.rs

@@ -261,9 +261,7 @@ impl ParserContext<'_> {
     fn parse_function_parameter(&mut self) -> Result<FunctionInput> {
         let mode = self.parse_function_parameter_mode()?;
         let (name, type_) = self.parse_typed_ident()?;
-        Ok(FunctionInput::Variable(FunctionInputVariable::new(
-            name, mode, type_, name.span,
-        )))
+        Ok(FunctionInput::new(name, mode, type_, name.span))
     }
 
     /// Returns `true` if the next token is Function or if it is a Const followed by Function.
@@ -279,14 +277,14 @@ impl ParserContext<'_> {
     fn parse_annotation(&mut self) -> Result<Annotation> {
         // Parse the `@` symbol and identifier.
         let start = self.expect(&Token::At)?;
-        let name = self.expect_identifier()?;
-        let span = start + name.span;
+        let identifier = self.expect_identifier()?;
+        let span = start + identifier.span;
 
         // TODO: Verify that this check is sound.
         // Check that there is no whitespace in between the `@` symbol and identifier.
-        match name.span.hi.0 - start.lo.0 > 1 + name.name.as_str().len() as u32 {
+        match identifier.span.hi.0 - start.lo.0 > 1 + identifier.name.as_str().len() as u32 {
             true => Err(ParserError::space_in_annotation(span).into()),
-            false => Ok(Annotation { name, span }),
+            false => Ok(Annotation { identifier, span }),
         }
     }
 

compiler/passes/src/code_generation/visit_program.rs

@@ -148,10 +148,9 @@ impl<'a> CodeGenerator<'a> {
             self.next_register += 1;
 
             self.variable_mapping
-                .insert(&input.get_variable().identifier.name, register_string.clone());
+                .insert(&input.identifier.name, register_string.clone());
 
-            let type_string =
-                self.visit_type_with_visibility(&input.get_variable().type_, Some(input.get_variable().mode()));
+            let type_string = self.visit_type_with_visibility(&input.type_, Some(input.mode()));
             writeln!(function_string, "    input {} as {};", register_string, type_string,)
                 .expect("failed to write to string");
         }

compiler/passes/src/static_single_assignment/rename_program.rs

@@ -18,8 +18,8 @@ use crate::StaticSingleAssigner;
 use itertools::Itertools;
 
 use leo_ast::{
-    Expression, Function, FunctionInput, ProgramReconstructor, ReturnStatement, Statement, StatementReconstructor,
-    TernaryExpression, TupleExpression,
+    Expression, Function, ProgramReconstructor, ReturnStatement, Statement, StatementReconstructor, TernaryExpression,
+    TupleExpression,
 };
 
 impl ProgramReconstructor for StaticSingleAssigner<'_> {
@@ -30,15 +30,9 @@ impl ProgramReconstructor for StaticSingleAssigner<'_> {
 
         // There is no need to reconstruct `function.inputs`.
         // However, for each input, we must add each symbol to the rename table.
-        for input in function.input.iter() {
-            match input {
-                FunctionInput::Variable(function_input_variable) => {
-                    self.rename_table.update(
-                        function_input_variable.identifier.name,
-                        function_input_variable.identifier.name,
-                    );
-                }
-            }
+        for input_variable in function.input.iter() {
+            self.rename_table
+                .update(input_variable.identifier.name, input_variable.identifier.name);
         }
 
         let mut block = self.reconstruct_block(function.block);

compiler/passes/src/type_checking/check_expressions.rs

@@ -552,7 +552,7 @@ impl<'a> ExpressionVisitor<'a> for TypeChecker<'a> {
                         .iter()
                         .zip(input.arguments.iter())
                         .for_each(|(expected, argument)| {
-                            self.visit_expression(argument, &Some(expected.get_variable().type_.clone()));
+                            self.visit_expression(argument, &Some(expected.type_.clone()));
                         });
 
                     Some(ret)

compiler/passes/src/type_checking/check_program.rs

@@ -26,6 +26,14 @@ use std::collections::HashSet;
 
 impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
     fn visit_function(&mut self, input: &'a Function) {
+        // Check that the function's annotations are valid.
+        for annotation in input.annotations.iter() {
+            match annotation.identifier.name {
+                sym::program => self.is_program_function = true,
+                _ => self.emit_err(TypeCheckerError::unknown_annotation(annotation, annotation.span)),
+            }
+        }
+
         let prev_st = std::mem::take(&mut self.symbol_table);
         self.symbol_table
             .swap(prev_st.borrow().lookup_fn_scope(input.name()).unwrap());
@@ -33,8 +41,7 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
 
         self.has_return = false;
         self.parent = Some(input.name());
-        input.input.iter().for_each(|i| {
-            let input_var = i.get_variable();
+        input.input.iter().for_each(|input_var| {
             self.assert_not_tuple(input_var.span, &input_var.type_);
 
             // Check for conflicting variable names.

compiler/passes/src/type_checking/checker.rs

@@ -30,6 +30,9 @@ pub struct TypeChecker<'a> {
     pub(crate) parent: Option<Symbol>,
     pub(crate) has_return: bool,
     pub(crate) negate: bool,
+    /// Are we traversing a program function?
+    /// A "program function" is a function that can be invoked by a user or another program.
+    pub(crate) is_program_function: bool,
 }
 
 const BOOLEAN_TYPE: Type = Type::Boolean;
@@ -63,6 +66,7 @@ impl<'a> TypeChecker<'a> {
     /// Returns a new type checker given a symbol table and error handler.
     pub fn new(symbol_table: SymbolTable, handler: &'a Handler) -> Self {
         Self {
+            is_program_function: false,
             symbol_table: RefCell::new(symbol_table),
             handler,
             parent: None,

compiler/span/src/symbol.rs

@@ -218,6 +218,7 @@ symbols! {
     owner,
     gates,
     _nonce,
+    program,
 
     // input file
     registers,

errors/src/errors/type_checker/type_checker_error.rs

@@ -273,4 +273,12 @@ create_messages!(
         msg: format!("Loop body contains a return statement or always returns."),
         help: Some("Remove the code in the loop body that always returns.".to_string()),
     }
+
+    // TODO: Consider emitting a warning instead of an error.
+    @formatted
+    unknown_annotation {
+        args: (annotation: impl Display),
+        msg: format!("Unknown annotation: `{annotation}`."),
+        help: Some("Use a valid annotation. The Leo compiler supports: `@program`".to_string()),
+    }
 );