feat(perf): Remove unused loads in mem2reg and last stores per function

[vezenovm]

Description

Problem*

Partially resolves #4535.

This does not necessarily make SSA "aware of loops" but handles the situation that we in the issue where we have multiple known stores and loads that we are unable to resolve across the function due to loops.

Summary*

The per function state and the logic run over it has been expanded. We store more information in last_loads, a new load_results map to track unused loads, and a utility map calls_reference_input to avoid deleting stores that are directly passed into an entry point with a mutable reference.

To count the remaining stores, I reference count both the results of the loads and the address of the loads. If the results of a load has a counter equal to 0 I remove that load instruction. I then keep track of each removed address with its own reference counter. If the removed loads counter is equal to the counter of the address of the loads, we know we have removed all loads to an address and thus can safely remove the store to that address.

We also added a check that store we want to remove is not used as an argument into a function call.

I have added a unit test remove_unused_loads_and_stores inside of mem2reg.rs.

The brillig_loop_size_regression test:

unconstrained fn main() -> pub Field {
    let mut emulated_enum = EnumEmulation { a: Option::some(1), b: Option::none(), c: Option::none() };

    for _ in 0..1 {
        assert_eq(emulated_enum.a.unwrap(), 1);
    }

    emulated_enum.a = Option::some(2);
    emulated_enum.a.unwrap()
}

now compiles to the optimal SSA after a single run of mem2reg:

After Mem2Reg:
brillig fn main f0 {
  b0():
    v1 = allocate
    v2 = allocate
    v3 = allocate
    v4 = allocate
    v5 = allocate
    v6 = allocate
    jmp b1(u32 0)
  b1(v0: u32):
    v8 = eq v0, u32 0
    jmpif v8 then: b3, else: b2
  b3():
    v11 = add v0, u32 1
    jmp b1(v11)
  b2():
    return Field 2
}

Additional Context

There is most likely other ways we can utilize this per function state. We can continue to iterate upon it in follow-ups.

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: 3866bc79ba1fd03e10672eebfa2eb1a58c4006ed, compared to commit: 052c4fe52a4df9d6492f9b0d6b449151b87b18d5

🧾 Summary (10% most significant diffs)

Program	Brillig opcodes (+/-)	%
poseidon2	-13 ✅	-4.32%
brillig_loop_size_regression	-35 ✅	-58.33%

---

Full diff report 👇

Program	Brillig opcodes (+/-)	%
references	295 (+8)	+2.79%
u128	3,356 (-4)	-0.12%
hashmap	30,846 (-105)	-0.34%
uhashmap	25,831 (-115)	-0.44%
fold_complex_outputs	803 (-5)	-0.62%
no_predicates_numeric_generic_poseidon	708 (-12)	-1.67%
fold_numeric_generic_poseidon	708 (-12)	-1.67%
fold_2_to_17	541 (-12)	-2.17%
bench_2_to_17	282 (-12)	-4.08%
poseidon2	288 (-13)	-4.32%
brillig_loop_size_regression	25 (-35)	-58.33%

[TomAFrench]

currently reviewing this.

[TomAFrench]

LGTM, some very minor nits

[jfecher]

Temporarily blocking just to give myself time to review

[vezenovm]

Can you update the doc comments for this pass to keep them up to date with how it works now?

I added some doc comments. Let me know if anything is unclear.