diff --git a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs
index f4faffcb7ae..0e097a012ec 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs
@@ -567,8 +567,10 @@ impl Loop {
let mut unroll_into = self.get_pre_header(function, cfg)?;
let mut jump_value = get_induction_variable(&function.dfg, unroll_into)?;
- while let Some(context) = self.unroll_header(function, unroll_into, jump_value)? {
- (unroll_into, jump_value) = context.unroll_loop_iteration();
+ while let Some((context, loop_header_id)) =
+ self.unroll_header(function, unroll_into, jump_value)?
+ {
+ (unroll_into, jump_value) = context.unroll_loop_iteration(loop_header_id);
}
Ok(())
@@ -600,20 +602,21 @@ impl Loop {
/// Unrolls the header block of the loop. This is the block that dominates all other blocks in the
/// loop and contains the jmpif instruction that lets us know if we should continue looping.
- /// Returns Some(iteration context) if we should perform another iteration.
+ /// Returns Some((iteration context, loop_header_id)) if we should perform another iteration.
fn unroll_header<'a>(
&'a self,
function: &'a mut Function,
unroll_into: BasicBlockId,
induction_value: ValueId,
- ) -> Result