feat: Initial work on rewriting closures to regular functions with hi…

[alehander92]

Description

This commit implements the following mechanism:

On a line where a lambda expression is encountered, we initialize a tuple for the captured lambda environment and we rewrite the lambda to a regular function taking this environment as an additional parameter. All calls to the closure are then modified to insert this hidden parameter.

In other words, the following code:

let x = some_value;
let closure = |a| x + a;

println(closure(10));
println(closure(20));

is rewritten to:

fn closure(env: (Field,), a: Field) -> Field {
  env.0 + a
}

let x = some_value;
let closure_env = (x,);

println(closure(closure_env, 10));
println(closure(closure_env, 20));

In the presence of nested closures, we propagate the captured variables implicitly through all intermediate closures:

let x = some_value;

# The hidden env here will include `x` even through
# it's not directly used by the outer closure:
let closure = |a, c|
  # here, `x` is initialized from the hidden env of the outer closure
  let inner_closure = |b|
    a + b + x

  inner_closure(c)

To make these transforms possible, the following changes were made to the logic of the HIR resolver and the monomorphization pass:

• In the HIR resolver pass, the code determines the precise list of variables captured by each lambda. Along with the list, we compute the index of each captured var within the parent closure's environment (when the capture is propagated).

• Introduction of a new Closure type in order to be able to recognize the call-sites that need the automatic environment variable treatment. It's a bit unfortunate that the Closure type is defined within the AST modules that are used to describe the output of the monomorphization pass, because we aim to eliminate all closures during the pass. A better solution would have been possible if the type check pass after HIR resolution was outputting types specific to the HIR pass (then the closures would exist only within this separate non-simplified type system).

• The majority of the work is in the Lambda processing step in the monomorphizer which performs the necessary transformations based on the above information.

Remaining things to do:

• There are a number of pending TODO items for various minor unresolved loose ends in the code.

• There are a lot of possible additional tests to be written.

• Update docs

Problem*

Resolves #1254

Documentation

• This PR requires documentation updates when merged.

  • I will submit a noir-lang/docs PR.
  • I will request for and support Dev Rel's help in documenting this PR.

PR Checklist*

• I have tested the changes locally.
• I have formatted the changes with Prettier and/or cargo fmt on default settings.
  : cargo fmt

[jfecher]

I've yet to go through the whole PR so far, but here's an initial review of mostly the type system related parts. Tomorrow I'll revisit this and go through the name resolution and monomorphization changes.

[kevaundray]

@alehander92 what is the status of this PR?

[alehander92]

@kevaundray we discussed some lang/impl questions that appeared while working on this PR with @jfecher through the last week, and I prepared some fixes for this PR and the initial closure conversion.(also rebased after upstream dropped the old ssa backend) (More advanced capabilities can be merged as a subsequent patch probably?): I'll try to push the current changes soon

[alehander92]

I've just pushed the rebased version of the original changes, soon I'll push some more fixes

[alehander92]

I refactored, so we can use only the Function variant for closures as well, we included some tests from the team, fixed some problems and tried to address most comments from the review. I found a regression in one case of returning a closure, however my understanding is advanced returning closure/correct array of closure capabilities can be added in a next MR

[jfecher]

I found a regression in one case of returning a closure, however my understanding is advanced returning closure/correct array of closure capabilities can be added in a next MR

Does this mean we can't return closures at all? I also noticed the current tests mostly test free functions rather than closures with an environment. What cases are expected to be implemented by this PR?

[jfecher]

Looks like this PR crashes for the input:

fn main() {
    let x: u32 = 32;

    let f = if x > 2 {
        || x
    } else {
        || x + 2342
    };

    dep::std::println(f());
}

Also, when trying:

use dep::std::println;

fn main() {
    let x: u32 = 32;

    let f = if x > 2 {
        || x
    } else {
        three
    };

    println(f());
}

fn three() -> u32 {
    3
}

I get the rather confusing error message:

error: Expected type fn((u32)) -> u32 [(u32)], found type fn() -> u32 [()]
   ┌─ /Users/jakefecher/Code/Noir/noir/short/src/main.nr:6:13
   │  
 6 │       let f = if x > 2 {
   │ ╭─────────────'
 7 │ │         || x
 8 │ │     } else {
 9 │ │         three
10 │ │     };
   │ ╰─────'
   │  
   = Expected the types of both if branches to be equal

Some suggestions to improve this:

1. It is confusing for users that the closure type includes its environment as a parameter. I think these should be kept separate in the type system. We should't need to prepend the environment parameter until monomorphization.
2. The enviornment type of the free function is displayed as [()], we should probably just omit this entirely.
3. It's probably fine to display the closure environment as [(u32)] since there isn't much prior art here. Users who aren't familiar with the internals of closures here may be confused though. Maybe we'd benefit from being more explicit/verbose about it? E.g. Expected type fn() -> u32 with closure environment (u32), found type fn() -> u32

[jfecher]

Does this mean we can't return closures at all? I also noticed the current tests mostly test free functions rather than closures with an environment. What cases are expected to be implemented by this PR?

I tested just now and it seems we're able to return lambdas but not closures that capture an environment. I'd like if that functionality were part of this PR as well. At the same time, I don't want to step on your toes and make things more difficult since I'm aware keeping up with merge conflicts gets increasingly difficult the larger the PR. So let me know if this would be too much. I'll try to be more consistent in reviewing these PRs in a more timely manner as well.

[alehander92]

..
Some suggestions to improve this:

1. It is confusing for users that the closure type includes its environment as a parameter. I think these should be kept separate in the type system. We should't need to prepend the environment parameter until monomorphization.

2. The enviornment type of the free function is displayed as `[()]`, we should probably just omit this entirely.

3. It's probably fine to display the closure environment as `[(u32)]` since there isn't much prior art here. Users who aren't familiar with the internals of closures here may be confused though. Maybe we'd benefit from being more explicit/verbose about it? E.g. `Expected type fn() -> u32 with closure environment (u32), found type fn() -> u32`

I agree with 2 and 3: i've fixed those in the last commit

for 1: i've fixed the unify/display logic, so it should behave correctly in type checking, however part of the logic itself remains before the monomorphization pass: should I move this as well?

now, with those fixes, both examples fail with the same kind of ICE, I'll try to look into it.

Overally, it would be good to fix the returning of closures, but in my understanding it's deeply related to the implicit genericness of functions with closures and instantiating them, which we planned for the next MR, so i think it would be easier to separate them like that(if we remain with the 2 MR plan)

[alehander92]

@jfecher thanks for the help and summary: sounds good.

sorry, had to fix one test again: now they all pass locally (at least cargo test --release --worskpace), I hope they pass the CI as well this time.

I also removed the : Vec<Type> params annotation and rebased: i guess because of the new commits you have to re-approve

[alehander92]

I would think if we're converting closures into tuples, then Self::convert_type should return a tuple then for closures instead of a function type. Perhaps this can be fixed in a later PR as well.

hm , we can try that in the new work

[Savio-Sou]

Did this go into v0.10.0? It doesn't seem to show up in the changelogs 🤔

Separately anything / what should go into the docs?

[alexvitkov]

There's few more things related to this PR that we'd like to get merged first:

1. This syntax change: feat: Syntax for environment types now works with generics #2383
2. Change the signatures of some standard library functions, so they work with closures Noir Closures Implementation Progress #2385 (comment)

Those are both small changes so it should be done today/tomorrow, after that I'll send a PR to docs describing the changes related to closures & give some examples on how to use them.

[alexvitkov]

The changes in Noir have been merged, I've submitted a PR for docs: noir-lang/docs#356