[asmgen] bugfixes in const_indexing_aggregates_function (#6834)

## Description

Fixes #6810, fixes #6812 and fixes #6813

---------

Co-authored-by: Joshua Batty <joshpbatty@gmail.com>

Author: vaivaswatha

sway-core/src/asm_generation/fuel/functions.rs

@@ -902,7 +902,9 @@ impl FuelAsmBuilder<'_, '_> {
             todo!("Enormous stack usage for locals.");
         }
         self.cur_bytecode.push(Op {
-            opcode: Either::Left(VirtualOp::CFEI(VirtualImmediate24 {
+            opcode: Either::Left(VirtualOp::CFEI(
+                VirtualRegister::Constant(ConstantRegister::StackPointer),
+                VirtualImmediate24 {
                 value: locals_size_bytes as u32 + (max_num_extra_args * 8) as u32,
             })),
             comment: format!("allocate {locals_size_bytes} bytes for locals and {max_num_extra_args} slots for call arguments"),
@@ -1053,7 +1055,9 @@ impl FuelAsmBuilder<'_, '_> {
             todo!("Enormous stack usage for locals.");
         }
         self.cur_bytecode.push(Op {
-            opcode: Either::Left(VirtualOp::CFSI(VirtualImmediate24 {
+            opcode: Either::Left(
+                VirtualOp::CFSI(VirtualRegister::Constant(ConstantRegister::StackPointer),
+                VirtualImmediate24 {
                 value: u32::try_from(locals_size_bytes + (max_num_extra_args * 8)).unwrap(),
             })),
             comment: format!("free {locals_size_bytes} bytes for locals and {max_num_extra_args} slots for extra call arguments"),

sway-core/src/asm_generation/fuel/optimizations.rs

@@ -145,25 +145,23 @@ impl AbstractInstructionSet {
                         }
                         VirtualOp::LW(dest, addr_reg, imm) => match reg_contents.get(addr_reg) {
                             Some(RegContents::BaseOffset(base_reg, offset))
-                                if get_def_version(&latest_version, &base_reg.reg)
-                                    == base_reg.ver
+                                if offset % 8 == 0
+                                    && get_def_version(&latest_version, &base_reg.reg)
+                                        == base_reg.ver
                                     && ((offset / 8) + imm.value as u64)
                                         < compiler_constants::TWELVE_BITS =>
                             {
-                                // bail if LW cannot read where this memory is
-                                if offset % 8 == 0 {
-                                    let new_imm = VirtualImmediate12::new_unchecked(
-                                        (offset / 8) + imm.value as u64,
-                                        "Immediate offset too big for LW",
-                                    );
-                                    let new_lw =
-                                        VirtualOp::LW(dest.clone(), base_reg.reg.clone(), new_imm);
-                                    // The register defined is no more useful for us. Forget anything from its past.
-                                    reg_contents.remove(dest);
-                                    record_new_def(&mut latest_version, dest);
-                                    // Replace the LW with a new one in-place.
-                                    *op = new_lw;
-                                }
+                                let new_imm = VirtualImmediate12::new_unchecked(
+                                    (offset / 8) + imm.value as u64,
+                                    "Immediate offset too big for LW",
+                                );
+                                let new_lw =
+                                    VirtualOp::LW(dest.clone(), base_reg.reg.clone(), new_imm);
+                                // The register defined is no more useful for us. Forget anything from its past.
+                                reg_contents.remove(dest);
+                                record_new_def(&mut latest_version, dest);
+                                // Replace the LW with a new one in-place.
+                                *op = new_lw;
                             }
                             _ => {
                                 reg_contents.remove(dest);
@@ -172,8 +170,9 @@ impl AbstractInstructionSet {
                         },
                         VirtualOp::SW(addr_reg, src, imm) => match reg_contents.get(addr_reg) {
                             Some(RegContents::BaseOffset(base_reg, offset))
-                                if get_def_version(&latest_version, &base_reg.reg)
-                                    == base_reg.ver
+                                if offset % 8 == 0
+                                    && get_def_version(&latest_version, &base_reg.reg)
+                                        == base_reg.ver
                                     && ((offset / 8) + imm.value as u64)
                                         < compiler_constants::TWELVE_BITS =>
                             {

sway-core/src/asm_generation/fuel/register_allocator.rs

@@ -714,11 +714,11 @@ fn spill(ops: &[Op], spills: &FxHashSet<VirtualRegister>) -> Vec<Op> {
     let mut cfs_idx_opt = None;
     for (op_idx, op) in ops.iter().enumerate() {
         match &op.opcode {
-            Either::Left(VirtualOp::CFEI(_)) => {
+            Either::Left(VirtualOp::CFEI(..)) => {
                 assert!(cfe_idx_opt.is_none(), "Found more than one stack extension");
                 cfe_idx_opt = Some(op_idx);
             }
-            Either::Left(VirtualOp::CFSI(_)) => {
+            Either::Left(VirtualOp::CFSI(..)) => {
                 assert!(cfs_idx_opt.is_none(), "Found more than one stack shrink");
                 cfs_idx_opt = Some(op_idx);
             }
@@ -728,9 +728,12 @@ fn spill(ops: &[Op], spills: &FxHashSet<VirtualRegister>) -> Vec<Op> {
 
     let cfe_idx = cfe_idx_opt.expect("Function does not have CFEI instruction for locals");
 
-    let Either::Left(VirtualOp::CFEI(VirtualImmediate24 {
-        value: locals_size_bytes,
-    })) = ops[cfe_idx].opcode
+    let Either::Left(VirtualOp::CFEI(
+        VirtualRegister::Constant(ConstantRegister::StackPointer),
+        VirtualImmediate24 {
+            value: locals_size_bytes,
+        },
+    )) = ops[cfe_idx].opcode
     else {
         panic!("Unexpected opcode");
     };
@@ -755,18 +758,24 @@ fn spill(ops: &[Op], spills: &FxHashSet<VirtualRegister>) -> Vec<Op> {
         if op_idx == cfe_idx {
             // This is the CFE instruction, use the new stack size.
             spilled.push(Op {
-                opcode: Either::Left(VirtualOp::CFEI(VirtualImmediate24 {
-                    value: new_locals_byte_size,
-                })),
+                opcode: Either::Left(VirtualOp::CFEI(
+                    VirtualRegister::Constant(ConstantRegister::StackPointer),
+                    VirtualImmediate24 {
+                        value: new_locals_byte_size,
+                    },
+                )),
                 comment: op.comment.clone() + &format!(" and {spills_size} bytes for spills"),
                 owning_span: op.owning_span.clone(),
             });
         } else if matches!(cfs_idx_opt, Some(cfs_idx) if cfs_idx == op_idx) {
             // This is the CFS instruction, use the new stack size.
             spilled.push(Op {
-                opcode: Either::Left(VirtualOp::CFSI(VirtualImmediate24 {
-                    value: new_locals_byte_size,
-                })),
+                opcode: Either::Left(VirtualOp::CFSI(
+                    VirtualRegister::Constant(ConstantRegister::StackPointer),
+                    VirtualImmediate24 {
+                        value: new_locals_byte_size,
+                    },
+                )),
                 comment: op.comment.clone() + &format!(" and {spills_size} bytes for spills"),
                 owning_span: op.owning_span.clone(),
             });

sway-core/src/asm_lang/mod.rs

@@ -120,7 +120,10 @@ impl Op {
         size_to_allocate_in_bytes: VirtualImmediate24,
     ) -> Self {
         Op {
-            opcode: Either::Left(VirtualOp::CFEI(size_to_allocate_in_bytes)),
+            opcode: Either::Left(VirtualOp::CFEI(
+                VirtualRegister::Constant(ConstantRegister::StackPointer),
+                size_to_allocate_in_bytes,
+            )),
             comment: String::new(),
             owning_span: None,
         }
@@ -481,23 +484,35 @@ impl Op {
             /* Memory Instructions */
             "aloc" => {
                 let r1 = single_reg(handler, args, immediate, whole_op_span)?;
-                VirtualOp::ALOC(r1)
+                VirtualOp::ALOC(VirtualRegister::Constant(ConstantRegister::HeapPointer), r1)
             }
             "cfei" => {
                 let imm = single_imm_24(handler, args, immediate, whole_op_span)?;
-                VirtualOp::CFEI(imm)
+                VirtualOp::CFEI(
+                    VirtualRegister::Constant(ConstantRegister::StackPointer),
+                    imm,
+                )
             }
             "cfsi" => {
                 let imm = single_imm_24(handler, args, immediate, whole_op_span)?;
-                VirtualOp::CFSI(imm)
+                VirtualOp::CFSI(
+                    VirtualRegister::Constant(ConstantRegister::StackPointer),
+                    imm,
+                )
             }
             "cfe" => {
                 let r1 = single_reg(handler, args, immediate, whole_op_span)?;
-                VirtualOp::CFE(r1)
+                VirtualOp::CFE(
+                    VirtualRegister::Constant(ConstantRegister::StackPointer),
+                    r1,
+                )
             }
             "cfs" => {
                 let r1 = single_reg(handler, args, immediate, whole_op_span)?;
-                VirtualOp::CFS(r1)
+                VirtualOp::CFS(
+                    VirtualRegister::Constant(ConstantRegister::StackPointer),
+                    r1,
+                )
             }
             "lb" => {
                 let (r1, r2, imm) = two_regs_imm_12(handler, args, immediate, whole_op_span)?;
@@ -1164,11 +1179,11 @@ impl fmt::Display for VirtualOp {
             RET(a) => write!(fmtr, "ret {a}"),
 
             /* Memory Instructions */
-            ALOC(a) => write!(fmtr, "aloc {a}"),
-            CFEI(a) => write!(fmtr, "cfei {a}"),
-            CFSI(a) => write!(fmtr, "cfsi {a}"),
-            CFE(a) => write!(fmtr, "cfe {a}"),
-            CFS(a) => write!(fmtr, "cfs {a}"),
+            ALOC(_hp, a) => write!(fmtr, "aloc {a}"),
+            CFEI(_sp, a) => write!(fmtr, "cfei {a}"),
+            CFSI(_sp, a) => write!(fmtr, "cfsi {a}"),
+            CFE(_sp, a) => write!(fmtr, "cfe {a}"),
+            CFS(_sp, a) => write!(fmtr, "cfs {a}"),
             LB(a, b, c) => write!(fmtr, "lb {a} {b} {c}"),
             LW(a, b, c) => write!(fmtr, "lw {a} {b} {c}"),
             MCL(a, b) => write!(fmtr, "mcl {a} {b}"),

sway-core/src/asm_lang/virtual_ops.rs

@@ -109,11 +109,11 @@ pub(crate) enum VirtualOp {
     RET(VirtualRegister),
 
     /* Memory Instructions */
-    ALOC(VirtualRegister),
-    CFEI(VirtualImmediate24),
-    CFSI(VirtualImmediate24),
-    CFE(VirtualRegister),
-    CFS(VirtualRegister),
+    ALOC(VirtualRegister, VirtualRegister),
+    CFEI(VirtualRegister, VirtualImmediate24),
+    CFSI(VirtualRegister, VirtualImmediate24),
+    CFE(VirtualRegister, VirtualRegister),
+    CFS(VirtualRegister, VirtualRegister),
     LB(VirtualRegister, VirtualRegister, VirtualImmediate12),
     LW(VirtualRegister, VirtualRegister, VirtualImmediate12),
     MCL(VirtualRegister, VirtualRegister),
@@ -290,11 +290,11 @@ impl VirtualOp {
             RET(r1) => vec![r1],
 
             /* Memory Instructions */
-            ALOC(r1) => vec![r1],
-            CFEI(_imm) => vec![],
-            CFSI(_imm) => vec![],
-            CFE(r1) => vec![r1],
-            CFS(r1) => vec![r1],
+            ALOC(hp, r1) => vec![hp, r1],
+            CFEI(sp, _imm) => vec![sp],
+            CFSI(sp, _imm) => vec![sp],
+            CFE(sp, r1) => vec![sp, r1],
+            CFS(sp, r1) => vec![sp, r1],
             LB(r1, r2, _i) => vec![r1, r2],
             LW(r1, r2, _i) => vec![r1, r2],
             MCL(r1, r2) => vec![r1, r2],
@@ -425,11 +425,11 @@ impl VirtualOp {
             | JNEI(_, _, _)
             | JNZI(_, _)
             | RET(_)
-            | ALOC(_)
-            | CFEI(_)
-            | CFSI(_)
-            | CFE(_)
-            | CFS(_)
+            | ALOC(..)
+            | CFEI(..)
+            | CFSI(..)
+            | CFE(..)
+            | CFS(..)
             | MCL(_, _)
             | MCLI(_, _)
             | MCP(_, _, _)
@@ -522,17 +522,17 @@ impl VirtualOp {
             | ED19(_, _, _, _)
              => vec![&VirtualRegister::Constant(Overflow), &VirtualRegister::Constant(Error)],
             FLAG(_) => vec![&VirtualRegister::Constant(Flags)],
+            | ALOC(hp, _) => vec![hp],
+            | CFEI(sp, _)
+            | CFSI(sp, _)
+            | CFE(sp, _)
+            | CFS(sp, _) => vec![sp],
             JMP(_)
             | JI(_)
             | JNE(_, _, _)
             | JNEI(_, _, _)
             | JNZI(_, _)
             | RET(_)
-            | ALOC(_)
-            | CFEI(_)
-            | CFSI(_)
-            | CFE(_)
-            | CFS(_)
             | LB(_, _, _)
             | LW(_, _, _)
             | MCL(_, _)
@@ -638,11 +638,11 @@ impl VirtualOp {
             RET(r1) => vec![r1],
 
             /* Memory Instructions */
-            ALOC(r1) => vec![r1],
-            CFEI(_imm) => vec![],
-            CFSI(_imm) => vec![],
-            CFE(r1) => vec![r1],
-            CFS(r1) => vec![r1],
+            ALOC(hp, r1) => vec![hp, r1],
+            CFEI(sp, _imm) => vec![sp],
+            CFSI(sp, _imm) => vec![sp],
+            CFE(sp, r1) => vec![sp, r1],
+            CFS(sp, r1) => vec![sp, r1],
             LB(_r1, r2, _i) => vec![r2],
             LW(_r1, r2, _i) => vec![r2],
             MCL(r1, r2) => vec![r1, r2],
@@ -760,11 +760,11 @@ impl VirtualOp {
             RET(_r1) => vec![],
 
             /* Memory Instructions */
-            ALOC(_r1) => vec![],
-            CFEI(_imm) => vec![],
-            CFSI(_imm) => vec![],
-            CFE(_r1) => vec![],
-            CFS(_r1) => vec![],
+            ALOC(hp, _r1) => vec![hp],
+            CFEI(sp, _imm) => vec![sp],
+            CFSI(sp, _imm) => vec![sp],
+            CFE(sp, _r1) => vec![sp],
+            CFS(sp, _r1) => vec![sp],
             LB(r1, _r2, _i) => vec![r1],
             LW(r1, _r2, _i) => vec![r1],
             MCL(_r1, _r2) => vec![],
@@ -1060,11 +1060,11 @@ impl VirtualOp {
             RET(r1) => Self::RET(update_reg(reg_to_reg_map, r1)),
 
             /* Memory Instructions */
-            ALOC(r1) => Self::ALOC(update_reg(reg_to_reg_map, r1)),
-            CFEI(i) => Self::CFEI(i.clone()),
-            CFSI(i) => Self::CFSI(i.clone()),
-            CFE(r1) => Self::CFE(update_reg(reg_to_reg_map, r1)),
-            CFS(r1) => Self::CFS(update_reg(reg_to_reg_map, r1)),
+            ALOC(hp, r1) => Self::ALOC(hp.clone(), update_reg(reg_to_reg_map, r1)),
+            CFEI(sp, i) => Self::CFEI(sp.clone(), i.clone()),
+            CFSI(sp, i) => Self::CFSI(sp.clone(), i.clone()),
+            CFE(sp, r1) => Self::CFE(sp.clone(), update_reg(reg_to_reg_map, r1)),
+            CFS(sp, r1) => Self::CFS(sp.clone(), update_reg(reg_to_reg_map, r1)),
             LB(r1, r2, i) => Self::LB(
                 update_reg(reg_to_reg_map, r1),
                 update_reg(reg_to_reg_map, r2),
@@ -1550,11 +1550,11 @@ impl VirtualOp {
             RET(reg) => AllocatedOpcode::RET(map_reg(&mapping, reg)),
 
             /* Memory Instructions */
-            ALOC(reg) => AllocatedOpcode::ALOC(map_reg(&mapping, reg)),
-            CFEI(imm) => AllocatedOpcode::CFEI(imm.clone()),
-            CFSI(imm) => AllocatedOpcode::CFSI(imm.clone()),
-            CFE(reg) => AllocatedOpcode::CFE(map_reg(&mapping, reg)),
-            CFS(reg) => AllocatedOpcode::CFS(map_reg(&mapping, reg)),
+            ALOC(_hp, reg) => AllocatedOpcode::ALOC(map_reg(&mapping, reg)),
+            CFEI(_sp, imm) => AllocatedOpcode::CFEI(imm.clone()),
+            CFSI(_sp, imm) => AllocatedOpcode::CFSI(imm.clone()),
+            CFE(_sp, reg) => AllocatedOpcode::CFE(map_reg(&mapping, reg)),
+            CFS(_sp, reg) => AllocatedOpcode::CFS(map_reg(&mapping, reg)),
             LB(reg1, reg2, imm) => AllocatedOpcode::LB(
                 map_reg(&mapping, reg1),
                 map_reg(&mapping, reg2),
@@ -1775,6 +1775,7 @@ fn update_reg(
     reg: &VirtualRegister,
 ) -> VirtualRegister {
     if let Some(r) = reg_to_reg_map.get(reg) {
+        assert!(reg.is_virtual(), "Only virtual registers should be updated");
         (*r).into()
     } else {
         reg.clone()

test/src/e2e_vm_tests/test_programs/should_pass/language/const_indexing_aggregates_asmgen/Forc.lock

@@ -0,0 +1,13 @@
+[[package]]
+name = "const_indexing_aggregates_asmgen"
+source = "member"
+dependencies = ["std"]
+
+[[package]]
+name = "core"
+source = "path+from-root-BA8D9B16C05D491F"
+
+[[package]]
+name = "std"
+source = "path+from-root-BA8D9B16C05D491F"
+dependencies = ["core"]

test/src/e2e_vm_tests/test_programs/should_pass/language/const_indexing_aggregates_asmgen/Forc.toml

@@ -0,0 +1,8 @@
+[project]
+authors = ["Fuel Labs <contact@fuel.sh>"]
+entry = "main.sw"
+license = "Apache-2.0"
+name = "const_indexing_aggregates_asmgen"
+
+[dependencies]
+std = { path = "../../../../reduced_std_libs/sway-lib-std-assert" }