memcpyopt: Ensure source local isn't clobbered before the new memcpy (F…

…uelLabs#4422)

## Description
When combining a load+store into a `memcpy`, the optimization was
missing a check that the source isn't clobbered (stored to) in between
the load and store.

The optimization thus is enabled for copy types too.

Issue FuelLabs#4345

## Checklist

- [x] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [x] I have updated the documentation where relevant (API docs, the
reference, and the Sway book).
- [x] I have added tests that prove my fix is effective or that my
feature works.
- [x] I have added (or requested a maintainer to add) the necessary
`Breaking*` or `New Feature` labels where relevant.
- [x] I have done my best to ensure that my PR adheres to [the Fuel Labs
Code Review
Standards](https://github.com/FuelLabs/rfcs/blob/master/text/code-standards/external-contributors.md).
- [x] I have requested a review from the relevant team or maintainers.

---------

Co-authored-by: Mohammad Fawaz <[email protected]>

sway-ir/src/block.rs

@@ -41,7 +41,7 @@ pub struct BlockContent {
     pub preds: FxHashSet<Block>,
 }
 
-#[derive(Debug, Clone, DebugWithContext)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, DebugWithContext)]
 pub struct BlockArgument {
     /// The block of which this is an argument.
     pub block: Block,
@@ -138,6 +138,17 @@ impl Block {
         idx
     }
 
+    pub fn set_arg(&self, context: &mut Context, arg: Value) {
+        match context.values[arg.0].value {
+            ValueDatum::Argument(BlockArgument { block, idx, ty: _ })
+                if block == *self && idx < context.blocks[self.0].args.len() =>
+            {
+                context.blocks[self.0].args[idx] = arg;
+            }
+            _ => panic!("Inconsistent block argument being set"),
+        }
+    }
+
     /// Add a block argument, asserts that `arg` is suitable here.
     pub fn add_arg(&self, context: &mut Context, arg: Value) {
         match context.values[arg.0].value {

sway-ir/src/instruction.rs

@@ -442,7 +442,7 @@ impl Instruction {
                 indices.iter_mut().for_each(replace);
             }
             Instruction::IntToPtr(value, _) => replace(value),
-            Instruction::Load(_) => (),
+            Instruction::Load(ptr) => replace(ptr),
             Instruction::MemCopyBytes {
                 dst_val_ptr,
                 src_val_ptr,
@@ -461,8 +461,12 @@ impl Instruction {
             Instruction::Nop => (),
             Instruction::PtrToInt(value, _) => replace(value),
             Instruction::Ret(ret_val, _) => replace(ret_val),
-            Instruction::Store { stored_val, .. } => {
+            Instruction::Store {
+                stored_val,
+                dst_val_ptr,
+            } => {
                 replace(stored_val);
+                replace(dst_val_ptr);
             }
 
             Instruction::FuelVm(fuel_vm_instr) => match fuel_vm_instr {

sway-ir/src/optimize/arg_demotion.rs

@@ -89,17 +89,6 @@ fn fn_arg_demotion(context: &mut Context, function: Function) -> Result<bool, Ir
     Ok(true)
 }
 
-// Match the mutable argument value and change its type.
-macro_rules! set_arg_type {
-    ($context: ident, $arg_val: ident, $new_ty: ident) => {
-        if let ValueDatum::Argument(BlockArgument { ty, .. }) =
-            &mut $context.values[$arg_val.0].value
-        {
-            *ty = $new_ty
-        }
-    };
-}
-
 fn demote_fn_signature(context: &mut Context, function: &Function, arg_idcs: &[(usize, Type)]) {
     // Change the types of the arg values in place to their pointer counterparts.
     let entry_block = function.get_entry_block(context);
@@ -108,31 +97,40 @@ fn demote_fn_signature(context: &mut Context, function: &Function, arg_idcs: &[(
         .map(|(arg_idx, arg_ty)| {
             let ptr_ty = Type::new_ptr(context, *arg_ty);
 
-            // Update the function signature.
-            let fn_args = &context.functions[function.0].arguments;
-            let (_name, fn_arg_val) = &fn_args[*arg_idx];
-            set_arg_type!(context, fn_arg_val, ptr_ty);
-
-            // Update the entry block signature.
+            // Create a new block arg, same as the old one but with a different type.
             let blk_arg_val = entry_block
                 .get_arg(context, *arg_idx)
                 .expect("Entry block args should be mirror of function args.");
-            set_arg_type!(context, blk_arg_val, ptr_ty);
+            let ValueDatum::Argument(block_arg) = context.values[blk_arg_val.0].value else {
+                panic!("Block argument is not of right Value kind");
+            };
+            let new_blk_arg_val = Value::new_argument(
+                context,
+                BlockArgument {
+                    ty: ptr_ty,
+                    ..block_arg
+                },
+            );
+
+            // Set both function and block arg to the new one.
+            entry_block.set_arg(context, new_blk_arg_val);
+            let (_name, fn_arg_val) = &mut context.functions[function.0].arguments[*arg_idx];
+            *fn_arg_val = new_blk_arg_val;
 
-            *fn_arg_val
+            (blk_arg_val, new_blk_arg_val)
         })
         .collect::<Vec<_>>();
 
     // For each of the old args, which have had their types changed, insert a `load` instruction.
     let arg_val_pairs = old_arg_vals
         .into_iter()
         .rev()
-        .map(|old_arg_val| {
-            let new_arg_val = Value::new_instruction(context, Instruction::Load(old_arg_val));
+        .map(|(old_arg_val, new_arg_val)| {
+            let load_from_new_arg = Value::new_instruction(context, Instruction::Load(new_arg_val));
             context.blocks[entry_block.0]
                 .instructions
-                .insert(0, new_arg_val);
-            (old_arg_val, new_arg_val)
+                .insert(0, load_from_new_arg);
+            (old_arg_val, load_from_new_arg)
         })
         .collect::<Vec<_>>();
 
@@ -237,9 +235,21 @@ fn demote_block_signature(context: &mut Context, function: &Function, block: Blo
         .rev()
         .map(|(_arg_idx, arg_val, arg_ty)| {
             let ptr_ty = Type::new_ptr(context, *arg_ty);
-            set_arg_type!(context, arg_val, ptr_ty);
 
-            let load_val = Value::new_instruction(context, Instruction::Load(*arg_val));
+            // Create a new block arg, same as the old one but with a different type.
+            let ValueDatum::Argument(block_arg) = context.values[arg_val.0].value else {
+                panic!("Block argument is not of right Value kind");
+            };
+            let new_blk_arg_val = Value::new_argument(
+                context,
+                BlockArgument {
+                    ty: ptr_ty,
+                    ..block_arg
+                },
+            );
+            block.set_arg(context, new_blk_arg_val);
+
+            let load_val = Value::new_instruction(context, Instruction::Load(new_blk_arg_val));
             let block_instrs = &mut context.blocks[block.0].instructions;
             block_instrs.insert(0, load_val);