typecheck that finalize_stub

compiler/passes/src/type_checking/check_program.rs

@@ -14,7 +14,7 @@
 // You should have received a copy of the GNU General Public License
 // along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
 
-use crate::{DiGraphError, TypeChecker, VariableSymbol, VariableType};
+use crate::{DiGraphError, TypeChecker};
 
 use leo_ast::*;
 use leo_errors::TypeCheckerError;
@@ -56,20 +56,9 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
     }
 
     fn visit_stub(&mut self, input: &'a Stub) {
-        // Cannot have mappings in stubs.
-        if input.mappings.len() != 0 {
-            self.emit_err(TypeCheckerError::stubs_can_only_have_records_and_transitions(
-                "mapping",
-                input.mappings.get(0).unwrap().1.span,
-            ));
-        }
-
         // Cannot have constant declarations in stubs.
-        if input.consts.len() != 0 {
+        if !input.consts.is_empty() {
-            self.emit_err(TypeCheckerError::stubs_can_only_have_records_and_transitions(
+            self.emit_err(TypeCheckerError::stubs_cannot_have_const_declarations(input.consts.get(0).unwrap().1.span));
-                "constant declaration",
-                input.consts.get(0).unwrap().1.span,
-            ));
         }
 
         // Typecheck the program's structs.
@@ -80,23 +69,11 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
     }
 
     fn visit_function_stub(&mut self, input: &'a FunctionStub) {
-        // Cannot have finalize scopes
+        // Must not be an inline function
-        if input.finalize.is_some() {
-            self.emit_err(TypeCheckerError::stub_functions_must_have_no_finalize(
-                input.finalize.as_ref().unwrap().span,
-            ));
-        }
-
-        // Must be transition functions
         if input.variant == Variant::Inline {
             self.emit_err(TypeCheckerError::stub_functions_must_not_be_inlines(input.span));
         }
 
-        // Must be empty
-        if !input.block.statements.is_empty() {
-            self.emit_err(TypeCheckerError::stub_functions_must_be_empty(input.block.span));
-        }
-
         // Lookup function metadata in the symbol table.
         // Note that this unwrap is safe since function metadata is stored in a prior pass.
         let function_index = self.symbol_table.borrow().lookup_fn_symbol(input.identifier.name).unwrap().id;
@@ -107,6 +84,19 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
         // Query helper function to type check function parameters and outputs.
         self.check_function_signature(&Function::from(input.clone()));
 
+        // Check that the finalize scope is valid
+        if input.finalize_stub.is_some() {
+            // Create a new child scope for the finalize block.
+            let scope_index = self.create_child_scope();
+
+            // Check the finalize signature.
+            let function = &Function::from(input.clone());
+            self.check_finalize_signature(function.finalize.as_ref().unwrap(), function);
+
+            // Exit the scope for the finalize block.
+            self.exit_scope(scope_index);
+        }
+
         // Exit the function's scope.
         self.exit_scope(function_index);
     }
@@ -114,10 +104,7 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
     fn visit_struct_stub(&mut self, input: &'a Struct) {
         // Allow records only.
         if !input.is_record {
-            self.emit_err(TypeCheckerError::stubs_can_only_have_records_and_transitions(
+            self.emit_err(TypeCheckerError::stubs_cannot_have_non_record_structs(input.span));
-                "non-record struct",
-                input.span,
-            ));
         }
 
         self.visit_struct(input);
@@ -326,90 +313,16 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
             self.is_finalize = true;
             // The function's finalize block does not have a return statement.
             self.has_return = false;
-            // The function;s finalize block does not have a finalize statement.
+            // The function's finalize block does not have a finalize statement.
             self.has_finalize = false;
 
-            // Check that the function is a transition function.
-            if !matches!(function.variant, Variant::Transition) {
-                self.emit_err(TypeCheckerError::only_transition_functions_can_have_finalize(finalize.span));
-            }
-
-            // Check that the name of the finalize block matches the function name.
-            if function.identifier.name != finalize.identifier.name {
-                self.emit_err(TypeCheckerError::finalize_name_mismatch(
-                    function.identifier.name,
-                    finalize.identifier.name,
-                    finalize.span,
-                ));
-            }
-
             // Create a new child scope for the finalize block.
             let scope_index = self.create_child_scope();
 
-            finalize.input.iter().for_each(|input_var| {
+            // Check the finalize signature.
-                // Check that the type of input parameter is defined.
+            self.check_finalize_signature(finalize, function);
-                if self.assert_type_is_valid(&input_var.type_(), input_var.span()) {
-                    // Check that the input parameter is not a tuple.
-                    if matches!(input_var.type_(), Type::Tuple(_)) {
-                        self.emit_err(TypeCheckerError::finalize_cannot_take_tuple_as_input(input_var.span()))
-                    }
-                    // Check that the input parameter is not a record.
-                    if let Type::Identifier(identifier) = input_var.type_() {
-                        // Note that this unwrap is safe, as the type is defined.
-                        if self.symbol_table.borrow().lookup_struct(identifier.name).unwrap().is_record {
-                            self.emit_err(TypeCheckerError::finalize_cannot_take_record_as_input(input_var.span()))
-                        }
-                    }
-                    // Check that the input parameter is not constant or private.
-                    if input_var.mode() == Mode::Constant || input_var.mode() == Mode::Private {
-                        self.emit_err(TypeCheckerError::finalize_input_mode_must_be_public(input_var.span()));
-                    }
-                    // Check for conflicting variable names.
-                    if let Err(err) =
-                        self.symbol_table.borrow_mut().insert_variable(input_var.identifier().name, VariableSymbol {
-                            type_: input_var.type_(),
-                            span: input_var.identifier().span(),
-                            declaration: VariableType::Input(input_var.mode()),
-                        })
-                    {
-                        self.handler.emit_err(err);
-                    }
-                }
-            });
 
-            // Check that the finalize block's return type is a unit type.
+            // TODO: Remove if this restriction is relaxed at Aleo instructions level
-            // Note: This is a temporary restriction to be compatible with the current version of snarkVM.
-            // Note: This restriction may be lifted in the future.
-            // Note: This check is still compatible with the other checks below.
-            if finalize.output_type != Type::Unit {
-                self.emit_err(TypeCheckerError::finalize_cannot_return_value(finalize.span));
-            }
-
-            // Type check the finalize block's return type.
-            // Note that checking that each of the component types are defined is sufficient to guarantee that the `output_type` is defined.
-            finalize.output.iter().for_each(|output_type| {
-                // Check that the type of output is defined.
-                if self.assert_type_is_valid(&output_type.type_(), output_type.span()) {
-                    // Check that the output is not a tuple. This is necessary to forbid nested tuples.
-                    if matches!(&output_type.type_(), Type::Tuple(_)) {
-                        self.emit_err(TypeCheckerError::nested_tuple_type(output_type.span()))
-                    }
-                    // Check that the output is not a record.
-                    if let Type::Identifier(identifier) = output_type.type_() {
-                        // Note that this unwrap is safe, as the type is defined.
-                        if self.symbol_table.borrow().lookup_struct(identifier.name).unwrap().is_record {
-                            self.emit_err(TypeCheckerError::finalize_cannot_output_record(output_type.span()))
-                        }
-                    }
-                    // Check that the mode of the output is valid.
-                    // Note that a finalize block can have only public outputs.
-                    if matches!(output_type.mode(), Mode::Constant | Mode::Private) {
-                        self.emit_err(TypeCheckerError::finalize_output_mode_must_be_public(output_type.span()));
-                    }
-                }
-            });
-
-            // TODO: Remove if this restriction is relaxed at Aleo instructions level.
             // Check that the finalize block is not empty.
             if finalize.block.statements.is_empty() {
                 self.emit_err(TypeCheckerError::finalize_block_must_not_be_empty(finalize.span));
@@ -418,9 +331,6 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
             // Type check the finalize block.
             self.visit_block(&finalize.block);
 
-            // Check that the return type is defined. Note that the component types are already checked.
-            self.assert_type_is_valid(&finalize.output_type, finalize.span);
-
             // If the function has a return type, then check that it has a return.
             if finalize.output_type != Type::Unit && !self.has_return {
                 self.emit_err(TypeCheckerError::missing_return(finalize.span));

compiler/passes/src/type_checking/checker.rs

@@ -19,6 +19,7 @@ use crate::{CallGraph, StructGraph, SymbolTable, TypeTable, VariableSymbol, Vari
 use leo_ast::{
     CoreConstant,
     CoreFunction,
+    Finalize,
     Function,
     Identifier,
     IntegerType,
@@ -1219,6 +1220,88 @@ impl<'a> TypeChecker<'a> {
             }
         });
     }
+
+    pub(crate) fn check_finalize_signature(&mut self, finalize: &Finalize, function: &Function) {
+        // Check that the function is a transition function.
+        if !matches!(function.variant, Variant::Transition) {
+            self.emit_err(TypeCheckerError::only_transition_functions_can_have_finalize(finalize.span));
+        }
+
+        // Check that the name of the finalize block matches the function name.
+        if function.identifier.name != finalize.identifier.name {
+            self.emit_err(TypeCheckerError::finalize_name_mismatch(
+                function.identifier.name,
+                finalize.identifier.name,
+                finalize.span,
+            ));
+        }
+
+        finalize.input.iter().for_each(|input_var| {
+            // Check that the type of input parameter is defined.
+            if self.assert_type_is_valid(&input_var.type_(), input_var.span()) {
+                // Check that the input parameter is not a tuple.
+                if matches!(input_var.type_(), Type::Tuple(_)) {
+                    self.emit_err(TypeCheckerError::finalize_cannot_take_tuple_as_input(input_var.span()))
+                }
+                // Check that the input parameter is not a record.
+                if let Type::Identifier(identifier) = input_var.type_() {
+                    // Note that this unwrap is safe, as the type is defined.
+                    if self.symbol_table.borrow().lookup_struct(identifier.name).unwrap().is_record {
+                        self.emit_err(TypeCheckerError::finalize_cannot_take_record_as_input(input_var.span()))
+                    }
+                }
+                // Check that the input parameter is not constant or private.
+                if input_var.mode() == Mode::Constant || input_var.mode() == Mode::Private {
+                    self.emit_err(TypeCheckerError::finalize_input_mode_must_be_public(input_var.span()));
+                }
+                // Check for conflicting variable names.
+                if let Err(err) =
+                    self.symbol_table.borrow_mut().insert_variable(input_var.identifier().name, VariableSymbol {
+                        type_: input_var.type_(),
+                        span: input_var.identifier().span(),
+                        declaration: VariableType::Input(input_var.mode()),
+                    })
+                {
+                    self.handler.emit_err(err);
+                }
+            }
+        });
+
+        // Check that the finalize block's return type is a unit type.
+        // Note: This is a temporary restriction to be compatible with the current version of snarkVM.
+        // Note: This restriction may be lifted in the future.
+        // Note: This check is still compatible with the other checks below.
+        if finalize.output_type != Type::Unit {
+            self.emit_err(TypeCheckerError::finalize_cannot_return_value(finalize.span));
+        }
+
+        // Type check the finalize block's return type.
+        // Note that checking that each of the component types are defined is sufficient to guarantee that the `output_type` is defined.
+        finalize.output.iter().for_each(|output_type| {
+            // Check that the type of output is defined.
+            if self.assert_type_is_valid(&output_type.type_(), output_type.span()) {
+                // Check that the output is not a tuple. This is necessary to forbid nested tuples.
+                if matches!(&output_type.type_(), Type::Tuple(_)) {
+                    self.emit_err(TypeCheckerError::nested_tuple_type(output_type.span()))
+                }
+                // Check that the output is not a record.
+                if let Type::Identifier(identifier) = output_type.type_() {
+                    // Note that this unwrap is safe, as the type is defined.
+                    if self.symbol_table.borrow().lookup_struct(identifier.name).unwrap().is_record {
+                        self.emit_err(TypeCheckerError::finalize_cannot_output_record(output_type.span()))
+                    }
+                }
+                // Check that the mode of the output is valid.
+                // Note that a finalize block can have only public outputs.
+                if matches!(output_type.mode(), Mode::Constant | Mode::Private) {
+                    self.emit_err(TypeCheckerError::finalize_output_mode_must_be_public(output_type.span()));
+                }
+            }
+        });
+
+        // Check that the return type is defined. Note that the component types are already checked.
+        self.assert_type_is_valid(&finalize.output_type, finalize.span);
+    }
 }
 
 fn types_to_string(types: &[Type]) -> String {