Typecheck helper functions; add codegen support

compiler/ast/src/functions/function_input.rs

@@ -22,6 +22,7 @@ use std::fmt;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
 pub enum ParamMode {
+    None,
     Const,
     Private,
     Public,
@@ -32,6 +33,7 @@ impl fmt::Display for ParamMode {
         use ParamMode::*;
 
         match self {
+            None => write!(f, ""),
             Const => write!(f, "const"),
             Private => write!(f, "private"),
             Public => write!(f, "public"),

compiler/parser/src/parser/file.rs

@@ -235,6 +235,7 @@ impl ParserContext<'_> {
 
     /// Returns a [`ParamMode`] AST node if the next tokens represent a function parameter mode.
     pub(super) fn parse_function_parameter_mode(&mut self) -> Result<ParamMode> {
+        // TODO: Allow explicit "private" mode.
         let public = self.eat(&Token::Public).then(|| self.prev_token.span);
         let constant = self.eat(&Token::Constant).then(|| self.prev_token.span);
         let const_ = self.eat(&Token::Const).then(|| self.prev_token.span);
@@ -246,7 +247,7 @@ impl ParserContext<'_> {
         match (public, constant, const_) {
             (None, Some(_), None) => Ok(ParamMode::Const),
             (None, None, Some(_)) => Ok(ParamMode::Const),
-            (None, None, None) => Ok(ParamMode::Private),
+            (None, None, None) => Ok(ParamMode::None),
             (Some(_), None, None) => Ok(ParamMode::Public),
             (Some(m1), Some(m2), None) | (Some(m1), None, Some(m2)) | (None, Some(m1), Some(m2)) => {
                 Err(ParserError::inputs_multiple_variable_types_specified(m1 + m2).into())
@@ -292,6 +293,7 @@ impl ParserContext<'_> {
     /// and function definition.
     fn parse_function(&mut self) -> Result<(Identifier, Function)> {
         // TODO: Handle dangling annotations.
+        // TODO: Handle duplicate annotations.
         // Parse annotations, if they exist.
         let mut annotations = Vec::new();
         while self.look_ahead(0, |t| &t.token) == &Token::At {

compiler/passes/src/code_generation/generator.rs

@@ -32,6 +32,9 @@ pub struct CodeGenerator<'a> {
     /// The first element of the tuple indicate whether the composite is a record or not.
     /// The second element of the tuple is a string modifier used for code generation.
     pub(crate) composite_mapping: IndexMap<&'a Symbol, (bool, String)>,
+    /// 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,
 }
 
 impl<'a> CodeGenerator<'a> {
@@ -43,6 +46,7 @@ impl<'a> CodeGenerator<'a> {
             current_function: None,
             variable_mapping: IndexMap::new(),
             composite_mapping: IndexMap::new(),
+            is_program_function: false,
         }
     }
 }

compiler/passes/src/code_generation/visit_program.rs

@@ -20,6 +20,7 @@ use leo_ast::{Circuit, CircuitMember, Function, Identifier, Program};
 
 use indexmap::IndexMap;
 use itertools::Itertools;
+use leo_span::sym;
 use std::fmt::Write as _;
 
 impl<'a> CodeGenerator<'a> {
@@ -134,6 +135,13 @@ impl<'a> CodeGenerator<'a> {
     }
 
     fn visit_function(&mut self, function: &'a Function) -> String {
+        // If the function is annotated with `@program`, then it is a program function.
+        for annotation in function.annotations.iter() {
+            if annotation.identifier.name == sym::program {
+                self.is_program_function = true;
+            }
+        }
+
         // Initialize the state of `self` with the appropriate values before visiting `function`.
         self.next_register = 0;
         self.variable_mapping = IndexMap::new();
@@ -150,7 +158,7 @@ impl<'a> CodeGenerator<'a> {
             self.variable_mapping
                 .insert(&input.identifier.name, register_string.clone());
 
-            let type_string = self.visit_type_with_visibility(&input.type_, Some(input.mode()));
+            let type_string = self.visit_type_with_visibility(&input.type_, input.mode());
             writeln!(function_string, "    input {} as {};", register_string, type_string,)
                 .expect("failed to write to string");
         }
@@ -159,6 +167,9 @@ impl<'a> CodeGenerator<'a> {
         let block_string = self.visit_block(&function.block);
         function_string.push_str(&block_string);
 
+        // Unset `is_program_function` after visiting `function`.
+        self.is_program_function = false;
+
         function_string
     }
 }

compiler/passes/src/code_generation/visit_statements.rs

@@ -18,7 +18,7 @@ use crate::CodeGenerator;
 
 use leo_ast::{
     AssignStatement, Block, ConditionalStatement, ConsoleStatement, DefinitionStatement, Expression,
-    IterationStatement, ReturnStatement, Statement,
+    IterationStatement, ParamMode, ReturnStatement, Statement,
 };
 
 use itertools::Itertools;
@@ -38,8 +38,8 @@ impl<'a> CodeGenerator<'a> {
 
     fn visit_return(&mut self, input: &'a ReturnStatement) -> String {
         let (operand, mut expression_instructions) = self.visit_expression(&input.expression);
-        let types = self.visit_return_type(&self.current_function.unwrap().output, None);
         // TODO: Bytecode functions have an associated output mode. Currently defaulting to private since we do not yet support this at the Leo level.
+        let types = self.visit_return_type(&self.current_function.unwrap().output, ParamMode::Private);
         let mut instructions = operand
             .split('\n')
             .into_iter()

compiler/passes/src/code_generation/visit_type.rs

@@ -52,22 +52,29 @@ impl<'a> CodeGenerator<'a> {
         }
     }
 
-    pub(crate) fn visit_type_with_visibility(&mut self, input: &'a Type, visibility: Option<ParamMode>) -> String {
+    pub(crate) fn visit_type_with_visibility(&mut self, input: &'a Type, visibility: ParamMode) -> String {
         let mut type_string = self.visit_type(input);
 
         if let Type::Identifier(_) = input {
             // Do not append anything for record and circuit types.
         } else {
-            // Append `.private` to return type.
             // todo: CAUTION private by default.
-            write!(type_string, ".{}", visibility.unwrap_or(ParamMode::Private)).expect("failed to write to string");
+            // Only program functions need a visibility associated with the input type.
+            if self.is_program_function {
+                // If a visibility is not provided in a program function, then it is private by default.
+                let visibility = match visibility {
+                    ParamMode::None => ParamMode::Private,
+                    _ => visibility,
+                };
+                write!(type_string, ".{}", visibility).expect("failed to write to string");
+            }
         }
 
         type_string
     }
 
     /// Returns one or more types equal to the number of return tuple members.
-    pub(crate) fn visit_return_type(&mut self, input: &'a Type, visibility: Option<ParamMode>) -> Vec<String> {
+    pub(crate) fn visit_return_type(&mut self, input: &'a Type, visibility: ParamMode) -> Vec<String> {
         // Handle return tuples.
         if let Type::Tuple(types) = input {
             types

compiler/passes/src/type_checking/check_program.rs

@@ -29,6 +29,7 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
         // Check that the function's annotations are valid.
         for annotation in input.annotations.iter() {
             match annotation.identifier.name {
+                // Set `is_program_function` to true if the corresponding annotation is found.
                 sym::program => self.is_program_function = true,
                 _ => self.emit_err(TypeCheckerError::unknown_annotation(annotation, annotation.span)),
             }
@@ -44,6 +45,13 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
         input.input.iter().for_each(|input_var| {
             self.assert_not_tuple(input_var.span, &input_var.type_);
 
+            // If the function is not a program function, then check that the parameters do not have an associated mode.
+            if !self.is_program_function && input_var.mode() != ParamMode::None {
+                self.emit_err(TypeCheckerError::helper_function_inputs_cannot_have_modes(
+                    input_var.span,
+                ));
+            }
+
             // Check for conflicting variable names.
             if let Err(err) = self.symbol_table.borrow_mut().insert_variable(
                 input_var.identifier.name,
@@ -72,6 +80,9 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
         let prev_st = *self.symbol_table.borrow_mut().parent.take().unwrap();
         self.symbol_table.swap(prev_st.lookup_fn_scope(input.name()).unwrap());
         self.symbol_table = RefCell::new(prev_st);
+
+        // Unset `is_program_function` flag.
+        self.is_program_function = false;
     }
 
     fn visit_circuit(&mut self, input: &'a Circuit) {

errors/src/errors/type_checker/type_checker_error.rs

@@ -281,4 +281,11 @@ create_messages!(
         msg: format!("Unknown annotation: `{annotation}`."),
         help: Some("Use a valid annotation. The Leo compiler supports: `@program`".to_string()),
     }
+
+    @formatted
+    helper_function_inputs_cannot_have_modes {
+        args: (),
+        msg: format!("Helper functions cannot have modes associated with their inputs."),
+        help: Some("Consider removing the mode or adding a `@program` annotation to the function.".to_string()),
+    }
 );