@@ -13,7 +13,7 @@ use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet};
1313use crate :: ssa:: {
1414 ir:: {
1515 basic_block:: BasicBlockId ,
16- function:: { Function , RuntimeType } ,
16+ function:: Function ,
1717 function_inserter:: FunctionInserter ,
1818 instruction:: { Instruction , InstructionId } ,
1919 types:: Type ,
@@ -27,12 +27,7 @@ use super::unrolling::{Loop, Loops};
2727impl Ssa {
2828 #[ tracing:: instrument( level = "trace" , skip( self ) ) ]
2929 pub ( crate ) fn loop_invariant_code_motion ( mut self ) -> Ssa {
30- let brillig_functions = self
31- . functions
32- . iter_mut ( )
33- . filter ( |( _, func) | matches ! ( func. runtime( ) , RuntimeType :: Brillig ( _) ) ) ;
34-
35- for ( _, function) in brillig_functions {
30+ for function in self . functions . values_mut ( ) {
3631 function. loop_invariant_code_motion ( ) ;
3732 }
3833
@@ -63,6 +58,7 @@ impl Loops {
6358 }
6459
6560 context. map_dependent_instructions ( ) ;
61+ context. inserter . map_data_bus_in_place ( ) ;
6662 }
6763}
6864
@@ -113,6 +109,22 @@ impl<'f> LoopInvariantContext<'f> {
113109
114110 if hoist_invariant {
115111 self . inserter . push_instruction ( instruction_id, pre_header) ;
112+
113+ // If we are hoisting a MakeArray instruction,
114+ // we need to issue an extra inc_rc in case they are mutated afterward.
115+ if matches ! (
116+ self . inserter. function. dfg[ instruction_id] ,
117+ Instruction :: MakeArray { .. }
118+ ) {
119+ let result =
120+ self . inserter . function . dfg . instruction_results ( instruction_id) [ 0 ] ;
121+ let inc_rc = Instruction :: IncrementRc { value : result } ;
122+ let call_stack = self . inserter . function . dfg . get_call_stack ( instruction_id) ;
123+ self . inserter
124+ . function
125+ . dfg
126+ . insert_instruction_and_results ( inc_rc, * block, None , call_stack) ;
127+ }
116128 } else {
117129 self . inserter . push_instruction ( instruction_id, * block) ;
118130 }
@@ -190,6 +202,7 @@ impl<'f> LoopInvariantContext<'f> {
190202 } ) ;
191203
192204 let can_be_deduplicated = instruction. can_be_deduplicated ( self . inserter . function , false )
205+ || matches ! ( instruction, Instruction :: MakeArray { .. } )
193206 || self . can_be_deduplicated_from_upper_bound ( & instruction) ;
194207
195208 is_loop_invariant && can_be_deduplicated
@@ -559,4 +572,94 @@ mod test {
559572 let ssa = ssa. loop_invariant_code_motion ( ) ;
560573 assert_normalized_ssa_equals ( ssa, expected) ;
561574 }
575+
576+ #[ test]
577+ fn insert_inc_rc_when_moving_make_array ( ) {
578+ // SSA for the following program:
579+ //
580+ // unconstrained fn main(x: u32, y: u32) {
581+ // let mut a1 = [1, 2, 3, 4, 5];
582+ // a1[x] = 64;
583+ // for i in 0 .. 5 {
584+ // let mut a2 = [1, 2, 3, 4, 5];
585+ // a2[y + i] = 128;
586+ // foo(a2);
587+ // }
588+ // foo(a1);
589+ // }
590+ //
591+ // We want to make sure move a loop invariant make_array instruction,
592+ // to account for whether that array has been marked as mutable.
593+ // To do so, we increment the reference counter on the array we are moving.
594+ // In the SSA below, we want to move `v42` out of the loop.
595+ let src = "
596+ brillig(inline) fn main f0 {
597+ b0(v0: u32, v1: u32):
598+ v8 = make_array [Field 1, Field 2, Field 3, Field 4, Field 5] : [Field; 5]
599+ v9 = allocate -> &mut [Field; 5]
600+ v11 = array_set v8, index v0, value Field 64
601+ v13 = add v0, u32 1
602+ store v11 at v9
603+ jmp b1(u32 0)
604+ b1(v2: u32):
605+ v16 = lt v2, u32 5
606+ jmpif v16 then: b3, else: b2
607+ b2():
608+ v17 = load v9 -> [Field; 5]
609+ call f1(v17)
610+ return
611+ b3():
612+ v19 = make_array [Field 1, Field 2, Field 3, Field 4, Field 5] : [Field; 5]
613+ v20 = allocate -> &mut [Field; 5]
614+ v21 = add v1, v2
615+ v23 = array_set v19, index v21, value Field 128
616+ call f1(v23)
617+ v25 = add v2, u32 1
618+ jmp b1(v25)
619+ }
620+ brillig(inline) fn foo f1 {
621+ b0(v0: [Field; 5]):
622+ return
623+ }
624+ " ;
625+
626+ let ssa = Ssa :: from_str ( src) . unwrap ( ) ;
627+
628+ // We expect the `make_array` at the top of `b3` to be replaced with an `inc_rc`
629+ // of the newly hoisted `make_array` at the end of `b0`.
630+ let expected = "
631+ brillig(inline) fn main f0 {
632+ b0(v0: u32, v1: u32):
633+ v8 = make_array [Field 1, Field 2, Field 3, Field 4, Field 5] : [Field; 5]
634+ v9 = allocate -> &mut [Field; 5]
635+ v11 = array_set v8, index v0, value Field 64
636+ v13 = add v0, u32 1
637+ store v11 at v9
638+ v14 = make_array [Field 1, Field 2, Field 3, Field 4, Field 5] : [Field; 5]
639+ jmp b1(u32 0)
640+ b1(v2: u32):
641+ v17 = lt v2, u32 5
642+ jmpif v17 then: b3, else: b2
643+ b2():
644+ v18 = load v9 -> [Field; 5]
645+ call f1(v18)
646+ return
647+ b3():
648+ inc_rc v14
649+ v20 = allocate -> &mut [Field; 5]
650+ v21 = add v1, v2
651+ v23 = array_set v14, index v21, value Field 128
652+ call f1(v23)
653+ v25 = add v2, u32 1
654+ jmp b1(v25)
655+ }
656+ brillig(inline) fn foo f1 {
657+ b0(v0: [Field; 5]):
658+ return
659+ }
660+ " ;
661+
662+ let ssa = ssa. loop_invariant_code_motion ( ) ;
663+ assert_normalized_ssa_equals ( ssa, expected) ;
664+ }
562665}
0 commit comments