Fix visibilities in codegen

compiler/passes/src/code_generation/generator.rs

@@ -16,7 +16,7 @@
 
 use leo_ast::Function;
 use leo_errors::emitter::Handler;
-use leo_span::{sym, Symbol};
+use leo_span::Symbol;
 
 use indexmap::IndexMap;
 
@@ -35,6 +35,8 @@ pub struct CodeGenerator<'a> {
     /// 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,
+    /// Are we traversing a finalize block?
+    pub(crate) in_finalize: bool,
 }
 
 impl<'a> CodeGenerator<'a> {
@@ -48,6 +50,7 @@ impl<'a> CodeGenerator<'a> {
             variable_mapping: IndexMap::new(),
             composite_mapping: IndexMap::new(),
             is_program_function: false,
+            in_finalize: false,
         }
     }
 }

compiler/passes/src/code_generation/visit_program.rs

@@ -16,7 +16,7 @@
 
 use crate::CodeGenerator;
 
-use leo_ast::{Circuit, CircuitMember, Function, Identifier, Mapping, Program, Type};
+use leo_ast::{Circuit, CircuitMember, Function, Identifier, Mapping, Mode, Program, Type};
 
 use indexmap::IndexMap;
 use itertools::Itertools;
@@ -188,7 +188,11 @@ impl<'a> CodeGenerator<'a> {
             self.variable_mapping
                 .insert(&input.identifier.name, register_string.clone());
 
-            let type_string = self.visit_type_with_visibility(&input.type_, input.mode);
+            let visibility = match (self.is_program_function, input.mode) {
+                (true, Mode::None) => Mode::Private,
+                _ => input.mode,
+            };
+            let type_string = self.visit_type_with_visibility(&input.type_, visibility);
             writeln!(function_string, "    input {} as {};", register_string, type_string,)
                 .expect("failed to write to string");
         }
@@ -201,6 +205,7 @@ impl<'a> CodeGenerator<'a> {
         if let Some(finalize) = &function.finalize {
             // Clear the register count.
             self.next_register = 0;
+            self.in_finalize = true;
 
             // Clear the variable mapping.
             // TODO: Figure out a better way to initialize.
@@ -217,13 +222,20 @@ impl<'a> CodeGenerator<'a> {
                 self.variable_mapping
                     .insert(&input.identifier.name, register_string.clone());
 
-                let type_string = self.visit_type_with_visibility(&input.type_, input.mode);
+                // A finalize block defaults to public visibility.
+                let visibility = match (self.is_program_function, input.mode) {
+                    (true, Mode::None) => Mode::Public,
+                    _ => unreachable!("Only program functions can have finalize blocks."),
+                };
+                let type_string = self.visit_type_with_visibility(&input.type_, visibility);
                 writeln!(function_string, "    input {} as {};", register_string, type_string,)
                     .expect("failed to write to string");
             }
 
             // Construct and append the finalize block body.
             function_string.push_str(&self.visit_block(&finalize.block));
+
+            self.in_finalize = false;
         }
 
         function_string

compiler/passes/src/code_generation/visit_statements.rs

@@ -42,17 +42,40 @@ impl<'a> CodeGenerator<'a> {
     }
 
     fn visit_return(&mut self, input: &'a ReturnStatement) -> String {
-        match &self.current_function.unwrap().output_type {
-            Type::Unit => String::new(),
-            output_type => {
+        match input.expression {
+            // Skip empty return statements.
+            Expression::Tuple(ref tuple) if tuple.elements.is_empty() => String::new(),
+            _ => {
                 let (operand, mut expression_instructions) = self.visit_expression(&input.expression);
-                // 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(output_type, Mode::Private);
                 let instructions = operand
                     .split('\n')
                     .into_iter()
-                    .zip(types.iter())
-                    .map(|(operand, type_)| format!("    output {} as {};\n", operand, type_))
+                    .zip(self.current_function.unwrap().output.iter())
+                    .enumerate()
+                    .map(|(i, (operand, output))| {
+                        let visibility = if self.is_program_function {
+                            match self.in_finalize {
+                                // If in finalize block, the default visibility is public.
+                                true => match output.mode {
+                                    Mode::None => Mode::Public,
+                                    mode => mode,
+                                },
+                                // If not in finalize block, the default visibility is private.
+                                false => match output.mode {
+                                    Mode::None => Mode::Private,
+                                    mode => mode,
+                                },
+                            }
+                        } else {
+                            // Only program functions have visibilities associated with their outputs.
+                            Mode::None
+                        };
+                        format!(
+                            "    output {} as {};\n",
+                            operand,
+                            self.visit_type_with_visibility(&output.type_, visibility)
+                        )
+                    })
                     .join("");
 
                 expression_instructions.push_str(&instructions);

compiler/passes/src/code_generation/visit_type.rs

@@ -17,7 +17,6 @@
 use crate::CodeGenerator;
 
 use leo_ast::{Mode, Type};
-use std::fmt::Write as _;
 
 impl<'a> CodeGenerator<'a> {
     fn visit_type(&mut self, input: &'a Type) -> String {
@@ -29,13 +28,7 @@ impl<'a> CodeGenerator<'a> {
             | Type::Scalar
             | Type::String
             | Type::Integer(..) => format!("{}", input),
-            Type::Identifier(ident) => {
-                if let Some((_, type_)) = self.composite_mapping.get(&ident.name) {
-                    format!("{}.{}", ident, type_)
-                } else {
-                    unreachable!("All composite types should be known at this phase of compilation")
-                }
-            }
+            Type::Identifier(ident) => format!("{}", ident),
             Type::Mapping(_) => {
                 unreachable!("Mapping types are not supported at this phase of compilation")
             }
@@ -47,25 +40,18 @@ impl<'a> CodeGenerator<'a> {
         }
     }
 
-    pub(crate) fn visit_type_with_visibility(&mut self, input: &'a Type, visibility: Mode) -> String {
-        let mut type_string = self.visit_type(input);
-
-        if let Type::Identifier(_) = input {
-            // Do not append anything for record and circuit types.
-        } else {
-            // todo: CAUTION private by default.
-            // 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 {
-                    Mode::None => Mode::Private,
-                    _ => visibility,
-                };
-                write!(type_string, ".{}", visibility).expect("failed to write to string");
+    pub(crate) fn visit_type_with_visibility(&mut self, type_: &'a Type, visibility: Mode) -> String {
+        match type_ {
+            // When the type is a record.
+            // Note that this unwrap is safe because all composite types have been added to the mapping.
+            Type::Identifier(identifier) if self.composite_mapping.get(&identifier.name).unwrap().0 => {
+                format!("{}.record", identifier)
             }
+            _ => match visibility {
+                Mode::None => self.visit_type(type_),
+                _ => format!("{}.{}", self.visit_type(type_), visibility),
+            },
         }
-
-        type_string
     }
 
     /// Returns one or more types equal to the number of return tuple members.