Move returns before goto when applicable.

crates/cairo-lang-lowering/src/optimizations/return_optimization.rs

@@ -12,7 +12,10 @@ use crate::{
     VarUsage, VariableId,
 };
 
-/// Optimizes EnumConstruct statements where all the arms reconstruct the enum and return it.
+/// Moves return earlier when applicable.
+/// Currently there are two cases:
+/// 1. EnumConstruct statements where all the arms reconstruct the enum and return it.
+/// 2. Goto that remaps the last variable and then returns it.
 pub fn return_optimization(db: &dyn LoweringGroup, lowered: &mut FlatLowered) {
     if !lowered.blocks.is_empty() {
         let ctx = ReturnOptimizerContext {
@@ -92,20 +95,29 @@ impl<'a> Analyzer<'a> for ReturnOptimizerContext<'_> {
     ) {
         match stmt {
             Statement::StructConstruct(StatementStructConstruct { inputs, output }) => {
-                if let Pattern::EnumConstruct { opt_input_var, variant, returned_vars } = info {
-                    if Some(*output) == *opt_input_var {
-                        *info = Pattern::StructConstruct {
-                            input_vars: inputs.to_vec(),
-                            variant,
-                            returned_vars,
-                        };
+                match info {
+                    Pattern::EnumConstruct { opt_input_var, variant, returned_vars } => {
+                        if Some(*output) == *opt_input_var {
+                            *info = Pattern::StructConstruct {
+                                input_vars: inputs.to_vec(),
+                                variant,
+                                returned_vars,
+                            };
+                        }
+                    }
+                    Pattern::Return { user_return, returned_vars } => {
+                        if user_return.var_id == *output
+                            || returned_vars.iter().any(|var_usage| &var_usage.var_id == output)
+                        {
+                            // If the output of the StructConstruct is returned we can't apply the
+                            // optimization.
+                            *info = Pattern::None;
+                        }
                     }
+                    _ => {}
                 }
 
-                // Keep the pattern across StructConstruct statements, this is ok, since
-                // we check the inputs to the return statement at the end.
-                // We need to allow this to handle the case where the StructConstruct is used
-                // to construct a unit type.
+                // Keep the pattern across StructConstruct statement, except for the cases above.
                 return;
             }
 
@@ -158,7 +170,7 @@ impl<'a> Analyzer<'a> for ReturnOptimizerContext<'_> {
     fn visit_goto(
         &mut self,
         info: &mut Self::Info,
-        _statement_location: StatementLocation,
+        (block_id, _statement_idx): StatementLocation,
         _target_block_id: BlockId,
         remapping: &VarRemapping,
     ) {
@@ -174,6 +186,13 @@ impl<'a> Analyzer<'a> for ReturnOptimizerContext<'_> {
             }
 
             if let Some(remapped_return) = remapping.get(&user_return.var_id) {
+                // In case user_return is remapped we can move the return earlier and avoid the
+                // remapping.
+                let mut return_vars = returned_vars.to_vec();
+                return_vars.pop();
+                return_vars.push(*remapped_return);
+                self.fixes.push(FixInfo { block_id, return_vars });
+
                 *info = Pattern::Return { user_return: *remapped_return, returned_vars };
             }
         } else if let Pattern::StructConstruct { input_vars, variant, returned_vars } = info {

crates/cairo-lang-lowering/src/optimizations/test_data/return_optimization

@@ -203,14 +203,14 @@ blk1:
 Statements:
   (v3: core::option::Option::<core::integer::u16>) <- Option::Some(v1)
 End:
-  Goto(blk3, {v3 -> v6})
+  Return(v3)
 
 blk2:
 Statements:
   (v4: ()) <- struct_construct()
   (v5: core::option::Option::<core::integer::u16>) <- Option::None(v4)
 End:
-  Goto(blk3, {v5 -> v6})
+  Return(v5)
 
 blk3:
 Statements:

crates/cairo-lang-lowering/src/test_data/assignment

@@ -33,15 +33,10 @@ blk1:
 Statements:
   (v1: core::felt252) <- 2u
 End:
-  Goto(blk3, {v1 -> v7})
+  Return(v1)
 
 blk2:
 Statements:
   (v3: core::felt252) <- 3u
 End:
-  Goto(blk3, {v3 -> v7})
-
-blk3:
-Statements:
-End:
-  Return(v7)
+  Return(v3)

crates/cairo-lang-lowering/src/test_data/destruct

@@ -50,18 +50,13 @@ Statements:
   (v5: ()) <- struct_construct()
   (v6: core::option::Option::<test::A>) <- Option::None(v5)
 End:
-  Goto(blk3, {v10 -> v11, v6 -> v7})
+  Return(v10, v6)
 
 blk2:
 Statements:
   (v3: core::option::Option::<test::A>) <- Option::Some(v1)
 End:
-  Goto(blk3, {v9 -> v11, v3 -> v7})
-
-blk3:
-Statements:
-End:
-  Return(v11, v7)
+  Return(v9, v3)
 
 //! > ==========================================================================
 

crates/cairo-lang-lowering/src/test_data/extern

@@ -174,22 +174,17 @@ End:
 blk1:
 Statements:
 End:
-  Goto(blk4, {v2 -> v6})
+  Return(v2)
 
 blk2:
 Statements:
 End:
-  Goto(blk4, {v3 -> v6})
+  Return(v3)
 
 blk3:
 Statements:
 End:
-  Goto(blk4, {v4 -> v6})
-
-blk4:
-Statements:
-End:
-  Return(v6)
+  Return(v4)
 
 //! > ==========================================================================
 
@@ -533,19 +528,14 @@ Statements:
   (v0: ()) <- struct_construct()
   (v1: core::option::Option::<()>) <- Option::Some(v0)
 End:
-  Goto(blk3, {v7 -> v11, v8 -> v12, v1 -> v4})
+  Return(v7, v8, v1)
 
 blk2:
 Statements:
   (v2: ()) <- struct_construct()
   (v3: core::option::Option::<()>) <- Option::None(v2)
 End:
-  Goto(blk3, {v9 -> v11, v10 -> v12, v3 -> v4})
-
-blk3:
-Statements:
-End:
-  Return(v11, v12, v4)
+  Return(v9, v10, v3)
 
 //! > ==========================================================================
 
@@ -583,16 +573,11 @@ Statements:
   (v0: ()) <- struct_construct()
   (v1: core::option::Option::<()>) <- Option::Some(v0)
 End:
-  Goto(blk3, {v7 -> v11, v8 -> v12, v1 -> v4})
+  Return(v7, v8, v1)
 
 blk2:
 Statements:
   (v2: ()) <- struct_construct()
   (v3: core::option::Option::<()>) <- Option::None(v2)
 End:
-  Goto(blk3, {v9 -> v11, v10 -> v12, v3 -> v4})
-
-blk3:
-Statements:
-End:
-  Return(v11, v12, v4)
+  Return(v9, v10, v3)

crates/cairo-lang-lowering/src/test_data/if

@@ -34,18 +34,13 @@ End:
 blk1:
 Statements:
 End:
-  Goto(blk3, {v1 -> v5})
+  Return(v1)
 
 blk2:
 Statements:
   (v3: core::felt252) <- 1u
 End:
-  Goto(blk3, {v3 -> v5})
-
-blk3:
-Statements:
-End:
-  Return(v5)
+  Return(v3)
 
 //! > ==========================================================================
 
@@ -86,17 +81,12 @@ blk1:
 Statements:
   (v1: core::felt252) <- 1u
 End:
-  Goto(blk3, {v1 -> v3})
+  Return(v1)
 
 blk2:
 Statements:
 End:
-  Goto(blk3, {v0 -> v3})
-
-blk3:
-Statements:
-End:
-  Return(v3)
+  Return(v0)
 
 //! > ==========================================================================
 
@@ -186,15 +176,10 @@ blk1:
 Statements:
   (v7: core::felt252) <- 0u
 End:
-  Goto(blk3, {v7 -> v10})
+  Return(v7)
 
 blk2:
 Statements:
   (v9: core::felt252) <- 1u
 End:
-  Goto(blk3, {v9 -> v10})
-
-blk3:
-Statements:
-End:
-  Return(v10)
+  Return(v9)

crates/cairo-lang-lowering/src/test_data/logical_operator

@@ -53,16 +53,11 @@ blk4:
 Statements:
   (v14: core::bool) <- test::get_bool()
 End:
-  Goto(blk6, {v14 -> v17})
+  Return(v14)
 
 blk5:
 Statements:
   (v12: ()) <- struct_construct()
   (v13: core::bool) <- bool::True(v12)
 End:
-  Goto(blk6, {v13 -> v17})
-
-blk6:
-Statements:
-End:
-  Return(v17)
+  Return(v13)