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

[asterite]

Coming back to this, I found that it's trying to unify a Constant against an InfixExpr with these values:

# Infix
lhs = Numeric(Shared(RefCell { value: Unbound('10, Numeric(u32)) }): u32)
op = Multiplication
rhs = ((6: numeric u32) / Numeric(Shared(RefCell { value: Unbound('10, Numeric(u32)) }): u32))

# Constant
value = 6

So it's unifying 6 against (6 / x) * x. And effectively the x is the same in both cases.

What ends up happening, because of

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

is trying to unify x against (6 / (6 / x)), and that's where the error happens.

Now I'm trying to see whre that (6 / x) * x is coming from, because I think that could be just 6.

[jfecher]

The 6 / x * x is definitely a result of trying to solve the equation previously. It's unfortunate we can't just optimize that to 6 anymore like we used to before we considered integer division.

I wonder if we added a version of division that tells the compiler to ignore rounding if that'd fix this. E.g. assuming the original is x = a * b and the translation is x / b = a then we already know x should be divisible by b so we should be able to optimize a theoretical x / b * b later on.

[asterite]

It's unfortunate we can't just optimize that to 6 anymore like we used to before we considered integer division.

Yeah... that's what I tried at first but found that a test about that failed.

I wonder where that (6 / x) * x came from, because in the code there's only N * NumFieldSegments. Like, here there aren't two type variables that are the same that came from the user, it's the compiler duplicating them and putting them in the equation...

I tried gating Type::InfixExpr(..) behind a method to see where's that coming from, but all I can see it's coming from canonicalize which already takes an InfixExpr like that, and I have no idea why I can't find where it all starts 😕

Some commits ago this PR worked, I guess because we didn't type-check as much as we do now (previously we'd only do it if there were lambdas). We could go back to that version, though I guess there would still be code that failed because of this if it was similar to the snippet that fails but also has a lambda argument in it (maybe uncommon). That is, I wonder if this bugs exists regardless of this feature and we could procede with this feature by only applying these changes if there are lambda arguments.

[asterite]

That code is currently erroring on type_check_call because func_type changed after we unified its types with the argument types.

I wanted to try to call type_check_call on a fresh func_type without all the type information we added from arguments... but I didn't know how. I cloned func_type before type-checking the arguments and passed but it didn't work.

[asterite]

I think I found a way to fix this, in a way that might be the "correct" way to do it.

[asterite]

@jfecher I pushed a commit with a fix that feels intuitive to me but I'm not totally sure it's correct.

Essentially, at one point the compiler is trying to unify 6 = a * b so it translates that to a = 6 / b. Later on it tries to unify 6 = b * (6 / b) and it fails because it can't simplify b * (6 / b) to b (as we know that's not always possible).

However, that 6 / b doesn't exist in code, it came from transforming 6 = a * b to a = 6 / b. So I changed the compiler to keep track of this fact (it's carried in InfixExpr).

Then, when the InfixExpr b * (6 / b) is created, if we know 6 / b was created as an "inversion" of another expression, we can be sure that b * (6 / b) is 6, because there was never a division in the first place (this is the part I'm slightly hesitant about, but feels intuitively correct to me).

[jfecher]

@asterite yep, that's the solution I was hinting at with:

I wonder if we added a version of division that tells the compiler to ignore rounding if that'd fix this. E.g. assuming the original is x = a * b and the translation is x / b = a then we already know x should be divisible by b so we should be able to optimize a theoretical x / b * b later on.

I'm not sure if we need it for more than just division but it should do the same thing either way so whichever approach is cleaner.

[asterite]

that's the solution I was hinting at with

Oh, cool! When I read it I thought it was a new language operator, not something we keep track internally.

I'm not sure if we need it for more than just division but it should do the same thing either way so whichever approach is cleaner.

Ah, right. I guess if it's 6 / (b * 6) that it can be simplifies as nothing is lost by first multiplying the dividing.

I'll try introducing a new division operator instead and see how the code looks like.

[asterite]

I started implementing it but I'm not sure... there's inverse() and approx_inverse(). For Multiplication in approx_inverse() now we have two operations that could undo that: both types of divisions. So we can't check anymore if an op is the inverse of just another one, we'd have to check all of them, or maybe have methods is_approx_inverse() and is_inverse().

Then, simplifying away stuff like 6 + (b - 6) into b instead of creating an infix expr for that might be good anyway, so I think I prefer the current version of the code.

[jfecher]

PR looks good.

While reviewing it I was wondering if it'd be cleaner to represent the inverse flag as an additional division operator instead like my earlier comment. I noticed a few things would be different:

• The unification check checks that op_a == op_b so we'd return false there if we ever had a a / b unified against the inversed division a // b. It seems like it'd be better to me if they were treated as equal? Not sure if this'd ever lead to any hard to find bugs with arithmetic types.
• We don't check the inverse flag on Type::eq or hash. Again probably fine since this is an internal flag only used for simplification of these types.
• Notably existing checks in inverse could return true that // is an inverse of * for canonicalizations that care about this.

And another point unrelated to the above

• This isn't needed for this PR (so I won't block on it) but I think the check to eliminate inversed types should be along with the other checks in canonicalize. The reason for this is that types can change when they're unified so what was a // b may become (b * c) // b which we'd want to simplify. This scenario may not be possible currently though since I think we only create these inversed operators during unify when we know we have constant = a <op> b rather than constant = a where a could later unify to an infix expression itself.

[jfecher]

I started implementing it but I'm not sure... there's inverse() and approx_inverse(). For Multiplication in approx_inverse() now we have two operations that could undo that: both types of divisions. So we can't check anymore if an op is the inverse of just another one, we'd have to check all of them, or maybe have methods is_approx_inverse() and is_inverse().

Yeah, we'd probably have to change the method so it accepts an operator and returns a boolean instead so it could match on multiple. The main concern was the existing canonicalizations that call inverse won't see an inversed division as the inverted form of a multiplication still. This is very unrelated to the original PR on lambda parameter inference though...

[asterite]

@jfecher Would you prefer me to try the new division operator in this PR, or leave it to a follow-up PR? I'm fine either way but, as you said, it's unrelated to the original PR so I don't know if this approach is good-enough for now and we can later improve it.

[jfecher]

@asterite I'm not sure it'd be an improvement or not given my past reasoning. We can leave it to a later PR