Add support for ASM control flow to/from addresses which don't fit in…

… immediate offsets. (FuelLabs#3109)

- Introduce a new control flow op `LoadLabel` which reads an address
  from the data section.
- Convert `JI`, `JNZI`, `JNEI` and `MoveAddress` to use `LoadLabel` and
  `JMP` or `JNE` when required.
- Introduce a `BLOB` instruction which is for testing purposes only, and
  efficiently & conveniently inserts a raft of `NOOP`s into the binary.

sway-ast/src/expr/op_code.rs

@@ -171,6 +171,7 @@ define_op_codes!(
     (Mul, MulOpcode, "mul", (ret: reg, lhs: reg, rhs: reg)),
     (Muli, MuliOpcode, "muli", (ret: reg, lhs: reg, rhs: imm)),
     (Noop, NoopOpcode, "noop", ()),
+    (Blob, BlobOpcode, "blob", (size: imm)),
     (Not, NotOpcode, "not", (ret: reg, arg: reg)),
     (Or, OrOpcode, "or", (ret: reg, lhs: reg, rhs: reg)),
     (Ori, OriOpcode, "ori", (ret: reg, lhs: reg, rhs: imm)),

sway-core/src/asm_generation/abstract_instruction_set.rs

@@ -190,7 +190,6 @@ impl RealizedAbstractInstructionSet {
                      opcode,
                      comment,
                      owning_span,
-                     offset: _,
                  }| {
                     AllocatedOp {
                         opcode,

sway-core/src/asm_generation/finalized_asm.rs

@@ -105,6 +105,7 @@ fn to_bytecode_mut(
                 {
                     acc + 8
                 }
+                AllocatedOpcode::BLOB(count) => acc + count.value as u64 * 4,
                 _ => acc + 4,
             })
             + 4;

sway-core/src/asm_generation/from_ir.rs

@@ -51,7 +51,12 @@ pub fn compile_ir_to_asm(
         println!("{allocated_program}");
     }
 
-    let final_program = allocated_program.into_final_program();
+    let final_program = check!(
+        CompileResult::from(allocated_program.into_final_program()),
+        return err(warnings, errors),
+        warnings,
+        errors
+    );
 
     if build_config
         .map(|cfg| cfg.print_finalized_asm)

sway-core/src/asm_generation/programs/allocated.rs

@@ -3,7 +3,7 @@ use super::{AllocatedProgram, FinalProgram};
 use crate::asm_generation::{AllocatedAbstractInstructionSet, InstructionSet};
 
 impl AllocatedProgram {
-    pub(crate) fn into_final_program(self) -> FinalProgram {
+    pub(crate) fn into_final_program(mut self) -> Result<FinalProgram, crate::CompileError> {
         // Concat the prologue and all the functions together.
         let abstract_ops = AllocatedAbstractInstructionSet {
             ops: std::iter::once(self.prologue.ops)
@@ -12,16 +12,16 @@ impl AllocatedProgram {
                 .collect(),
         };
 
-        let realized_ops = abstract_ops.realize_labels(&self.data_section);
+        let realized_ops = abstract_ops.realize_labels(&mut self.data_section)?;
         let ops = InstructionSet {
             ops: realized_ops.pad_to_even(),
         };
 
-        FinalProgram {
+        Ok(FinalProgram {
             kind: self.kind,
             data_section: self.data_section,
             ops,
-        }
+        })
     }
 }
 

sway-core/src/asm_lang/allocated_ops.rs

@@ -92,10 +92,11 @@ pub(crate) enum AllocatedOpcode {
     SUBI(AllocatedRegister, AllocatedRegister, VirtualImmediate12),
     XOR(AllocatedRegister, AllocatedRegister, AllocatedRegister),
     XORI(AllocatedRegister, AllocatedRegister, VirtualImmediate12),
-    JMP(AllocatedRegister),
     JI(VirtualImmediate24),
     JNEI(AllocatedRegister, AllocatedRegister, VirtualImmediate12),
     JNZI(AllocatedRegister, VirtualImmediate18),
+    JMP(AllocatedRegister),
+    JNE(AllocatedRegister, AllocatedRegister, AllocatedRegister),
     RET(AllocatedRegister),
     RETD(AllocatedRegister, AllocatedRegister),
     CFEI(VirtualImmediate24),
@@ -166,6 +167,7 @@ pub(crate) enum AllocatedOpcode {
     K256(AllocatedRegister, AllocatedRegister, AllocatedRegister),
     S256(AllocatedRegister, AllocatedRegister, AllocatedRegister),
     NOOP,
+    BLOB(VirtualImmediate24),
     FLAG(AllocatedRegister),
     GM(AllocatedRegister, VirtualImmediate18),
     Undefined,
@@ -209,10 +211,11 @@ impl AllocatedOpcode {
             SUBI(r1, _r2, _i) => vec![r1],
             XOR(r1, _r2, _r3) => vec![r1],
             XORI(r1, _r2, _i) => vec![r1],
-            JMP(_r1) => vec![],
             JI(_im) => vec![],
             JNEI(_r1, _r2, _i) => vec![],
             JNZI(_r1, _i) => vec![],
+            JMP(_r1) => vec![],
+            JNE(_r1, _r2, _r3) => vec![],
             RET(_r1) => vec![],
             RETD(_r1, _r2) => vec![],
             CFEI(_imm) => vec![],
@@ -253,6 +256,7 @@ impl AllocatedOpcode {
             K256(_r1, _r2, _r3) => vec![],
             S256(_r1, _r2, _r3) => vec![],
             NOOP => vec![],
+            BLOB(_imm) => vec![],
             FLAG(_r1) => vec![],
             GM(r1, _imm) => vec![r1],
             Undefined | DataSectionOffsetPlaceholder => vec![],
@@ -301,10 +305,11 @@ impl fmt::Display for AllocatedOpcode {
             SUBI(a, b, c) => write!(fmtr, "subi {} {} {}", a, b, c),
             XOR(a, b, c) => write!(fmtr, "xor  {} {} {}", a, b, c),
             XORI(a, b, c) => write!(fmtr, "xori {} {} {}", a, b, c),
-            JMP(a) => write!(fmtr, "jmp {}", a),
             JI(a) => write!(fmtr, "ji   {}", a),
             JNEI(a, b, c) => write!(fmtr, "jnei {} {} {}", a, b, c),
             JNZI(a, b) => write!(fmtr, "jnzi {} {}", a, b),
+            JMP(a) => write!(fmtr, "jmp {}", a),
+            JNE(a, b, c) => write!(fmtr, "jne  {a} {b} {c}"),
             RET(a) => write!(fmtr, "ret  {}", a),
             RETD(a, b) => write!(fmtr, "retd  {} {}", a, b),
             CFEI(a) => write!(fmtr, "cfei {}", a),
@@ -345,6 +350,7 @@ impl fmt::Display for AllocatedOpcode {
             K256(a, b, c) => write!(fmtr, "k256 {} {} {}", a, b, c),
             S256(a, b, c) => write!(fmtr, "s256 {} {} {}", a, b, c),
             NOOP => write!(fmtr, "noop"),
+            BLOB(a) => write!(fmtr, "blob {a}"),
             FLAG(a) => write!(fmtr, "flag {}", a),
             GM(a, b) => write!(fmtr, "gm {} {}", a, b),
             Undefined => write!(fmtr, "undefined op"),
@@ -426,17 +432,18 @@ impl AllocatedOp {
             SUBI(a, b, c)   => VmOp::SUBI(a.to_register_id(), b.to_register_id(), c.value),
             XOR (a, b, c)   => VmOp::XOR (a.to_register_id(), b.to_register_id(), c.to_register_id()),
             XORI(a, b, c)   => VmOp::XORI(a.to_register_id(), b.to_register_id(), c.value),
-            JMP(a)          => VmOp::JMP(a.to_register_id()),
             JI  (a)         => VmOp::JI  (a.value),
             JNEI(a, b, c)   => VmOp::JNEI(a.to_register_id(), b.to_register_id(), c.value),
             JNZI(a, b)      => VmOp::JNZI(a.to_register_id(), b.value),
+            JMP (a)         => VmOp::JMP(a.to_register_id()),
+            JNE (a, b, c)   => VmOp::JNE(a.to_register_id(), b.to_register_id(), c.to_register_id()),
             RET (a)         => VmOp::RET (a.to_register_id()),
             RETD(a, b)      => VmOp::RETD (a.to_register_id(), b.to_register_id()),
             CFEI(a)         => VmOp::CFEI(a.value),
             CFSI(a)         => VmOp::CFSI(a.value),
             LB  (a, b, c)   => VmOp::LB  (a.to_register_id(), b.to_register_id(), c.value),
-            LWDataId  (a, b)=> return Either::Left(realize_lw(a, b, data_section, offset_to_data_section)),
-            LW (a, b, c)    => VmOp::LW(a.to_register_id(), b.to_register_id(), c.value),
+            LW  (a, b, c)   => VmOp::LW(a.to_register_id(), b.to_register_id(), c.value),
+            LWDataId (a, b) => return Either::Left(realize_lw(a, b, data_section, offset_to_data_section)),
             ALOC(a)         => VmOp::ALOC(a.to_register_id()),
             MCL (a, b)      => VmOp::MCL (a.to_register_id(), b.to_register_id()),
             MCLI(a, b)      => VmOp::MCLI(a.to_register_id(), b.value),
@@ -445,7 +452,7 @@ impl AllocatedOp {
             MEQ (a, b, c, d)=> VmOp::MEQ (a.to_register_id(), b.to_register_id(), c.to_register_id(), d.to_register_id()),
             SB  (a, b, c)   => VmOp::SB  (a.to_register_id(), b.to_register_id(), c.value),
             SW  (a, b, c)   => VmOp::SW  (a.to_register_id(), b.to_register_id(), c.value),
-            BAL  (a, b, c)  => VmOp::BAL (a.to_register_id(), b.to_register_id(), c.to_register_id()),
+            BAL (a, b, c)   => VmOp::BAL (a.to_register_id(), b.to_register_id(), c.to_register_id()),
             BHSH(a, b)      => VmOp::BHSH(a.to_register_id(), b.to_register_id()),
             BHEI(a)         => VmOp::BHEI(a.to_register_id()),
             BURN(a)         => VmOp::BURN(a.to_register_id()),
@@ -470,8 +477,9 @@ impl AllocatedOp {
             K256(a, b, c)   => VmOp::K256(a.to_register_id(), b.to_register_id(), c.to_register_id()),
             S256(a, b, c)   => VmOp::S256(a.to_register_id(), b.to_register_id(), c.to_register_id()),
             NOOP            => VmOp::NOOP,
+            BLOB(a)         => return Either::Left(std::iter::repeat(VmOp::NOOP).take(a.value as usize).collect()),
             FLAG(a)         => VmOp::FLAG(a.to_register_id()),
-            GM(a, b)         => VmOp::GM(a.to_register_id(), b.value),
+            GM  (a, b)      => VmOp::GM(a.to_register_id(), b.value),
             Undefined       => VmOp::Undefined,
             DataSectionOffsetPlaceholder => return Either::Right(offset_to_data_section.to_be_bytes()),
             DataSectionRegisterLoadPlaceholder => VmOp::LW(crate::asm_generation::compiler_constants::DATA_SECTION_REGISTER as fuel_asm::RegisterId, ConstantRegister::InstructionStart.to_register_id(), 1),

sway-core/src/asm_lang/mod.rs

@@ -62,7 +62,6 @@ pub(crate) struct RealizedOp {
     /// A descriptive comment for ASM readability
     pub(crate) comment: String,
     pub(crate) owning_span: Option<Span>,
-    pub(crate) offset: u64,
 }
 
 impl Op {
@@ -941,6 +940,15 @@ impl Op {
                     VirtualOp::S256(r1, r2, r3)
                 }
                 "noop" => VirtualOp::NOOP,
+                "blob" => {
+                    let imm = check!(
+                        single_imm_24(args, immediate, whole_op_span),
+                        return err(warnings, errors),
+                        warnings,
+                        errors
+                    );
+                    VirtualOp::BLOB(imm)
+                }
                 "flag" => {
                     let r1 = check!(
                         single_reg(args, immediate, whole_op_span),
@@ -1443,6 +1451,7 @@ impl fmt::Display for VirtualOp {
             K256(a, b, c) => write!(fmtr, "k256 {} {} {}", a, b, c),
             S256(a, b, c) => write!(fmtr, "s256 {} {} {}", a, b, c),
             NOOP => Ok(()),
+            BLOB(a) => write!(fmtr, "blob {a}"),
             FLAG(a) => write!(fmtr, "flag {}", a),
             GM(a, b) => write!(fmtr, "gm {} {}", a, b),
 
@@ -1492,6 +1501,8 @@ pub(crate) enum ControlFlowOp<Reg> {
     MoveAddress(Reg, Label),
     // placeholder for the DataSection offset
     DataSectionOffsetPlaceholder,
+    // Placeholder for loading an address from the data section.
+    LoadLabel(Reg, Label),
     // Save all currently live general purpose registers, using a label as a handle.
     PushAll(Label),
     // Restore all previously saved general purpose registers.
@@ -1516,6 +1527,7 @@ impl<Reg: fmt::Display> fmt::Display for ControlFlowOp<Reg> {
                 MoveAddress(r1, lab) => format!("mova {} {}", r1, lab),
                 DataSectionOffsetPlaceholder =>
                     "DATA SECTION OFFSET[0..32]\nDATA SECTION OFFSET[32..64]".into(),
+                LoadLabel(r1, lab) => format!("lwlab {r1} {lab}"),
                 PushAll(lab) => format!("pusha {lab}"),
                 PopAll(lab) => format!("popa {lab}"),
             }
@@ -1536,7 +1548,7 @@ impl<Reg: Clone + Eq + Ord + Hash> ControlFlowOp<Reg> {
             | PopAll(_) => vec![],
 
             JumpIfNotEq(r1, r2, _) => vec![r1, r2],
-            JumpIfNotZero(r1, _) | MoveAddress(r1, _) => vec![r1],
+            JumpIfNotZero(r1, _) | MoveAddress(r1, _) | LoadLabel(r1, _) => vec![r1],
         })
         .into_iter()
         .collect()
@@ -1551,6 +1563,7 @@ impl<Reg: Clone + Eq + Ord + Hash> ControlFlowOp<Reg> {
             | Call(_)
             | MoveAddress(..)
             | DataSectionOffsetPlaceholder
+            | LoadLabel(..)
             | PushAll(_)
             | PopAll(_) => vec![],
 
@@ -1564,7 +1577,7 @@ impl<Reg: Clone + Eq + Ord + Hash> ControlFlowOp<Reg> {
     pub(crate) fn def_registers(&self) -> BTreeSet<&Reg> {
         use ControlFlowOp::*;
         (match self {
-            MoveAddress(reg, _) => vec![reg],
+            MoveAddress(reg, _) | LoadLabel(reg, _) => vec![reg],
 
             Label(_)
             | Comment
@@ -1601,6 +1614,7 @@ impl<Reg: Clone + Eq + Ord + Hash> ControlFlowOp<Reg> {
             JumpIfNotEq(r1, r2, label) => Self::JumpIfNotEq(update_reg(r1), update_reg(r2), *label),
             JumpIfNotZero(r1, label) => Self::JumpIfNotZero(update_reg(r1), *label),
             MoveAddress(r1, label) => Self::MoveAddress(update_reg(r1), *label),
+            LoadLabel(r1, label) => Self::LoadLabel(update_reg(r1), *label),
         }
     }
 
@@ -1619,6 +1633,7 @@ impl<Reg: Clone + Eq + Ord + Hash> ControlFlowOp<Reg> {
             | Call(_)
             | MoveAddress(..)
             | DataSectionOffsetPlaceholder
+            | LoadLabel(..)
             | PushAll(_)
             | PopAll(_) => (),
 
@@ -1688,6 +1703,7 @@ impl ControlFlowOp<VirtualRegister> {
             JumpIfNotEq(r1, r2, label) => JumpIfNotEq(map_reg(r1), map_reg(r2), *label),
             JumpIfNotZero(r1, label) => JumpIfNotZero(map_reg(r1), *label),
             MoveAddress(r1, label) => MoveAddress(map_reg(r1), *label),
+            LoadLabel(r1, label) => LoadLabel(map_reg(r1), *label),
         }
     }
 }

sway-core/src/asm_lang/virtual_ops.rs

@@ -140,6 +140,7 @@ pub(crate) enum VirtualOp {
     K256(VirtualRegister, VirtualRegister, VirtualRegister),
     S256(VirtualRegister, VirtualRegister, VirtualRegister),
     NOOP,
+    BLOB(VirtualImmediate24),
     FLAG(VirtualRegister),
     GM(VirtualRegister, VirtualImmediate18),
     Undefined,
@@ -227,6 +228,7 @@ impl VirtualOp {
             K256(r1, r2, r3) => vec![r1, r2, r3],
             S256(r1, r2, r3) => vec![r1, r2, r3],
             NOOP => vec![],
+            BLOB(_imm) => vec![],
             FLAG(r1) => vec![r1],
             GM(r1, _imm) => vec![r1],
             Undefined | DataSectionOffsetPlaceholder => vec![],
@@ -319,6 +321,7 @@ impl VirtualOp {
             K256(r1, r2, r3) => vec![r1, r2, r3],
             S256(r1, r2, r3) => vec![r1, r2, r3],
             NOOP => vec![],
+            BLOB(_imm) => vec![],
             FLAG(r1) => vec![r1],
             GM(_r1, _imm) => vec![],
             Undefined | DataSectionOffsetPlaceholder => vec![],
@@ -411,6 +414,7 @@ impl VirtualOp {
             K256(_r1, _r2, _r3) => vec![],
             S256(_r1, _r2, _r3) => vec![],
             NOOP => vec![],
+            BLOB(_imm) => vec![],
             FLAG(_r1) => vec![],
             GM(r1, _imm) => vec![r1],
             Undefined | DataSectionOffsetPlaceholder => vec![],
@@ -760,6 +764,7 @@ impl VirtualOp {
                 update_reg(reg_to_reg_map, r3),
             ),
             NOOP => Self::NOOP,
+            BLOB(i) => Self::BLOB(i.clone()),
             FLAG(r1) => Self::FLAG(update_reg(reg_to_reg_map, r1)),
             GM(r1, i) => Self::GM(update_reg(reg_to_reg_map, r1), i.clone()),
             Undefined => Self::Undefined,
@@ -1140,6 +1145,7 @@ impl VirtualOp {
                 map_reg(&mapping, reg3),
             ),
             NOOP => AllocatedOpcode::NOOP,
+            BLOB(imm) => AllocatedOpcode::BLOB(imm.clone()),
             FLAG(reg) => AllocatedOpcode::FLAG(map_reg(&mapping, reg)),
             GM(reg, imm) => AllocatedOpcode::GM(map_reg(&mapping, reg), imm.clone()),
             Undefined => AllocatedOpcode::Undefined,

sway-parse/src/expr/op_code.rs

@@ -45,6 +45,7 @@ define_op_codes!(
     (Mul, MulOpcode, "mul", (ret, lhs, rhs)),
     (Muli, MuliOpcode, "muli", (ret, lhs, rhs)),
     (Noop, NoopOpcode, "noop", ()),
+    (Blob, BlobOpcode, "blob", (size)),
     (Not, NotOpcode, "not", (ret, arg)),
     (Or, OrOpcode, "or", (ret, lhs, rhs)),
     (Ori, OriOpcode, "ori", (ret, lhs, rhs)),

test/src/e2e_vm_tests/harness.rs

@@ -87,7 +87,11 @@ pub(crate) fn runs_in_vm(
     let maturity = 1;
     let script_data = script_data.unwrap_or_default();
     let block_height = (u32::MAX >> 1) as u64;
-    let params = &ConsensusParameters::DEFAULT;
+    let params = &ConsensusParameters {
+        // The default max length is 1MB which isn't enough for the bigger tests.
+        max_script_length: 64 * 1024 * 1024,
+        ..ConsensusParameters::DEFAULT
+    };
 
     let tx = TransactionBuilder::script(script.bytecode.clone(), script_data)
         .add_unsigned_coin_input(rng.gen(), rng.gen(), 1, Default::default(), rng.gen(), 0)

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

@@ -0,0 +1,3 @@
+[[package]]
+name = 'far_jumps'
+source = 'root'

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

@@ -0,0 +1,6 @@
+[project]
+authors = ["Fuel Labs <[email protected]>"]
+entry = "main.sw"
+license = "Apache-2.0"
+name = "far_jumps"
+implicit-std = false

test/src/e2e_vm_tests/test_programs/should_pass/language/far_jumps/src/main.sw

@@ -0,0 +1,18 @@
+script;
+
+fn main() -> u64 {
+    asm() {
+        blob i262144;
+    }
+    if t() {
+        111
+    } else {
+        222
+    }
+}
+
+fn t() -> bool {
+    asm() {
+        one: bool
+    }
+}

test/src/e2e_vm_tests/test_programs/should_pass/language/far_jumps/test.toml

@@ -0,0 +1,3 @@
+category = "run"
+expected_result = { action = "return", value = 111 }
+validate_abi = false