Skip to content

Commit

Permalink
chore: remove original return from aztec fns (#4804)
Browse files Browse the repository at this point in the history
The following function:

```rust
#[aztec(private)]
fn do_work_conditionally() -> u64 {
  let condition = storage.flag;
  if condition {
     2
  } else {
     42
  }
}
```

Will be transformed by the macro, adding a push with the whole if/else
expression to `context.return_values`...but without removing the
statement in the first place, leading to duplicate code execution.

This PR forces the removal of the last function statement in case it is
replaced by a push to the context.
  • Loading branch information
Thunkar authored Feb 28, 2024
1 parent efd70f4 commit 9e246c1
Showing 1 changed file with 6 additions and 4 deletions.
10 changes: 6 additions & 4 deletions noir/aztec_macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -648,6 +648,10 @@ fn transform_function(

// Abstract return types such that they get added to the kernel's return_values
if let Some(return_values) = abstract_return_values(func) {
// In case we are pushing return values to the context, we remove the statement that originated it
// This avoids running duplicate code, since blocks like if/else can be value returning statements
func.def.body.0.pop();
// Add the new return statement
func.def.body.0.push(return_values);
}

Expand Down Expand Up @@ -1255,15 +1259,13 @@ fn create_avm_context() -> Result<Statement, AztecMacroError> {
/// Any primitive type that can be cast will be casted to a field and pushed to the context.
fn abstract_return_values(func: &NoirFunction) -> Option<Statement> {
let current_return_type = func.return_type().typ;
let len = func.def.body.len();
let last_statement = &func.def.body.0[len - 1];
let last_statement = func.def.body.0.last();

// TODO: (length, type) => We can limit the size of the array returned to be limited by kernel size
// Doesn't need done until we have settled on a kernel size
// TODO: support tuples here and in inputs -> convert into an issue

// Check if the return type is an expression, if it is, we can handle it
match last_statement {
match last_statement? {
Statement { kind: StatementKind::Expression(expression), .. } => {
match current_return_type {
// Call serialize on structs, push the whole array, calling push_array
Expand Down

0 comments on commit 9e246c1

Please sign in to comment.