Add new statements to passes

compiler/parser/src/parser/file.rs

@@ -333,14 +333,52 @@ impl ParserContext<'_> {
         let (inputs, ..) = self.parse_paren_comma_list(|p| p.parse_function_parameter().map(Some))?;
 
         // Parse return type.
-        self.expect(&Token::Arrow)?;
+        let output = match self.eat(&Token::Arrow) {
+            false => Type::Unit,
+            true => {
                 self.disallow_circuit_construction = true;
                 let output = self.parse_type()?.0;
                 self.disallow_circuit_construction = false;
+                output
+            }
+        };
 
         // Parse the function body.
         let block = self.parse_block()?;
 
+        // Parse the `finalize` block if it exists.
+        let finalize = match self.eat(&Token::Finalize) {
+            false => None,
+            true => {
+                // Get starting span.
+                let start = self.prev_token.span;
+
+                // Parse parameters.
+                let (input, ..) = self.parse_paren_comma_list(|p| p.parse_function_parameter().map(Some))?;
+
+                // Parse return type.
+                let output = match self.eat(&Token::Arrow) {
+                    false => Type::Unit,
+                    true => {
+                        self.disallow_circuit_construction = true;
+                        let output = self.parse_type()?.0;
+                        self.disallow_circuit_construction = false;
+                        output
+                    }
+                };
+
+                // Parse the finalize body.
+                let block = self.parse_block()?;
+
+                Some(Finalize {
+                    input,
+                    output,
+                    span: start + block.span,
+                    block,
+                })
+            }
+        };
+
         Ok((
             name,
             Function {
@@ -349,7 +387,7 @@ impl ParserContext<'_> {
                 input: inputs,
                 output,
                 span: start + block.span,
-                finalize: None,
+                finalize,
                 block,
             },
         ))

compiler/passes/src/code_generation/visit_statements.rs

@@ -17,8 +17,9 @@
 use crate::CodeGenerator;
 
 use leo_ast::{
-    AssignStatement, Block, ConditionalStatement, ConsoleFunction, ConsoleStatement, DefinitionStatement, Expression,
-    IterationStatement, ParamMode, ReturnStatement, Statement,
+    AssignStatement, Block, ConditionalStatement, ConsoleFunction, ConsoleStatement, DecrementStatement,
+    DefinitionStatement, Expression, FinalizeStatement, IncrementStatement, IterationStatement, ParamMode,
+    ReturnStatement, Statement,
 };
 
 use itertools::Itertools;
@@ -26,13 +27,16 @@ use itertools::Itertools;
 impl<'a> CodeGenerator<'a> {
     fn visit_statement(&mut self, input: &'a Statement) -> String {
         match input {
-            Statement::Return(stmt) => self.visit_return(stmt),
-            Statement::Definition(stmt) => self.visit_definition(stmt),
             Statement::Assign(stmt) => self.visit_assign(stmt),
-            Statement::Conditional(stmt) => self.visit_conditional(stmt),
-            Statement::Iteration(stmt) => self.visit_iteration(stmt),
-            Statement::Console(stmt) => self.visit_console(stmt),
             Statement::Block(stmt) => self.visit_block(stmt),
+            Statement::Conditional(stmt) => self.visit_conditional(stmt),
+            Statement::Console(stmt) => self.visit_console(stmt),
+            Statement::Decrement(stmt) => self.visit_decrement(stmt),
+            Statement::Definition(stmt) => self.visit_definition(stmt),
+            Statement::Finalize(stmt) => self.visit_finalize(stmt),
+            Statement::Increment(stmt) => self.visit_increment(stmt),
+            Statement::Iteration(stmt) => self.visit_iteration(stmt),
+            Statement::Return(stmt) => self.visit_return(stmt),
         }
     }
 
@@ -60,6 +64,18 @@ impl<'a> CodeGenerator<'a> {
         unreachable!("DefinitionStatement's should not exist in SSA form.")
     }
 
+    fn visit_increment(&mut self, _input: &'a IncrementStatement) -> String {
+        todo!()
+    }
+
+    fn visit_decrement(&mut self, _input: &'a DecrementStatement) -> String {
+        todo!()
+    }
+
+    fn visit_finalize(&mut self, _input: &'a FinalizeStatement) -> String {
+        todo!()
+    }
+
     fn visit_assign(&mut self, input: &'a AssignStatement) -> String {
         match &input.place {
             Expression::Identifier(identifier) => {

compiler/passes/src/loop_unrolling/unroll_statement.rs

@@ -20,6 +20,29 @@ use crate::unroller::Unroller;
 use crate::{VariableSymbol, VariableType};
 
 impl StatementReconstructor for Unroller<'_> {
+    fn reconstruct_block(&mut self, input: Block) -> Block {
+        let scope_index = self.current_scope_index();
+
+        // Enter the block scope.
+        self.enter_block_scope(scope_index);
+        self.block_index = 0;
+
+        let block = Block {
+            statements: input
+                .statements
+                .into_iter()
+                .map(|s| self.reconstruct_statement(s))
+                .collect(),
+            span: input.span,
+        };
+
+        // Exit the block scope.
+        self.exit_block_scope(scope_index);
+        self.block_index = scope_index + 1;
+
+        block
+    }
+
     fn reconstruct_definition(&mut self, input: DefinitionStatement) -> Statement {
         // If we are unrolling a loop, then we need to repopulate the symbol table.
         if self.is_unrolling {
@@ -71,27 +94,4 @@ impl StatementReconstructor for Unroller<'_> {
             _ => Statement::Iteration(Box::from(input)),
         }
     }
-
-    fn reconstruct_block(&mut self, input: Block) -> Block {
-        let scope_index = self.current_scope_index();
-
-        // Enter the block scope.
-        self.enter_block_scope(scope_index);
-        self.block_index = 0;
-
-        let block = Block {
-            statements: input
-                .statements
-                .into_iter()
-                .map(|s| self.reconstruct_statement(s))
-                .collect(),
-            span: input.span,
-        };
-
-        // Exit the block scope.
-        self.exit_block_scope(scope_index);
-        self.block_index = scope_index + 1;
-
-        block
-    }
 }

compiler/passes/src/static_single_assignment/rename_statement.rs

@@ -18,8 +18,9 @@ use crate::{RenameTable, StaticSingleAssigner};
 
 use leo_ast::{
     AssignStatement, BinaryExpression, BinaryOperation, Block, ConditionalStatement, ConsoleFunction, ConsoleStatement,
-    DefinitionStatement, Expression, ExpressionConsumer, Identifier, IterationStatement, Node, ReturnStatement,
-    Statement, StatementConsumer, TernaryExpression, UnaryExpression, UnaryOperation,
+    DecrementStatement, DefinitionStatement, Expression, ExpressionConsumer, FinalizeStatement, Identifier,
+    IncrementStatement, IterationStatement, Node, ReturnStatement, Statement, StatementConsumer, TernaryExpression,
+    UnaryExpression, UnaryOperation,
 };
 use leo_span::Symbol;
 
@@ -28,54 +29,6 @@ use indexmap::IndexSet;
 impl StatementConsumer for StaticSingleAssigner<'_> {
     type Output = Vec<Statement>;
 
-    /// Transforms a `ReturnStatement` into an empty `BlockStatement`,
-    /// storing the expression and the associated guard in `self.early_returns`.
-    /// Note that type checking guarantees that there is at most one `ReturnStatement` in a block.
-    fn consume_return(&mut self, input: ReturnStatement) -> Self::Output {
-        // Construct the associated guard.
-        let guard = match self.condition_stack.is_empty() {
-            true => None,
-            false => {
-                let (first, rest) = self.condition_stack.split_first().unwrap();
-                Some(rest.iter().cloned().fold(first.clone(), |acc, condition| {
-                    Expression::Binary(BinaryExpression {
-                        op: BinaryOperation::And,
-                        left: Box::new(acc),
-                        right: Box::new(condition),
-                        span: Default::default(),
-                    })
-                }))
-            }
-        };
-
-        // Consume the expression and add it to `early_returns`.
-        let (expression, statements) = self.consume_expression(input.expression);
-        // Note that this is the only place where `self.early_returns` is mutated.
-        // Furthermore, `expression` will always be an identifier or tuple expression.
-        self.early_returns.push((guard, expression));
-
-        statements
-    }
-
-    /// Consumes the `DefinitionStatement` into an `AssignStatement`, renaming the left-hand-side as appropriate.
-    fn consume_definition(&mut self, definition: DefinitionStatement) -> Self::Output {
-        // First consume the right-hand-side of the definition.
-        let (value, mut statements) = self.consume_expression(definition.value);
-
-        // Then assign a new unique name to the left-hand-side of the definition.
-        // Note that this order is necessary to ensure that the right-hand-side uses the correct name when consuming a complex assignment.
-        self.is_lhs = true;
-        let identifier = match self.consume_identifier(definition.variable_name).0 {
-            Expression::Identifier(identifier) => identifier,
-            _ => unreachable!("`self.consume_identifier` will always return an `Identifier`."),
-        };
-        self.is_lhs = false;
-
-        statements.push(Self::simple_assign_statement(Expression::Identifier(identifier), value));
-
-        statements
-    }
-
     /// Consume all `AssignStatement`s, renaming as necessary.
     fn consume_assign(&mut self, assign: AssignStatement) -> Self::Output {
         // First consume the right-hand-side of the assignment.
@@ -93,6 +46,15 @@ impl StatementConsumer for StaticSingleAssigner<'_> {
         statements
     }
 
+    /// Consumes a `Block`, flattening its constituent `ConditionalStatement`s.
+    fn consume_block(&mut self, block: Block) -> Self::Output {
+        block
+            .statements
+            .into_iter()
+            .flat_map(|statement| self.consume_statement(statement))
+            .collect()
+    }
+
     /// Consumes a `ConditionalStatement`, producing phi functions for variables written in the then-block and otherwise-block.
     /// Furthermore a new `AssignStatement` is introduced for non-trivial expressions in the condition of `ConditionalStatement`s.
     /// For example,
@@ -187,11 +149,6 @@ impl StatementConsumer for StaticSingleAssigner<'_> {
         statements
     }
 
-    // TODO: Error message
-    fn consume_iteration(&mut self, _input: IterationStatement) -> Self::Output {
-        unreachable!("`IterationStatement`s should not be in the AST at this phase of compilation.");
-    }
-
     // TODO: Where do we handle console statements.
     fn consume_console(&mut self, input: ConsoleStatement) -> Self::Output {
         let (function, mut statements) = match input.function {
@@ -230,12 +187,68 @@ impl StatementConsumer for StaticSingleAssigner<'_> {
         statements
     }
 
-    /// Consumes a `Block`, flattening its constituent `ConditionalStatement`s.
-    fn consume_block(&mut self, block: Block) -> Self::Output {
-        block
-            .statements
-            .into_iter()
-            .flat_map(|statement| self.consume_statement(statement))
-            .collect()
+    fn consume_decrement(&mut self, _input: DecrementStatement) -> Self::Output {
+        todo!()
+    }
+
+    /// Consumes the `DefinitionStatement` into an `AssignStatement`, renaming the left-hand-side as appropriate.
+    fn consume_definition(&mut self, definition: DefinitionStatement) -> Self::Output {
+        // First consume the right-hand-side of the definition.
+        let (value, mut statements) = self.consume_expression(definition.value);
+
+        // Then assign a new unique name to the left-hand-side of the definition.
+        // Note that this order is necessary to ensure that the right-hand-side uses the correct name when consuming a complex assignment.
+        self.is_lhs = true;
+        let identifier = match self.consume_identifier(definition.variable_name).0 {
+            Expression::Identifier(identifier) => identifier,
+            _ => unreachable!("`self.consume_identifier` will always return an `Identifier`."),
+        };
+        self.is_lhs = false;
+
+        statements.push(Self::simple_assign_statement(Expression::Identifier(identifier), value));
+
+        statements
+    }
+
+    fn consume_finalize(&mut self, _input: FinalizeStatement) -> Self::Output {
+        todo!()
+    }
+
+    fn consume_increment(&mut self, _input: IncrementStatement) -> Self::Output {
+        todo!()
+    }
+
+    // TODO: Error message
+    fn consume_iteration(&mut self, _input: IterationStatement) -> Self::Output {
+        unreachable!("`IterationStatement`s should not be in the AST at this phase of compilation.");
+    }
+
+    /// Transforms a `ReturnStatement` into an empty `BlockStatement`,
+    /// storing the expression and the associated guard in `self.early_returns`.
+    /// Note that type checking guarantees that there is at most one `ReturnStatement` in a block.
+    fn consume_return(&mut self, input: ReturnStatement) -> Self::Output {
+        // Construct the associated guard.
+        let guard = match self.condition_stack.is_empty() {
+            true => None,
+            false => {
+                let (first, rest) = self.condition_stack.split_first().unwrap();
+                Some(rest.iter().cloned().fold(first.clone(), |acc, condition| {
+                    Expression::Binary(BinaryExpression {
+                        op: BinaryOperation::And,
+                        left: Box::new(acc),
+                        right: Box::new(condition),
+                        span: Default::default(),
+                    })
+                }))
+            }
+        };
+
+        // Consume the expression and add it to `early_returns`.
+        let (expression, statements) = self.consume_expression(input.expression);
+        // Note that this is the only place where `self.early_returns` is mutated.
+        // Furthermore, `expression` will always be an identifier or tuple expression.
+        self.early_returns.push((guard, expression));
+
+        statements
     }
 }

compiler/passes/src/type_checking/check_statements.rs

@@ -29,52 +29,16 @@ impl<'a> StatementVisitor<'a> for TypeChecker<'a> {
         }
 
         match input {
-            Statement::Return(stmt) => self.visit_return(stmt),
-            Statement::Definition(stmt) => self.visit_definition(stmt),
             Statement::Assign(stmt) => self.visit_assign(stmt),
-            Statement::Conditional(stmt) => self.visit_conditional(stmt),
-            Statement::Iteration(stmt) => self.visit_iteration(stmt),
-            Statement::Console(stmt) => self.visit_console(stmt),
             Statement::Block(stmt) => self.visit_block(stmt),
-        }
-    }
-
-    fn visit_return(&mut self, input: &'a ReturnStatement) {
-        // we can safely unwrap all self.parent instances because
-        // statements should always have some parent block
-        let parent = self.parent.unwrap();
-        let return_type = &self
-            .symbol_table
-            .borrow()
-            .lookup_fn_symbol(parent)
-            .map(|f| f.output.clone());
-
-        self.has_return = true;
-
-        self.visit_expression(&input.expression, return_type);
-    }
-
-    fn visit_definition(&mut self, input: &'a DefinitionStatement) {
-        let declaration = if input.declaration_type == DeclarationType::Const {
-            VariableType::Const
-        } else {
-            VariableType::Mut
-        };
-
-        // Check that the type of the definition is valid.
-        self.assert_type_is_valid(input.span, &input.type_);
-
-        self.visit_expression(&input.value, &Some(input.type_.clone()));
-
-        if let Err(err) = self.symbol_table.borrow_mut().insert_variable(
-            input.variable_name.name,
-            VariableSymbol {
-                type_: input.type_.clone(),
-                span: input.span(),
-                declaration,
-            },
-        ) {
-            self.handler.emit_err(err);
+            Statement::Conditional(stmt) => self.visit_conditional(stmt),
+            Statement::Console(stmt) => self.visit_console(stmt),
+            Statement::Decrement(stmt) => self.visit_decrement(stmt),
+            Statement::Definition(stmt) => self.visit_definition(stmt),
+            Statement::Finalize(stmt) => self.visit_finalize(stmt),
+            Statement::Increment(stmt) => self.visit_increment(stmt),
+            Statement::Iteration(stmt) => self.visit_iteration(stmt),
+            Statement::Return(stmt) => self.visit_return(stmt),
         }
     }
 
@@ -108,6 +72,26 @@ impl<'a> StatementVisitor<'a> for TypeChecker<'a> {
         }
     }
 
+    fn visit_block(&mut self, input: &'a Block) {
+        // Creates a new sub-scope since we are in a block.
+        let scope_index = self.symbol_table.borrow_mut().insert_block();
+        let previous_symbol_table = std::mem::take(&mut self.symbol_table);
+        self.symbol_table.swap(
+            previous_symbol_table
+                .borrow()
+                .lookup_scope_by_index(scope_index)
+                .unwrap(),
+        );
+        self.symbol_table.borrow_mut().parent = Some(Box::new(previous_symbol_table.into_inner()));
+
+        input.statements.iter().for_each(|stmt| self.visit_statement(stmt));
+
+        let previous_symbol_table = *self.symbol_table.borrow_mut().parent.take().unwrap();
+        self.symbol_table
+            .swap(previous_symbol_table.lookup_scope_by_index(scope_index).unwrap());
+        self.symbol_table = RefCell::new(previous_symbol_table);
+    }
+
     fn visit_conditional(&mut self, input: &'a ConditionalStatement) {
         self.visit_expression(&input.condition, &Some(Type::Boolean));
 
@@ -136,6 +120,58 @@ impl<'a> StatementVisitor<'a> for TypeChecker<'a> {
         self.has_return = previous_has_return || (then_block_has_return && otherwise_block_has_return);
     }
 
+    fn visit_console(&mut self, input: &'a ConsoleStatement) {
+        match &input.function {
+            ConsoleFunction::Assert(expr) => {
+                let type_ = self.visit_expression(expr, &Some(Type::Boolean));
+                self.assert_bool_type(&type_, expr.span());
+            }
+            ConsoleFunction::AssertEq(left, right) | ConsoleFunction::AssertNeq(left, right) => {
+                let t1 = self.visit_expression(left, &None);
+                let t2 = self.visit_expression(right, &None);
+
+                // Check that the types are equal.
+                self.check_eq_types(&t1, &t2, input.span());
+            }
+        }
+    }
+
+    fn visit_decrement(&mut self, _input: &'a DecrementStatement) {
+        todo!()
+    }
+
+    fn visit_definition(&mut self, input: &'a DefinitionStatement) {
+        let declaration = if input.declaration_type == DeclarationType::Const {
+            VariableType::Const
+        } else {
+            VariableType::Mut
+        };
+
+        // Check that the type of the definition is valid.
+        self.assert_type_is_valid(input.span, &input.type_);
+
+        self.visit_expression(&input.value, &Some(input.type_.clone()));
+
+        if let Err(err) = self.symbol_table.borrow_mut().insert_variable(
+            input.variable_name.name,
+            VariableSymbol {
+                type_: input.type_.clone(),
+                span: input.span(),
+                declaration,
+            },
+        ) {
+            self.handler.emit_err(err);
+        }
+    }
+
+    fn visit_finalize(&mut self, _input: &'a FinalizeStatement) {
+        todo!()
+    }
+
+    fn visit_increment(&mut self, _input: &'a IncrementStatement) {
+        todo!()
+    }
+
     fn visit_iteration(&mut self, input: &'a IterationStatement) {
         let iter_type = &Some(input.type_.clone());
         self.assert_int_type(iter_type, input.variable.span);
@@ -190,39 +226,18 @@ impl<'a> StatementVisitor<'a> for TypeChecker<'a> {
         }
     }
 
-    fn visit_console(&mut self, input: &'a ConsoleStatement) {
-        match &input.function {
-            ConsoleFunction::Assert(expr) => {
-                let type_ = self.visit_expression(expr, &Some(Type::Boolean));
-                self.assert_bool_type(&type_, expr.span());
-            }
-            ConsoleFunction::AssertEq(left, right) | ConsoleFunction::AssertNeq(left, right) => {
-                let t1 = self.visit_expression(left, &None);
-                let t2 = self.visit_expression(right, &None);
-
-                // Check that the types are equal.
-                self.check_eq_types(&t1, &t2, input.span());
-            }
-        }
-    }
-
-    fn visit_block(&mut self, input: &'a Block) {
-        // Creates a new sub-scope since we are in a block.
-        let scope_index = self.symbol_table.borrow_mut().insert_block();
-        let previous_symbol_table = std::mem::take(&mut self.symbol_table);
-        self.symbol_table.swap(
-            previous_symbol_table
+    fn visit_return(&mut self, input: &'a ReturnStatement) {
+        // we can safely unwrap all self.parent instances because
+        // statements should always have some parent block
+        let parent = self.parent.unwrap();
+        let return_type = &self
+            .symbol_table
             .borrow()
-                .lookup_scope_by_index(scope_index)
-                .unwrap(),
-        );
-        self.symbol_table.borrow_mut().parent = Some(Box::new(previous_symbol_table.into_inner()));
+            .lookup_fn_symbol(parent)
+            .map(|f| f.output.clone());
 
-        input.statements.iter().for_each(|stmt| self.visit_statement(stmt));
+        self.has_return = true;
 
-        let previous_symbol_table = *self.symbol_table.borrow_mut().parent.take().unwrap();
-        self.symbol_table
-            .swap(previous_symbol_table.lookup_scope_by_index(scope_index).unwrap());
-        self.symbol_table = RefCell::new(previous_symbol_table);
+        self.visit_expression(&input.expression, return_type);
     }
 }

compiler/span/src/symbol.rs

@@ -186,6 +186,7 @@ symbols! {
     context,
     CoreFunction,
     console,
+    decrement,
     Else: "else",
     finalize,
     For: "for",
@@ -193,6 +194,7 @@ symbols! {
     If: "if",
     In: "in",
     import,
+    increment,
     input,
     Let: "let",
     leo,