diff --git a/compiler/passes/src/code_generation/generator.rs b/compiler/passes/src/code_generation/generator.rs index de994e20d0..858e10ebea 100644 --- a/compiler/passes/src/code_generation/generator.rs +++ b/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, } } } diff --git a/compiler/passes/src/code_generation/visit_program.rs b/compiler/passes/src/code_generation/visit_program.rs index ed41bfd1c9..0a8d5b9f34 100644 --- a/compiler/passes/src/code_generation/visit_program.rs +++ b/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 diff --git a/compiler/passes/src/code_generation/visit_statements.rs b/compiler/passes/src/code_generation/visit_statements.rs index 9f1ad55b8c..abe52a85ed 100644 --- a/compiler/passes/src/code_generation/visit_statements.rs +++ b/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); diff --git a/compiler/passes/src/code_generation/visit_type.rs b/compiler/passes/src/code_generation/visit_type.rs index 477923961f..4949769967 100644 --- a/compiler/passes/src/code_generation/visit_type.rs +++ b/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.