Implement bug fix

This commit is contained in:
Pranav Gaddamadugu 2022-08-02 23:50:08 -07:00
parent 3a9b7c276b
commit 5be72596db
2 changed files with 42 additions and 24 deletions

View File

@ -15,10 +15,11 @@
// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
use crate::StaticSingleAssigner;
use itertools::Itertools;
use leo_ast::{
Expression, Function, FunctionInput, ProgramReconstructor, ReturnStatement, Statement, StatementReconstructor,
TernaryExpression,
TernaryExpression, TupleExpression,
};
impl ProgramReconstructor for StaticSingleAssigner<'_> {
@ -48,20 +49,40 @@ impl ProgramReconstructor for StaticSingleAssigner<'_> {
// Type checking guarantees that there exists at least one return statement in the function body.
let (_, last_return_expression) = returns.pop().unwrap();
// TODO: Document handling tuples
// Fold all return expressions into a single ternary expression.
let expression =
returns
.into_iter()
.rev()
.fold(last_return_expression, |acc, (guard, expression)| match guard {
None => unreachable!("All return statements except for the last one must have a guard."),
Some(guard) => Expression::Ternary(TernaryExpression {
let expression = returns
.into_iter()
.rev()
.fold(last_return_expression, |acc, (guard, expr)| match guard {
None => unreachable!("All return statements except for the last one must have a guard."),
Some(guard) => match (acc, expr) {
(Expression::Tuple(acc_tuple), Expression::Tuple(expr_tuple)) => {
Expression::Tuple(TupleExpression {
elements: acc_tuple
.elements
.into_iter()
.zip_eq(expr_tuple.elements.into_iter())
.map(|(if_true, if_false)| {
Expression::Ternary(TernaryExpression {
condition: Box::new(guard.clone()),
if_true: Box::new(if_true),
if_false: Box::new(if_false),
span: Default::default(),
})
})
.collect(),
span: Default::default(),
})
}
(acc, expr) => Expression::Ternary(TernaryExpression {
condition: Box::new(guard),
if_true: Box::new(expression),
if_false: Box::new(acc),
if_true: Box::new(acc),
if_false: Box::new(expr),
span: Default::default(),
}),
});
},
});
// Add the `ReturnStatement` to the end of the block.
block.statements.push(Statement::Return(ReturnStatement {

View File

@ -26,19 +26,12 @@ use leo_span::Symbol;
use indexmap::IndexSet;
impl StatementReconstructor for StaticSingleAssigner<'_> {
/// Transforms a `ReturnStatement` into an `AssignStatement`,
/// storing the variable and the associated guard in `self.early_returns`.
/// Transforms a `ReturnStatement` into an empty `BlockStatement`,
/// storing the expression and the associated guard in `self.early_returns`.
///
/// Note that this pass assumes that there is at most one `ReturnStatement` in a block.
/// Note that type checking guarantees that there is at most one `ReturnStatement` in a block.
fn reconstruct_return(&mut self, input: ReturnStatement) -> Statement {
// Create a fresh name for the expression in the return statement.
let symbol = self.unique_symbol("$return");
self.rename_table.update(symbol, symbol);
// Initialize a new `AssignStatement` for the return expression.
let place = Expression::Identifier(Identifier::new(symbol));
// Add the variable and associated guard.
// Construct the associated guard.
let guard = match self.condition_stack.is_empty() {
true => None,
false => {
@ -53,9 +46,13 @@ impl StatementReconstructor for StaticSingleAssigner<'_> {
}))
}
};
self.early_returns.push((guard, place.clone()));
Self::simple_assign_statement(place, self.reconstruct_expression(input.expression).0)
// Reconstruct the expression and add it to the early returns.
let expression = self.reconstruct_expression(input.expression).0;
self.early_returns.push((guard, expression));
// Return an empty block.
Statement::dummy(input.span)
}
/// Reconstructs the `DefinitionStatement` into an `AssignStatement`, renaming the left-hand-side as appropriate.