From e03c8e1e794f1b0c2362aa5ab37a1bda937f1378 Mon Sep 17 00:00:00 2001
From: ilyalesokhin-starkware <ilya@starkware.co>
Date: Thu, 13 Feb 2025 09:20:21 +0200
Subject: [PATCH] Avoid merge in loop functions. (#7262)

---
 crates/cairo-lang-lowering/src/lower/mod.rs   | 29 +++++++++++++++----
 .../src/lower/test_data/for                   |  8 ++---
 .../src/lower/test_data/loop                  |  8 ++---
 .../src/lower/test_data/while                 |  8 ++---
 4 files changed, 35 insertions(+), 18 deletions(-)

diff --git a/crates/cairo-lang-lowering/src/lower/mod.rs b/crates/cairo-lang-lowering/src/lower/mod.rs
index 93b43deb9f6..03b8828e0c5 100644
--- a/crates/cairo-lang-lowering/src/lower/mod.rs
+++ b/crates/cairo-lang-lowering/src/lower/mod.rs
@@ -241,10 +241,12 @@ pub fn lower_for_loop(
         ty: unit_ty,
         location: ctx.get_location(some_block.stable_ptr.untyped()),
     });
-    let sealed_none = lowered_expr_to_block_scope_end(
+    let none_subscope_block_id = none_subscope.block_id;
+    let sealed_none = lower_early_return(
         ctx,
-        none_subscope.clone(),
-        Ok(LoweredExpr::Tuple { exprs: vec![], location: for_location }),
+        none_subscope,
+        LoweredExpr::Tuple { exprs: vec![], location: for_location },
+        for_location,
     )
     .map_err(LoweringFlowError::Failed)?;
 
@@ -259,7 +261,7 @@ pub fn lower_for_loop(
             },
             MatchArm {
                 arm_selector: MatchArmSelector::VariantId(none_variant),
-                block_id: none_subscope.block_id,
+                block_id: none_subscope_block_id,
                 var_ids: vec![none_var_id],
             },
         ],
@@ -323,10 +325,11 @@ pub fn lower_while_loop(
     let subscope_else = create_subscope(ctx, builder);
     let block_else_id = subscope_else.block_id;
     let else_block_input_var_id = ctx.new_var(VarRequest { ty: unit_ty, location: while_location });
-    let block_else = lowered_expr_to_block_scope_end(
+    let block_else = lower_early_return(
         ctx,
         subscope_else,
-        Ok(LoweredExpr::Tuple { exprs: vec![], location: while_location }),
+        LoweredExpr::Tuple { exprs: vec![], location: while_location },
+        while_location,
     )
     .map_err(LoweringFlowError::Failed)?;
 
@@ -622,6 +625,20 @@ pub fn lowered_expr_to_block_scope_end(
     })
 }
 
+/// Converts [`LoweringResult<LoweredExpr>`] into `BlockScopeEnd`.
+pub fn lower_early_return(
+    ctx: &mut LoweringContext<'_, '_>,
+    mut builder: BlockBuilder,
+    ret_expr: LoweredExpr,
+    location: LocationId,
+) -> Maybe<SealedBlockBuilder> {
+    let lowered_expr = (|| {
+        let ret_var_usage = ret_expr.as_var_usage(ctx, &mut builder)?;
+        Err(LoweringFlowError::Return(ret_var_usage, location))
+    })();
+    lowered_expr_to_block_scope_end(ctx, builder, lowered_expr)
+}
+
 /// Lowers a semantic statement.
 pub fn lower_statement(
     ctx: &mut LoweringContext<'_, '_>,
diff --git a/crates/cairo-lang-lowering/src/lower/test_data/for b/crates/cairo-lang-lowering/src/lower/test_data/for
index 821ba799055..f3376b1475a 100644
--- a/crates/cairo-lang-lowering/src/lower/test_data/for
+++ b/crates/cairo-lang-lowering/src/lower/test_data/for
@@ -113,18 +113,18 @@ Statements:
   (v6: core::felt252) <- core::Felt252Add::add(v1, v2)
   (v8: core::array::SpanIter::<core::felt252>, v9: core::felt252, v7: ()) <- test::foo[118-164](v4, v6, v2)
 End:
-  Goto(blk3, {v9 -> v12, v8 -> v13, v7 -> v11})
+  Goto(blk3, {v9 -> v13, v8 -> v14, v7 -> v12})
 
 blk2:
 Statements:
-  (v14: ()) <- struct_construct()
+  (v11: ()) <- struct_construct()
 End:
-  Goto(blk3, {v1 -> v12, v4 -> v13, v14 -> v11})
+  Return(v4, v1, v11)
 
 blk3:
 Statements:
 End:
-  Return(v13, v12, v11)
+  Return(v14, v13, v12)
 
 
 Final lowering:
diff --git a/crates/cairo-lang-lowering/src/lower/test_data/loop b/crates/cairo-lang-lowering/src/lower/test_data/loop
index 1d02f4e028f..26e87bd36b7 100644
--- a/crates/cairo-lang-lowering/src/lower/test_data/loop
+++ b/crates/cairo-lang-lowering/src/lower/test_data/loop
@@ -1464,18 +1464,18 @@ Statements:
   (v15: core::felt252, v14: ()) <- core::ops::arith::DeprecatedAddAssign::<core::felt252, core::Felt252AddEq>::add_assign(v2, v13)
   (v17: test::A, v18: core::felt252, v16: ()) <- test::foo[58-134](v15, v11)
 End:
-  Goto(blk3, {v17 -> v21, v18 -> v22, v16 -> v20})
+  Goto(blk3, {v17 -> v22, v18 -> v23, v16 -> v21})
 
 blk2:
 Statements:
-  (v23: ()) <- struct_construct()
+  (v20: ()) <- struct_construct()
 End:
-  Goto(blk3, {v1 -> v21, v2 -> v22, v23 -> v20})
+  Return(v1, v2, v20)
 
 blk3:
 Statements:
 End:
-  Return(v21, v22, v20)
+  Return(v22, v23, v21)
 
 
 Final lowering:
diff --git a/crates/cairo-lang-lowering/src/lower/test_data/while b/crates/cairo-lang-lowering/src/lower/test_data/while
index 3539ef0a197..087080ff423 100644
--- a/crates/cairo-lang-lowering/src/lower/test_data/while
+++ b/crates/cairo-lang-lowering/src/lower/test_data/while
@@ -88,18 +88,18 @@ Statements:
   (v9: core::felt252) <- core::Felt252Sub::sub(v1, v8)
   (v11: core::felt252, v10: ()) <- test::foo[51-95](v9)
 End:
-  Goto(blk3, {v11 -> v14, v10 -> v13})
+  Goto(blk3, {v11 -> v15, v10 -> v14})
 
 blk2:
 Statements:
-  (v15: ()) <- struct_construct()
+  (v13: ()) <- struct_construct()
 End:
-  Goto(blk3, {v1 -> v14, v15 -> v13})
+  Return(v1, v13)
 
 blk3:
 Statements:
 End:
-  Return(v14, v13)
+  Return(v15, v14)
 
 
 Final lowering: