fix(ssa): Track all local allocations during flattening

[vezenovm]

Description

Problem*

Resolves #6434 (comment)

Summary*

When there is a duplicate store in the same block without a load or call in between we used to remove that store in mem2reg. After #6434 reactivating the following snippet:

                if let Some(last_store) = references.last_stores.get(&address) {
                    self.instructions_to_remove.insert(*last_store);
                }

gives this error:

In #6434, when handing side effects during flattening we switched to doing prev_value = load foo; store (if c then 3 else prev_value) in foo. However, the first store to an allocation has no previous value (foo = allocate; store 3 in foo). So if we replace the first store with a load, the load is not going to be removed by mem2reg as it now does not have a known value. Now that the load is not being removed the allocation will also remain and not be removed by DIE as expected. We then finally error in ACIR gen for having a remaining allocation.

Fix

I brought back the set of local allocations we had prior to #6434. In #6434 we also tried to account for this first store to an allocation case, however, we are currently only looking at stores which immediately follow an allocation instruction.

Additional Context

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: a19cd1b47401aa55003551396abf4bb9708ab2a9, compared to commit: 7d612425728a9d5852ec26a96af629aec7cd0a21

🧾 Summary (10% most significant diffs)

Program	Brillig opcodes (+/-)	%
closures_mut_ref	-7 ✅	-21.88%
import	-16 ✅	-26.67%
tuples	-16 ✅	-29.09%
loop	-21 ✅	-45.65%
sha256_brillig_performance_regression	-1,607 ✅	-49.58%
keccak256	-1,964 ✅	-52.35%
brillig_loop_size_regression	-22 ✅	-52.38%
reference_only_used_as_alias	-135 ✅	-75.42%

---

Full diff report 👇

Program	Brillig opcodes (+/-)	%
u128	2,795 (-1)	-0.04%
poseidonsponge_x5_254	4,265 (-4)	-0.09%
array_dynamic	318 (-1)	-0.31%
sha256_regression	6,544 (-24)	-0.37%
poseidon_bn254_hash_width_3	5,426 (-22)	-0.40%
poseidon_bn254_hash	5,426 (-22)	-0.40%
ram_blowup_regression	956 (-4)	-0.42%
regression_5252	4,632 (-22)	-0.47%
array_dynamic_blackbox_input	1,022 (-5)	-0.49%
sha256_var_padding_regression	4,769 (-24)	-0.50%
regression_4449	725 (-4)	-0.55%
sha256_var_size_regression	1,706 (-11)	-0.64%
array_sort	302 (-2)	-0.66%
6_array	406 (-3)	-0.73%
regression_struct_array_conditional	531 (-4)	-0.75%
ecdsa_secp256k1	896 (-7)	-0.78%
brillig_cow_regression	2,143 (-17)	-0.79%
array_dynamic_nested_blackbox_input	865 (-7)	-0.80%
sha256_var_witness_const_regression	1,233 (-11)	-0.88%
conditional_1	1,183 (-11)	-0.92%
sha256	2,213 (-25)	-1.12%
conditional_regression_short_circuit	1,188 (-14)	-1.16%
6	1,115 (-14)	-1.24%
sha2_byte	2,774 (-40)	-1.42%
uhashmap	13,303 (-225)	-1.66%
regression_capacity_tracker	234 (-4)	-1.68%
strings	881 (-16)	-1.78%
eddsa	10,434 (-253)	-2.37%
slice_dynamic_index	2,576 (-64)	-2.42%
hashmap	20,020 (-505)	-2.46%
regression	712 (-20)	-2.73%
brillig_cow	368 (-12)	-3.16%
brillig_rc_regression_6123	173 (-6)	-3.35%
brillig_pedersen	552 (-26)	-4.50%
pedersen_check	552 (-26)	-4.50%
no_predicates_numeric_generic_poseidon	758 (-36)	-4.53%
fold_numeric_generic_poseidon	758 (-36)	-4.53%
integer_array_indexing	61 (-3)	-4.69%
poseidon2	345 (-18)	-4.96%
bench_2_to_17	338 (-18)	-5.06%
fold_complex_outputs	480 (-26)	-5.14%
3_add	53 (-3)	-5.36%
simple_2d_array	106 (-6)	-5.36%
fold_2_to_17	581 (-36)	-5.83%
signed_arithmetic	252 (-16)	-5.97%
derive	138 (-9)	-6.12%
higher_order_functions	665 (-46)	-6.47%
nested_array_in_slice	1,113 (-87)	-7.25%
1_mul	63 (-5)	-7.35%
regression_6451	48 (-4)	-7.69%
generics	95 (-10)	-9.52%
slice_loop	257 (-28)	-9.82%
array_to_slice	736 (-82)	-10.02%
7_function	535 (-60)	-10.08%
missing_closure_env	24 (-3)	-11.11%
slices	1,766 (-271)	-13.30%
brillig_constant_reference_regression	84 (-15)	-15.15%
references	161 (-42)	-20.69%
closures_mut_ref	25 (-7)	-21.88%
import	44 (-16)	-26.67%
tuples	39 (-16)	-29.09%
loop	25 (-21)	-45.65%
sha256_brillig_performance_regression	1,634 (-1,607)	-49.58%
keccak256	1,788 (-1,964)	-52.35%
brillig_loop_size_regression	20 (-22)	-52.38%
reference_only_used_as_alias	44 (-135)	-75.42%

[github-actions]

Changes to number of Brillig opcodes executed

Generated at commit: a19cd1b47401aa55003551396abf4bb9708ab2a9, compared to commit: 7d612425728a9d5852ec26a96af629aec7cd0a21

🧾 Summary (10% most significant diffs)

Program	Brillig opcodes (+/-)	%
brillig_constant_reference_regression	-18 ✅	-18.18%
slice_loop	-216 ✅	-22.18%
closures_mut_ref	-7 ✅	-25.00%
import	-16 ✅	-29.63%
tuples	-16 ✅	-35.56%
loop	-21 ✅	-50.00%
brillig_loop_size_regression	-22 ✅	-55.00%
reference_only_used_as_alias	-135 ✅	-67.50%

---

Full diff report 👇

Program	Brillig opcodes (+/-)	%
ram_blowup_regression	778,681 (-4)	-0.00%
sha256_var_padding_regression	219,719 (-24)	-0.01%
brillig_cow_regression	519,040 (-63)	-0.01%
sha256_regression	116,180 (-24)	-0.02%
u128	25,430 (-8)	-0.03%
sha256_var_size_regression	16,345 (-11)	-0.07%
ecdsa_secp256k1	6,774 (-7)	-0.10%
regression_4449	199,604 (-280)	-0.14%
array_dynamic_nested_blackbox_input	4,514 (-7)	-0.15%
sha256_var_witness_const_regression	6,741 (-11)	-0.16%
sha256	13,845 (-25)	-0.18%
conditional_1	5,718 (-11)	-0.19%
conditional_regression_short_circuit	7,163 (-14)	-0.20%
6	7,085 (-14)	-0.20%
array_dynamic	497 (-1)	-0.20%
array_dynamic_blackbox_input	18,181 (-72)	-0.39%
6_array	1,667 (-7)	-0.42%
regression_capacity_tracker	908 (-4)	-0.44%
array_sort	582 (-4)	-0.68%
regression_struct_array_conditional	1,577 (-12)	-0.76%
uhashmap	147,011 (-1,133)	-0.76%
eddsa	732,419 (-6,199)	-0.84%
strings	1,734 (-16)	-0.91%
brillig_cow	1,135 (-12)	-1.05%
hashmap	55,892 (-831)	-1.47%
poseidonsponge_x5_254	188,510 (-3,960)	-2.06%
regression_5252	939,100 (-20,030)	-2.09%
poseidon_bn254_hash	167,074 (-3,879)	-2.27%
poseidon_bn254_hash_width_3	167,074 (-3,879)	-2.27%
brillig_pedersen	957 (-26)	-2.64%
pedersen_check	957 (-26)	-2.64%
brillig_rc_regression_6123	292 (-8)	-2.67%
sha2_byte	48,585 (-1,351)	-2.71%
fold_complex_outputs	766 (-26)	-3.28%
fold_numeric_generic_poseidon	5,132 (-175)	-3.30%
no_predicates_numeric_generic_poseidon	5,132 (-175)	-3.30%
7_function	2,496 (-91)	-3.52%
poseidon2	705 (-26)	-3.56%
higher_order_functions	1,242 (-46)	-3.57%
integer_array_indexing	75 (-3)	-3.85%
array_to_slice	1,790 (-82)	-4.38%
bench_2_to_17	589,848 (-27,770)	-4.50%
simple_2d_array	125 (-6)	-4.58%
fold_2_to_17	1,093,915 (-52,820)	-4.61%
keccak256	33,591 (-1,964)	-5.52%
nested_array_in_slice	1,471 (-87)	-5.58%
derive	135 (-9)	-6.25%
slice_dynamic_index	4,374 (-307)	-6.56%
generics	133 (-10)	-6.99%
3_add	39 (-3)	-7.14%
signed_arithmetic	205 (-16)	-7.24%
regression	2,642 (-209)	-7.33%
1_mul	52 (-5)	-8.77%
regression_6451	38 (-4)	-9.52%
slices	2,884 (-402)	-12.23%
sha256_brillig_performance_regression	22,993 (-3,381)	-12.82%
missing_closure_env	20 (-3)	-13.04%
references	252 (-53)	-17.38%
brillig_constant_reference_regression	81 (-18)	-18.18%
slice_loop	758 (-216)	-22.18%
closures_mut_ref	21 (-7)	-25.00%
import	38 (-16)	-29.63%
tuples	29 (-16)	-35.56%
loop	21 (-21)	-50.00%
brillig_loop_size_regression	18 (-22)	-55.00%
reference_only_used_as_alias	65 (-135)	-67.50%

[github-actions]

Changes to circuit sizes

Generated at commit: a19cd1b47401aa55003551396abf4bb9708ab2a9, compared to commit: 7d612425728a9d5852ec26a96af629aec7cd0a21

🧾 Summary (10% most significant diffs)

Program	ACIR opcodes (+/-)	%	Circuit size (+/-)	%
poseidonsponge_x5_254	-294 ✅	-22.04%	-176 ✅	-12.30%
regression_5252	-99,291 ✅	-75.12%	-133,894 ✅	-75.16%

---

Full diff report 👇

Program	ACIR opcodes (+/-)	%	Circuit size (+/-)	%
conditional_1	6,150 (+2)	+0.03%	16,734 (-9)	-0.05%
schnorr	1,471 (-42)	-2.78%	23,864 (-42)	-0.18%
side_effects_constrain_array	5 (-4)	-44.44%	2,764 (-5)	-0.18%
regression	105 (-12)	-10.26%	3,647 (-8)	-0.22%
sha256_var_size_regression	11,441 (-532)	-4.44%	70,436 (-1,434)	-2.00%
slices	738 (-49)	-6.23%	3,838 (-94)	-2.39%
hashmap	24,105 (-1,552)	-6.05%	62,573 (-2,080)	-3.22%
slice_dynamic_index	992 (-269)	-21.33%	6,394 (-513)	-7.43%
references	7 (-3)	-30.00%	23 (-3)	-11.54%
poseidonsponge_x5_254	1,040 (-294)	-22.04%	1,255 (-176)	-12.30%
regression_5252	32,889 (-99,291)	-75.12%	44,252 (-133,894)	-75.16%

[TomAFrench]

This looks very reasonable but it's late so will approve tomorrow after I've had another look.