arg and ret demotion passes should be module passes (#7343)

These passes modify functions other than the one they receive as an
argument, so they cannot be function passes. Make them module passes
instead.

Author: vaivaswatha

sway-ir/src/optimize/arg_demotion.rs

@@ -4,7 +4,8 @@
 /// specific parameters.
 use crate::{
     AnalysisResults, Block, BlockArgument, Context, Function, InstOp, Instruction,
-    InstructionInserter, IrError, Pass, PassMutability, ScopedPass, Type, Value, ValueDatum,
+    InstructionInserter, IrError, Module, Pass, PassMutability, ScopedPass, Type, Value,
+    ValueDatum,
 };
 
 use rustc_hash::FxHashMap;
@@ -16,20 +17,25 @@ pub fn create_arg_demotion_pass() -> Pass {
         name: ARG_DEMOTION_NAME,
         descr: "Demotion of by-value function arguments to by-reference",
         deps: Vec::new(),
-        runner: ScopedPass::FunctionPass(PassMutability::Transform(arg_demotion)),
+        runner: ScopedPass::ModulePass(PassMutability::Transform(arg_demotion)),
     }
 }
 
 pub fn arg_demotion(
     context: &mut Context,
     _: &AnalysisResults,
-    function: Function,
+    module: Module,
 ) -> Result<bool, IrError> {
-    let mut result = fn_arg_demotion(context, function)?;
-
-    // We also need to be sure that block args within this function are demoted.
-    for block in function.block_iter(context) {
-        result |= demote_block_signature(context, &function, block);
+    let mut result = false;
+    // This is a module pass because modifying the signature of a function may affect the
+    // call sites in other functions, requiring their modification as well.
+    for function in module.function_iter(context) {
+        result |= fn_arg_demotion(context, function)?;
+
+        // We also need to be sure that block args within this function are demoted.
+        for block in function.block_iter(context) {
+            result |= demote_block_signature(context, &function, block);
+        }
     }
 
     Ok(result)

sway-ir/src/optimize/ret_demotion.rs

@@ -7,7 +7,7 @@
 /// return value is mem_copied to the new argument instead of being returned by value.
 use crate::{
     AnalysisResults, BlockArgument, Context, Function, InstOp, Instruction, InstructionInserter,
-    IrError, Pass, PassMutability, ScopedPass, Type, Value,
+    IrError, Module, Pass, PassMutability, ScopedPass, Type, Value,
 };
 
 pub const RET_DEMOTION_NAME: &str = "ret-demotion";
@@ -17,91 +17,104 @@ pub fn create_ret_demotion_pass() -> Pass {
         name: RET_DEMOTION_NAME,
         descr: "Demotion of by-value function return values to by-reference",
         deps: Vec::new(),
-        runner: ScopedPass::FunctionPass(PassMutability::Transform(ret_val_demotion)),
+        runner: ScopedPass::ModulePass(PassMutability::Transform(ret_val_demotion)),
     }
 }
 
 pub fn ret_val_demotion(
     context: &mut Context,
     _: &AnalysisResults,
-    function: Function,
+    module: Module,
 ) -> Result<bool, IrError> {
-    // Reject non-candidate.
-    let ret_type = function.get_return_type(context);
-    if !super::target_fuel::is_demotable_type(context, &ret_type) {
-        // Return type fits in a register.
-        return Ok(false);
-    }
+    // This is a module pass because we need to update all the callers of a function if we change
+    // its signature.
+    let mut changed = false;
+    for function in module.function_iter(context) {
+        // Reject non-candidate.
+        let ret_type = function.get_return_type(context);
+        if !super::target_fuel::is_demotable_type(context, &ret_type) {
+            // Return type fits in a register.
+            continue;
+        }
+
+        changed = true;
+
+        // Change the function signature.  It now returns a pointer.
+        let ptr_ret_type = Type::new_typed_pointer(context, ret_type);
+        function.set_return_type(context, ptr_ret_type);
+
+        // The storage for the return value must be determined.  For entry-point functions it's a new
+        // local and otherwise it's an extra argument.
+        let entry_block = function.get_entry_block(context);
+        let ptr_arg_val = if function.is_entry(context) {
+            let ret_var = function.new_unique_local_var(
+                context,
+                "__ret_value".to_owned(),
+                ret_type,
+                None,
+                false,
+            );
+
+            // Insert the return value pointer at the start of the entry block.
+            let get_ret_var =
+                Value::new_instruction(context, entry_block, InstOp::GetLocal(ret_var));
+            entry_block.prepend_instructions(context, vec![get_ret_var]);
+            get_ret_var
+        } else {
+            let ptr_arg_val = Value::new_argument(
+                context,
+                BlockArgument {
+                    block: entry_block,
+                    idx: function.num_args(context),
+                    ty: ptr_ret_type,
+                },
+            );
+            function.add_arg(context, "__ret_value", ptr_arg_val);
+            entry_block.add_arg(context, ptr_arg_val);
+            ptr_arg_val
+        };
 
-    // Change the function signature.  It now returns a pointer.
-    let ptr_ret_type = Type::new_typed_pointer(context, ret_type);
-    function.set_return_type(context, ptr_ret_type);
-
-    // The storage for the return value must be determined.  For entry-point functions it's a new
-    // local and otherwise it's an extra argument.
-    let entry_block = function.get_entry_block(context);
-    let ptr_arg_val = if function.is_entry(context) {
-        let ret_var =
-            function.new_unique_local_var(context, "__ret_value".to_owned(), ret_type, None, false);
-
-        // Insert the return value pointer at the start of the entry block.
-        let get_ret_var = Value::new_instruction(context, entry_block, InstOp::GetLocal(ret_var));
-        entry_block.prepend_instructions(context, vec![get_ret_var]);
-        get_ret_var
-    } else {
-        let ptr_arg_val = Value::new_argument(
-            context,
-            BlockArgument {
-                block: entry_block,
-                idx: function.num_args(context),
-                ty: ptr_ret_type,
-            },
-        );
-        function.add_arg(context, "__ret_value", ptr_arg_val);
-        entry_block.add_arg(context, ptr_arg_val);
-        ptr_arg_val
-    };
-
-    // Gather the blocks which are returning.
-    let ret_blocks = function
-        .block_iter(context)
-        .filter_map(|block| {
-            block.get_terminator(context).and_then(|term| {
-                if let InstOp::Ret(ret_val, _ty) = term.op {
-                    Some((block, ret_val))
-                } else {
-                    None
-                }
+        // Gather the blocks which are returning.
+        let ret_blocks = function
+            .block_iter(context)
+            .filter_map(|block| {
+                block.get_terminator(context).and_then(|term| {
+                    if let InstOp::Ret(ret_val, _ty) = term.op {
+                        Some((block, ret_val))
+                    } else {
+                        None
+                    }
+                })
             })
-        })
-        .collect::<Vec<_>>();
-
-    // Update each `ret` to store the return value to the 'out' arg and then return the pointer.
-    for (ret_block, ret_val) in ret_blocks {
-        // This is a special case where we're replacing the terminator.  We can just pop it off the
-        // end of the block and add new instructions.
-        let last_instr_pos = ret_block.num_instructions(context) - 1;
-        let orig_ret_val = ret_block.get_instruction_at(context, last_instr_pos);
-        ret_block.remove_instruction_at(context, last_instr_pos);
-        let md_idx = orig_ret_val.and_then(|val| val.get_metadata(context));
-
-        ret_block
-            .append(context)
-            .store(ptr_arg_val, ret_val)
-            .add_metadatum(context, md_idx);
-        ret_block
-            .append(context)
-            .ret(ptr_arg_val, ptr_ret_type)
-            .add_metadatum(context, md_idx);
-    }
-
-    // If the function isn't an entry point we need to update all the callers to pass the extra
-    // argument.
-    if !function.is_entry(context) {
-        update_callers(context, function, ret_type);
+            .collect::<Vec<_>>();
+
+        // Update each `ret` to store the return value to the 'out' arg and then return the pointer.
+        for (ret_block, ret_val) in ret_blocks {
+            // This is a special case where we're replacing the terminator.  We can just pop it off the
+            // end of the block and add new instructions.
+            let last_instr_pos = ret_block.num_instructions(context) - 1;
+            let orig_ret_val = ret_block.get_instruction_at(context, last_instr_pos);
+            ret_block.remove_instruction_at(context, last_instr_pos);
+            let md_idx = orig_ret_val.and_then(|val| val.get_metadata(context));
+
+            ret_block
+                .append(context)
+                .store(ptr_arg_val, ret_val)
+                .add_metadatum(context, md_idx);
+            ret_block
+                .append(context)
+                .ret(ptr_arg_val, ptr_ret_type)
+                .add_metadatum(context, md_idx);
+        }
+
+        // If the function isn't an entry point we need to update all the callers to pass the extra
+        // argument.
+        if !function.is_entry(context) {
+            update_callers(context, function, ret_type);
+        }
     }
 
-    Ok(true)
+    Ok(changed)
 }
 
 fn update_callers(context: &mut Context, function: Function, ret_type: Type) {

test/src/e2e_vm_tests/test_programs/should_pass/language/const_generics/stdout.snap

@@ -1,5 +1,6 @@
 ---
 source: test/src/snapshot/mod.rs
+assertion_line: 162
 ---
 > forc build --path test/src/e2e_vm_tests/test_programs/should_pass/language/const_generics --release
 exit status: 1
@@ -443,12 +444,12 @@ warning
 ____
 
   Compiled script "const_generics" with 2 warnings.
-    Finished debug [unoptimized + fuel] target(s) [9.552 KB] in ???
+    Finished debug [unoptimized + fuel] target(s) [9.56 KB] in ???
      Running 1 test, filtered 0 tests
 
 tested -- const_generics
 
-      test run_main ... ok (???, 17570 gas)
+      test run_main ... ok (???, 17574 gas)
            debug output:
 [src/main.sw:105:13] a = [1, 2]
 [src/main.sw:109:13] [C {}].len() = 1

test/src/e2e_vm_tests/test_programs/should_pass/language/panic_expression/panicking_lib/stdout.snap

@@ -8,7 +8,7 @@ output:
     Building test/src/e2e_vm_tests/test_programs/should_pass/language/panic_expression/panicking_lib
    Compiling library std (test/src/e2e_vm_tests/reduced_std_libs/sway-lib-std-core)
    Compiling library panicking_lib (test/src/e2e_vm_tests/test_programs/should_pass/language/panic_expression/panicking_lib)
-    Finished release [optimized + fuel] target(s) [5.656 KB] in ???
+    Finished release [optimized + fuel] target(s) [5.672 KB] in ???
      Running 18 tests, filtered 0 tests
 
 tested -- panicking_lib
@@ -65,14 +65,14 @@ AsciiString { data: "generic panic with string" }, log rb: 10098701174489624218
             └─ panicked in:   panicking_lib, src/lib.sw:74:5
            decoded log values:
 AsciiString { data: "generic panic different string" }, log rb: 10098701174489624218
-      test test_generic_panic_with_error_type_enum_variant ... ok (???, 293 gas)
+      test test_generic_panic_with_error_type_enum_variant ... ok (???, 297 gas)
            revert code: ffffffff00000004
             ├─ panic message: Error A.
             ├─ panic value:   A
             └─ panicked in:   panicking_lib, src/lib.sw:74:5
            decoded log values:
 A, log rb: 2721958641300806892
-      test test_generic_panic_with_error_type_enum_different_variant_same_revert_code ... ok (???, 293 gas)
+      test test_generic_panic_with_error_type_enum_different_variant_same_revert_code ... ok (???, 297 gas)
            revert code: ffffffff00000004
             ├─ panic message: Error A.
             ├─ panic value:   A