feat(experimental): try to infer lambda argument types inside calls

[asterite]

Description

Problem

Fixes #6802

Summary

I've been thinking for days how we could have lambda parameter types be "inferred" from the call they are being passed to.

The first thing that came to my mind is that lambdas are most commonly passed as callbacks after invoking a method on self, where either self or a part of self is given to the lambda (like Option::map, BoundedVec::map, etc.). So the first thing that I tried in this PR is to eagerly unify a method call's function type with the object type. Then, when elaborating a lambda as a method call argument we pass the potential parameter types of a function type that is in that call position:

fn map<U>(&self, f: fn(T) -> U) { ... }
                       ^
                     this

self.map(|x| ...)
          ^
    has to unify with this

Then these types are unified, without erroring because later we'll check the type of the lambda against the argument anyway.

And... that worked! And that already covers a lot of cases.

Then I did the same thing for function calls, except that there's no self, but at least it now works if a callback has a concrete type, like if it's:

fn foo(f: fn(Foo) -> ...

And that worked too! Though I'm not sure there are many uses of that...

BUT: it didn't work in the code Nico shared in Slack, because it's a function call where the first argument given is like a self type, except that it's not a method all:

for_each_in_bounded_vec(notes, |note, _| {
})

So the final thing I did was to eagerly try to unify argument types as we elaborate them against the target function type. And that made that example work!

It won't work if the lambda comes before the argument (which works in Rust) but I think that pattern is uncommon (though we could try to make it work in the future).

Additional Context

I don't know if this is the right way to approach this.

I also don't know if unifying eagerly would cause any issues. One thing that's not done here is using "unify_with_coercions", but given that we don't issue the errors that happens in these eager checks, maybe it's fine (maybe it won't work in cases where an array is automatically converted to a slice, though I guess we could make it work in the future).

One more thing: the changes in the stdlib and programs aren't really necessary, but I wanted to see if the code compiles with those changes... and in many cases the code is simplified a bit.

And finally: the code is not the best as I was just experimenting. We should clean it up.

Documentation

Check one:

• No documentation needed.
• Documentation included in this PR.
• [For Experimental Features] Documentation to be submitted in a separate PR.

PR Checklist

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

[github-actions]

Changes to Brillig bytecode sizes

Generated at commit: 4a6e10312457646be325bb564e0c5f983418761d, compared to commit: a6685befe3b103b236997843f69afc07b6ea0a8c

🧾 Summary (10% most significant diffs)

Program	Brillig opcodes (+/-)	%
higher_order_functions_inliner_zero	+26 ❌	+3.90%
slices_inliner_min	-18 ✅	-0.68%

---

Full diff report 👇

Program	Brillig opcodes (+/-)	%
higher_order_functions_inliner_zero	692 (+26)	+3.90%
higher_order_functions_inliner_min	1,448 (+4)	+0.28%
slices_inliner_min	2,615 (-18)	-0.68%

[github-actions]

Changes to number of Brillig opcodes executed

Generated at commit: 4a6e10312457646be325bb564e0c5f983418761d, compared to commit: a6685befe3b103b236997843f69afc07b6ea0a8c

🧾 Summary (10% most significant diffs)

Program	Brillig opcodes (+/-)	%
higher_order_functions_inliner_min	-165 ✅	-5.86%
higher_order_functions_inliner_zero	-90 ✅	-6.79%

---

Full diff report 👇

Program	Brillig opcodes (+/-)	%
slices_inliner_min	4,549 (-75)	-1.62%
higher_order_functions_inliner_min	2,651 (-165)	-5.86%
higher_order_functions_inliner_zero	1,236 (-90)	-6.79%

[github-actions]

Compilation Report

Program	Compilation Time	%
sha256_regression	1.050s	5%
regression_4709	0.826s	4%
ram_blowup_regression	16.500s	3%
rollup-root	3.572s	-6%
rollup-merge	2.038s	-7%
rollup-block-root-single-tx	150.000s	7%
rollup-block-root-empty	2.086s	-4%
rollup-block-root	139.000s	-1%
rollup-block-merge	3.728s	1%
rollup-base-public	29.180s	3%
rollup-base-private	10.002s	0%
private-kernel-tail	0.947s	-4%
private-kernel-reset	6.020s	2%
private-kernel-inner	1.952s	-7%

[github-actions]

Execution Report

Program	Execution Time	%
sha256_regression	0.052s	1%
regression_4709	0.001s	0%
ram_blowup_regression	0.619s	3%
rollup-root	0.104s	0%
rollup-merge	0.006s	0%
rollup-block-root	36.800s	1%
rollup-block-merge	0.104s	-1%
rollup-base-public	1.218s	0%
rollup-base-private	0.449s	-1%
private-kernel-tail	0.019s	0%
private-kernel-reset	0.311s	0%
private-kernel-inner	0.068s	0%

[asterite]

I didn't expect this to have any repercussion in SSA 😮

Lambdas that involve math operations somehow compile before this change without type annotations:

let myarray: [i32; 3] = [1, 2, 3];

// Compiles fine without saying `n: i32`
assert(myarray.any(|n| n > 2));

I think it's because of the >.

But in the end n ends up being i32 so I don't understand how it could change... I'll investigate it later today. Though I guess if it ends up in more optimizations that's good, though maybe it's suspicious and this introduced a bug... 🤔

[github-actions]

Execution Memory Report

Program	Peak Memory	%
keccak256	74.63M	0%
workspace	123.72M	0%
regression_4709	315.93M	-1%
ram_blowup_regression	512.57M	0%
rollup-root	498.20M	0%
rollup-merge	472.94M	-1%
rollup-block-root	1.22G	0%
rollup-block-merge	498.21M	0%
rollup-base-public	734.09M	0%
rollup-base-private	590.40M	0%
private-kernel-tail	180.82M	0%
private-kernel-reset	245.44M	0%
private-kernel-inner	208.84M	0%

[github-actions]

Compilation Memory Report

Program	Peak Memory
keccak256	77.56M
workspace	123.76M
regression_4709	424.07M
ram_blowup_regression	1.46G
rollup-root	601.17M
rollup-merge	494.15M
rollup-block-root-single-tx	16.06G
rollup-block-root-empty	488.32M
rollup-block-root	16.07G
rollup-block-merge	601.16M
rollup-base-public	2.38G
rollup-base-private	1.14G
private-kernel-tail	207.34M
private-kernel-reset	584.40M
private-kernel-inner	294.60M

[asterite]

Also compilation memory and execution memory went up 🤔

Now it's back to normal. Actually the numbers are the same, it's just that maybe the base numbers were outdated...

We could undo commits 79f4076 and 2cffe19 if we wanted, though they are probably good checks to improve performance.

[asterite]

It seems with this PR some functions get inlined right from the beginning, while in master some are not. I'm not sure why... but if this only happens for higher-order functions (maybe uncommon?) and if it leads to an optimization, maybe it's good.

[jfecher]

Instead of unifying afterward we can push the type down and that should be sufficient. This is what Rust does as well since it uses a different algorithm called "bidirectional type inference" which splits up type checking functions into fn infer(T) -> Type which we have, and fn check(T, Type) which we don't have. Certain constructs like literals are always inferred, and others like lambdas are usually checked.

For our purposes though we can just use your existing check arg type function but pass down the expected type instead of the entire function type and argument index.

[asterite]

Instead of unifying afterward we can push the type down and that should be sufficient

Some weeks ago I tried something similar to what I did here but I can't remember what (it didn't work so I deleted the branch). I thought I was assigning types instead of unifying, and getting errors like "can't find method foo for T" and that's why I tried unifying here. But it's probably the case that I had a bug or something else and that's why it wasn't working. I'll try pushing the type down and using it when the type is unspecified 👍

[asterite]

Hm, now we get a failure on code like this:

struct U60Repr<let N: u32, let NumSegments: u32> {}

impl<let N: u32, let NumSegments: u32> U60Repr<N, NumSegments> {
    fn new<let NumFieldSegments: u32>(_: [Field; N * NumFieldSegments]) -> Self {
        U60Repr {}
    }
}

fn main() {
    let input: [Field; 6] = [0; 6];
    let _: U60Repr<3, 6> = U60Repr::new(input);
}

The error is this:

error: Expected type [Field; ((6 / _) * _)], found type [Field; 6]

let u60repr: U60Repr<3, 6> = unsafe { U60Repr::new(input) };
                                                   -----

Looking at the types that are compares, these are the ones before this PR:

arg = [Field; (6: numeric u32)]
param = [Field; (Numeric(Shared(RefCell { value: Unbound('8, Numeric(u32)) }): u32) * Numeric(Shared(RefCell { value: Unbound('10, Numeric(u32)) }): u32))]

And these are the ones in this PR:

arg = [Field; (6: numeric u32)]
param = [Field; ((6 / _) * Numeric(Shared(RefCell { value: Unbound('10, Numeric(u32)) }): u32))]

It seems the compiler is trying to solve the equation and that's why there's a division. I don't know if this is a bug in the unification code or if it's a bug introduced in this PR.

[asterite]

So what I found is happening is...

When checking arg against param, eventually it'll try the math substitution here:

https://github.com/noir-lang/noir/blob/c172880ae47ec4906cda662801bd4b7866c9586b/compiler/noirc_frontend/src/hir_def/types.rs#L1729C47-L1736

Code:

// Handle cases like `4 = a + b` by trying to solve to `a = 4 - b`
let new_type = InfixExpr(
    Box::new(Constant(*value, kind.clone())),
    inverse,
    rhs.clone(),
);
new_type.try_unify(lhs, bindings)?;
Ok(())

That will eventually fail because of an occurs check in Type::try_bind_to:

// Check if the target id occurs within `this` before binding. Otherwise this could
// cause infinitely recursive types
if this.occurs(target_id) {
    Err(UnificationError)
} else {
    bindings.insert(target_id, (var.clone(), this.kind(), this.clone()));
    Ok(())
}

Maybe it's failing because we are binding these types multiple times now?

Something that caught my attention is the ? followed by Ok(()) in the first snippet:

new_type.try_unify(lhs, bindings)?;
Ok(())

That could just be:

new_type.try_unify(lhs, bindings)

I wonder if the intention of that code was trying to solve the math equation, but always returning Ok(()) if we couldn't unify it, but it ended up accidentally returning error because of the ?.

If in the occurs check I return Ok(()) instead of Err(UnificationError) the issue is also solved and no test fails, though of course I don't know if that's okay to do.

@jfecher Thoughts?

[jfecher]

@asterite we can't remove the occurs check, that'd break the check for infinite types like a = (a, a). I think the error is stemming from our inability to solve 6 = ? * ? because N and NumSegments are not also being pushed down, since they're only known from the let statement. My guess is before this change on arguments we did ?a = ?b * ?c instead and could just set the value of ?a to equal ?b * ?c in that case. Then when they're constrained further later it'd be all fine.

[asterite]

I see, thanks.

I think the error is stemming from our inability to solve 6 = ? * ?

In the code before this PR, [Field; 6] was checked against [Field; ?a * ?b] and that somehow didn't seem to break. Maybe it's that it solved ?a to be 6 / ?b? I didn't check.

I think now it's checking 6 = (6 / ?a) * ?b, maybe we don't have that case covered. But that equation could be turned into 6 / ?b = 6 / ?a so ?a and ?b should be the same... not sure! I jumped into working on the loop statement, but I'll come back to this after that.

[jfecher]

// Handle cases like 4 = a + b by trying to solve to a = 4 - b

I don't see anything wrong with that case aside from the fact the last Ok(()) is unnecessary as you mentioned. I don't think it's the cause of the issue here at least.

I think now it's checking 6 = (6 / ?a) * ?b, maybe we don't have that case covered.

This looks incorrect to me since as you mentioned it'd require a = b which shouldn't follow from the original equation of 6 = a * b. So the error probably lies in how that constraint comes about