fix(defunctionalize): Higher order functions (HOF) dynamic dispatch and HOFs in arrays

[vezenovm]

Description

Problem*

Working towards #5503 but not that issue.

The case this fixes has been reduced to the following:

fn main(x: u32) {
    println(lambdas_with_input_and_return_values(x));
}
fn lambdas_with_input_and_return_values(x: u32) -> u32 {
    let f1 = |value: u32| value;
    let f2 = |value: u32| value + 1;
    let f3 = |value: u32| value + 2;
    let f = if x == 0 {
        f1
    } else if x == 1 {
        f2
    } else {
        f3
    };
    f(x)
}

This is the following SSA before defunctionalization:

Details

acir(inline) fn main f0 {
  b0(v0: u32):
    v2 = call f2(v0) -> u32
    call f1(v2)
    return
}
acir(inline) fn println f1 {
  b0(v0: u32):
    call f6(u1 1, v0)
    return
}
acir(inline) fn lambdas_with_input_and_return_values f2 {
  b0(v0: u32):
    v4 = eq v0, u32 0
    jmpif v4 then: b1, else: b2
  b1():
    jmp b3(f3)
  b2():
    v6 = eq v0, u32 1
    jmpif v6 then: b4, else: b5
  b3(v1: function):
    v10 = call v1(v0) -> u32
    return v10
  b4():
    jmp b6(f4)
  b5():
    jmp b6(f5)
  b6(v2: function):
    jmp b3(v2)
}
acir(inline) fn lambda f3 {
  b0(v0: u32):
    return v0
}
acir(inline) fn lambda f4 {
  b0(v0: u32):
    v2 = add v0, u32 1
    return v2
}
acir(inline) fn lambda f5 {
  b0(v0: u32):
    v2 = add v0, u32 2
    return v2
}
brillig(inline) fn print_unconstrained f6 {
  b0(v0: u1, v1: u32):
    v21 = make_array b"{\"kind\":\"unsignedinteger\",\"width\":32}"
    call print(v0, v1, v21, u1 0)
    return
}

On master that gives the following apply function:

Details

acir(inline_always) fn apply f7 {
  b0(v0: Field, v1: u32):
    v4 = eq v0, Field 3
    jmpif v4 then: b3, else: b2
  b1(v2: u32):
    return v2
  b2():
    v8 = eq v0, Field 4
    jmpif v8 then: b5, else: b4
  b3():
    v6 = call f3(v1) -> u32
    jmp b1(v6)
  b4():
    constrain v0 == Field 5
    v13 = call f5(v1) -> u32
    jmp b1(v13)
  b5():
    v10 = call f4(v1) -> u32
    jmp b1(v10)
}

This logic panics during branch analysis:

The application panicked (crashed).
Message:  internal error: entered unreachable code: return encountered before a join point was found. This can only happen if early-return was added to the language without implementing it by jmping to a join block first
Location: compiler/noirc_evaluator/src/ssa/opt/flatten_cfg/branch_analysis.rs:106

Follow-ups: #8674 then #8677

Summary*

@jfecher had fixed this problem for match statements here #7570. I basically copied the logic in the PR to here. I wanted to make a utility in the function builder for this code, but the one during ssa_gen uses the Tree<Value> struct that is restricted in the ssa_gen with pub(super). I felt it was simpler to duplicate code in this case. If the reviewers disagree I can work on a shared utility here. Although in general this is a workaround to handle flattening restrictionx and we can remove this code once flattening is fixed to not required to merge into one block.

This is the new apply function for the snippet above:

Details

acir(inline_always) fn apply f7 {
  b0(v0: Field, v1: u32):
    v9 = eq v0, Field 3
    jmpif v9 then: b3, else: b2
  b1(v2: u32):
    return v2
  b2():
    v13 = eq v0, Field 4
    jmpif v13 then: b6, else: b5
  b3():
    v11 = call f3(v1) -> u32
    jmp b4(v11)
  b4(v3: u32):
    jmp b10(v3)
  b5():
    constrain v0 == Field 5
    v18 = call f5(v1) -> u32
    jmp b8(v18)
  b6():
    v15 = call f4(v1) -> u32
    jmp b7(v15)
  b7(v4: u32):
    jmp b9(v4)
  b8(v5: u32):
    jmp b9(v5)
  b9(v6: u32):
    jmp b10(v6)
  b10(v7: u32):
    jmp b1(v7)
}

This is currently still a draft as I generate this larger apply function for both ACIR/Brillig runtimes. I want to see the affect on Brillig functions and if it is significant I will gate this change by runtime.
EDIT: It looks to affect the Brillig byte code by at most ~1.5%. We can gate this in the future if it starts to cause larger regressions.

Additional Context

Do we want to gate this change to ACIR runtimes? #8672 (comment)

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: baa120804ad736c206186d4ea935beba6a1eb38b, compared to commit: b57c80e1c819b49cb88e3b948baea9c9f2e0e9ab

🧾 Summary (10% most significant diffs)

Program	Brillig opcodes (+/-)	%
higher_order_functions_inliner_min	+12 ❌	+1.08%
higher_order_functions_inliner_zero	+4 ❌	+0.62%

---

Full diff report 👇

Program	Brillig opcodes (+/-)	%
higher_order_functions_inliner_min	1,122 (+12)	+1.08%
higher_order_functions_inliner_zero	650 (+4)	+0.62%
hashmap_inliner_min	8,789 (+40)	+0.46%
uhashmap_inliner_zero	6,870 (+24)	+0.35%
uhashmap_inliner_min	7,285 (+24)	+0.33%
hashmap_inliner_zero	7,724 (+8)	+0.10%

[github-actions]

Changes to number of Brillig opcodes executed

Generated at commit: baa120804ad736c206186d4ea935beba6a1eb38b, compared to commit: b57c80e1c819b49cb88e3b948baea9c9f2e0e9ab

🧾 Summary (10% most significant diffs)

Program	Brillig opcodes (+/-)	%
higher_order_functions_inliner_min	+26 ❌	+1.46%
higher_order_functions_inliner_zero	+12 ❌	+1.04%

---

Full diff report 👇

Program	Brillig opcodes (+/-)	%
higher_order_functions_inliner_min	1,812 (+26)	+1.46%
higher_order_functions_inliner_zero	1,163 (+12)	+1.04%
hashmap_inliner_min	73,340 (+108)	+0.15%
hashmap_inliner_zero	66,783 (+60)	+0.09%
uhashmap_inliner_zero	168,780 (+108)	+0.06%
uhashmap_inliner_min	176,694 (+108)	+0.06%

[vezenovm]

Looks to be some minor regressions for using the new apply function vs. the one with a common return in Brillig. Do we think it is worth it to gate this change to ACIR runtimes?

[github-actions]

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark 'Test Suite Duration'.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 1.20.

Benchmark suite	Current: 8b4cd72	Previous: eb3eef9	Ratio
test_report_zkpassport_noir_rsa_	3 s	2 s	1.50

This comment was automatically generated by workflow using github-action-benchmark.

CC: @TomAFrench

[github-actions]

ACVM Benchmarks

Details

Benchmark suite	Current: 8274228	Previous: b57c80e	Ratio
purely_sequential_opcodes	253900 ns/iter (± 329)	267550 ns/iter (± 2599)	0.95
perfectly_parallel_opcodes	222416 ns/iter (± 475)	236248 ns/iter (± 545)	0.94
perfectly_parallel_batch_inversion_opcodes	3209303 ns/iter (± 22109)	3217535 ns/iter (± 59094)	1.00

This comment was automatically generated by workflow using github-action-benchmark.

[github-actions]

Execution Time

Details

Benchmark suite	Current: 8274228	Previous: b57c80e	Ratio
private-kernel-inner	0.028 s	0.028 s	1
private-kernel-reset	0.162 s	0.161 s	1.01
private-kernel-tail	0.011 s	0.011 s	1
rollup-base-private	0.312 s	0.31 s	1.01
rollup-base-public	0.198 s	0.205 s	0.97
rollup-block-root	11.6 s	11.2 s	1.04
rollup-merge	0.004 s	0.004 s	1
rollup-root	0.009 s	0.009 s	1
semaphore-depth-10	0.023 s	0.019 s	1.21

This comment was automatically generated by workflow using github-action-benchmark.

[github-actions]

Compilation Time

Details

Benchmark suite	Current: 8274228	Previous: b57c80e	Ratio
private-kernel-inner	2.394 s	2.47 s	0.97
private-kernel-reset	6.536 s	6.526 s	1.00
private-kernel-tail	1.154 s	1.058 s	1.09
rollup-base-private	16.26 s	16.3 s	1.00
rollup-base-public	14.74 s	14.6 s	1.01
rollup-block-root-empty	1.376 s	1.384 s	0.99
rollup-block-root-single-tx	126 s	126 s	1
rollup-block-root	127 s	126 s	1.01
rollup-merge	1.168 s	1.158 s	1.01
rollup-root	1.698 s	1.63 s	1.04
semaphore-depth-10	0.843 s	0.829 s	1.02

This comment was automatically generated by workflow using github-action-benchmark.

[github-actions]

Artifact Size

Details

Benchmark suite	Current: 8274228	Previous: b57c80e	Ratio
private-kernel-inner	1132.7 KB	1132.6 KB	1.00
private-kernel-reset	2051.1 KB	2051 KB	1.00
private-kernel-tail	586.6 KB	586.4 KB	1.00
rollup-base-private	5129.2 KB	5129.1 KB	1.00
rollup-base-public	4070.6 KB	4070.5 KB	1.00
rollup-block-root-empty	256.7 KB	256.7 KB	1
rollup-block-root-single-tx	25697.8 KB	25697.8 KB	1
rollup-block-root	25721.5 KB	25721.5 KB	1
rollup-merge	181.6 KB	181.6 KB	1
rollup-root	416.5 KB	416.5 KB	1
semaphore-depth-10	636.3 KB	636.3 KB	1

This comment was automatically generated by workflow using github-action-benchmark.

[github-actions]

Test Suite Duration

Details

Benchmark suite	Current: 8274228	Previous: b57c80e	Ratio
test_report_AztecProtocol_aztec-packages_noir-projects_aztec-nr	65 s	67 s	0.97
test_report_AztecProtocol_aztec-packages_noir-projects_noir-contracts	97 s	99 s	0.98
test_report_AztecProtocol_aztec-packages_noir-projects_noir-protocol-circuits_crates_blob	45 s	44 s	1.02
test_report_AztecProtocol_aztec-packages_noir-projects_noir-protocol-circuits_crates_private-kernel-lib	189 s	195 s	0.97
test_report_AztecProtocol_aztec-packages_noir-projects_noir-protocol-circuits_crates_rollup-lib	184 s	197 s	0.93
test_report_AztecProtocol_aztec-packages_noir-projects_noir-protocol-circuits_crates_types	59 s	60 s	0.98
test_report_noir-lang_noir-bignum_	370 s	378 s	0.98
test_report_noir-lang_noir_bigcurve_	254 s	229 s	1.11
test_report_noir-lang_sha512_	32 s	30 s	1.07
test_report_zkpassport_noir_rsa_	2 s	3 s	0.67

This comment was automatically generated by workflow using github-action-benchmark.

[github-actions]

Compilation Memory

Details

Benchmark suite	Current: 8274228	Previous: b57c80e	Ratio
private-kernel-inner	333.01 MB	333.13 MB	1.00
private-kernel-reset	572.38 MB	572.37 MB	1.00
private-kernel-tail	232.88 MB	232.89 MB	1.00
rollup-base-private	1470 MB	1470 MB	1
rollup-base-public	1630 MB	1630 MB	1
rollup-block-root-empty	429.53 MB	429.52 MB	1.00
rollup-block-root-single-tx	7240 MB	7240 MB	1
rollup-block-root	7250 MB	7240 MB	1.00
rollup-merge	413.8 MB	413.81 MB	1.00
rollup-root	468.07 MB	468.03 MB	1.00
semaphore_depth_10	128.37 MB	128.37 MB	1

This comment was automatically generated by workflow using github-action-benchmark.

[github-actions]

Execution Memory

Details

Benchmark suite	Current: 8274228	Previous: b57c80e	Ratio
private-kernel-inner	242.05 MB	242.04 MB	1.00
private-kernel-reset	264.09 MB	264.09 MB	1
private-kernel-tail	217.23 MB	217.23 MB	1
rollup-base-private	577.89 MB	577.89 MB	1
rollup-base-public	574.21 MB	574.21 MB	1
rollup-block-root	1460 MB	1460 MB	1
rollup-merge	399.73 MB	399.73 MB	1
rollup-root	405.29 MB	405.29 MB	1
semaphore_depth_10	92.93 MB	92.93 MB	1

This comment was automatically generated by workflow using github-action-benchmark.

[jfecher]

Looks to be some minor regressions for using the new apply function vs. the one with a common return in Brillig. Do we think it is worth it to gate this change to ACIR runtimes?

For a maximum regression of ~1.5% I think the bug fix is more important than gating it

[vezenovm]

For a maximum regression of ~1.5% I think the bug fix is more important than gating it

Yeah that is what I was thinking. If larger regressions are revealed in the future we can gate it.

[jfecher]

I just merged the 2nd PR into this one first - should be faster than merging both into master separately.

[vezenovm]

I just merged the 2nd PR into this one first - should be faster than merging both into master separately.

Agreed, I missed your 2nd approval.

[github-actions]

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark 'Execution Time'.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 1.20.

Benchmark suite	Current: 8274228	Previous: b57c80e	Ratio
semaphore-depth-10	0.023 s	0.019 s	1.21

This comment was automatically generated by workflow using github-action-benchmark.

CC: @TomAFrench