Disallow control flow opcodes in asm blocks (FuelLabs#3279)

sway-ast/src/expr/op_code.rs

@@ -178,12 +178,12 @@ define_op_codes!(
     (Xor, XorOpcode, "xor", (ret: reg, lhs: reg, rhs: reg)),
     (Xori, XoriOpcode, "xori", (ret: reg, lhs: reg, rhs: imm)),
     /* Control Flow Instructions */
-    // (Jmp, JmpOpcode, "jmp", (offset: imm)),
-    // (Ji, JiOpcode, "ji", (offset: imm)),
-    // (JNE, JneOpcode, "jne", (lhs: reg, rhs: reg, offset: imm)),
-    // (Jnei, JneiOpcode, "jnei", (lhs: reg, rhs: reg, offset: imm)),
-    // (JNZI, JnziOpcode, "jnzi", (arg: reg, offset: imm)),
-    // (RET, RetOpcode, "ret", (value: reg)),
+    (Jmp, JmpOpcode, "jmp", (offset: reg)),
+    (Ji, JiOpcode, "ji", (offset: imm)),
+    (Jne, JneOpcode, "jne", (lhs: reg, rhs: reg, offset: reg)),
+    (Jnei, JneiOpcode, "jnei", (lhs: reg, rhs: reg, offset: imm)),
+    (Jnzi, JnziOpcode, "jnzi", (arg: reg, offset: imm)),
+    (Ret, RetOpcode, "ret", (value: reg)),
     /* Memory Instructions */
     (Aloc, AlocOpcode, "aloc", (size: reg)),
     (Cfei, CfeiOpcode, "cfei", (size: imm)),
@@ -246,8 +246,8 @@ define_op_codes!(
         (reg_a: reg, reg_b: reg, addr: reg, size: reg)
     ),
     (Mint, MintOpcode, "mint", (coins: reg)),
-    // (Retd, RetdOpcode, "retd", (addr: reg, size: reg)),
-    // (Rvrt, RvrtOpcode, "rvrt", (value: reg)),
+    (Retd, RetdOpcode, "retd", (addr: reg, size: reg)),
+    (Rvrt, RvrtOpcode, "rvrt", (value: reg)),
     (
         Smo,
         SmoOpcode,

sway-core/src/semantic_analysis/ast_node/expression/typed_expression.rs

@@ -15,7 +15,7 @@ pub(crate) use self::{
 };
 
 use crate::{
-    asm_lang::virtual_register::VirtualRegister,
+    asm_lang::{virtual_ops::VirtualOp, virtual_register::VirtualRegister},
     declaration_engine::declaration_engine::*,
     error::*,
     language::{parsed::*, ty, *},
@@ -675,6 +675,18 @@ impl ty::TyExpression {
     ) -> CompileResult<ty::TyExpression> {
         let mut warnings = vec![];
         let mut errors = vec![];
+
+        // Various checks that we can catch early to check that the assembly is valid. For now,
+        // this includes two checks:
+        // 1. Check that no control flow opcodes are used.
+        // 2. Check that initialized registers are not reassigned in the `asm` block.
+        check!(
+            check_asm_block_validity(&asm),
+            return err(warnings, errors),
+            warnings,
+            errors
+        );
+
         let asm_span = asm
             .returns
             .clone()
@@ -716,14 +728,6 @@ impl ty::TyExpression {
             )
             .collect();
 
-        // Make sure that all registers that are initialized are *not* assigned again.
-        check!(
-            disallow_assigning_initialized_registers(&asm),
-            return err(warnings, errors),
-            warnings,
-            errors
-        );
-
         let exp = ty::TyExpression {
             expression: ty::TyExpressionVariant::AsmExpression {
                 whole_block_span: asm.whole_block_span,
@@ -1949,18 +1953,10 @@ mod tests {
     }
 }
 
-fn disallow_assigning_initialized_registers(asm: &AsmExpression) -> CompileResult<()> {
+fn check_asm_block_validity(asm: &AsmExpression) -> CompileResult<()> {
     let mut errors = vec![];
     let mut warnings = vec![];
 
-    // Collect all registers that have initializers in the list of arguments
-    let initialized_registers = asm
-        .registers
-        .iter()
-        .filter(|reg| reg.initializer.is_some())
-        .map(|reg| VirtualRegister::Virtual(reg.name.to_string()))
-        .collect::<FxHashSet<_>>();
-
     // Collect all asm block instructions in the form of `VirtualOp`s
     let mut opcodes = vec![];
     for op in &asm.body {
@@ -1970,35 +1966,74 @@ fn disallow_assigning_initialized_registers(asm: &AsmExpression) -> CompileResul
             .map(|reg_name| VirtualRegister::Virtual(reg_name.to_string()))
             .collect::<Vec<VirtualRegister>>();
 
-        opcodes.push(check!(
-            crate::asm_lang::Op::parse_opcode(
-                &op.op_name,
-                &registers,
-                &op.immediate,
-                op.span.clone(),
+        opcodes.push((
+            check!(
+                crate::asm_lang::Op::parse_opcode(
+                    &op.op_name,
+                    &registers,
+                    &op.immediate,
+                    op.span.clone(),
+                ),
+                return err(warnings, errors),
+                warnings,
+                errors
             ),
-            return err(warnings, errors),
-            warnings,
-            errors
+            op.op_name.clone(),
+            op.span.clone(),
         ));
     }
 
-    // From the list of `VirtualOp`s, figure out what registers are assigned
+    // Check #1: Disallow control flow instructions
+    //
+    errors.extend(
+        opcodes
+            .iter()
+            .filter(|op| {
+                matches!(
+                    op.0,
+                    VirtualOp::JMP(_)
+                        | VirtualOp::JI(_)
+                        | VirtualOp::JNE(..)
+                        | VirtualOp::JNEI(..)
+                        | VirtualOp::JNZI(..)
+                        | VirtualOp::RET(_)
+                        | VirtualOp::RETD(..)
+                        | VirtualOp::RVRT(..)
+                )
+            })
+            .map(|op| CompileError::DisallowedControlFlowInstruction {
+                name: op.1.to_string(),
+                span: op.2.clone(),
+            })
+            .collect::<Vec<_>>(),
+    );
+
+    // Check #2: Disallow initialized registers from being reassigned in the asm block
+    //
+    // 1. Collect all registers that have initializers in the list of arguments
+    let initialized_registers = asm
+        .registers
+        .iter()
+        .filter(|reg| reg.initializer.is_some())
+        .map(|reg| VirtualRegister::Virtual(reg.name.to_string()))
+        .collect::<FxHashSet<_>>();
+
+    // 2. From the list of `VirtualOp`s, figure out what registers are assigned
     let assigned_registers: FxHashSet<VirtualRegister> =
         opcodes.iter().fold(FxHashSet::default(), |mut acc, op| {
-            for u in op.def_registers() {
+            for u in op.0.def_registers() {
                 acc.insert(u.clone());
             }
             acc
         });
 
-    // Intersect the list of assigned registers with the list of initialized registers
+    // 3. Intersect the list of assigned registers with the list of initialized registers
     let initialized_and_assigned_registers = assigned_registers
         .intersection(&initialized_registers)
         .collect::<FxHashSet<_>>();
 
-    // Form all the compile errors given the violating registers above. Obtain span information
-    // from the original `asm.registers` vector.
+    // 4. Form all the compile errors given the violating registers above. Obtain span information
+    //    from the original `asm.registers` vector.
     errors.extend(
         asm.registers
             .iter()

sway-error/src/error.rs

@@ -358,18 +358,6 @@ pub enum CompileError {
     Immediate18TooLarge { val: u64, span: Span },
     #[error("The value \"{val}\" is too large to fit in this 24-bit immediate spot.")]
     Immediate24TooLarge { val: u64, span: Span },
-    #[error(
-        "The opcode \"ji\" is not valid in inline assembly. Try using function calls instead."
-    )]
-    DisallowedJi { span: Span },
-    #[error("The opcode \"jnei\" is not valid in inline assembly. Use an enclosing if expression instead.")]
-    DisallowedJnei { span: Span },
-    #[error("The opcode \"jnzi\" is not valid in inline assembly. Use an enclosing if expression instead.")]
-    DisallowedJnzi { span: Span },
-    #[error(
-        "The opcode \"lw\" is not valid in inline assembly. Try assigning a static value to a variable instead."
-    )]
-    DisallowedLw { span: Span },
     #[error(
         "This op expects {expected} register(s) as arguments, but you provided {received} register(s)."
     )]
@@ -675,6 +663,8 @@ pub enum CompileError {
             Consider assigning to a different register inside the ASM block."
     )]
     InitializedRegisterReassignment { name: String, span: Span },
+    #[error("Control flow VM instructions are not allowed in assembly blocks.")]
+    DisallowedControlFlowInstruction { name: String, span: Span },
 }
 
 impl std::convert::From<TypeError> for CompileError {
@@ -758,10 +748,6 @@ impl Spanned for CompileError {
             Immediate12TooLarge { span, .. } => span.clone(),
             Immediate18TooLarge { span, .. } => span.clone(),
             Immediate24TooLarge { span, .. } => span.clone(),
-            DisallowedJi { span, .. } => span.clone(),
-            DisallowedJnei { span, .. } => span.clone(),
-            DisallowedJnzi { span, .. } => span.clone(),
-            DisallowedLw { span, .. } => span.clone(),
             IncorrectNumberOfAsmRegisters { span, .. } => span.clone(),
             UnnecessaryImmediate { span, .. } => span.clone(),
             AmbiguousPath { span, .. } => span.clone(),
@@ -857,6 +843,7 @@ impl Spanned for CompileError {
             RefMutableNotAllowedInMain { param_name } => param_name.span(),
             PointerReturnNotAllowedInMain { span } => span.clone(),
             InitializedRegisterReassignment { span, .. } => span.clone(),
+            DisallowedControlFlowInstruction { span, .. } => span.clone(),
         }
     }
 }

sway-parse/src/expr/op_code.rs

@@ -57,12 +57,12 @@ define_op_codes!(
     (Xor, XorOpcode, "xor", (ret, lhs, rhs)),
     (Xori, XoriOpcode, "xori", (ret, lhs, rhs)),
     /* Control Flow Instructions */
-    // (Jmp, JmpOpcode, "jmp", (offset)),
-    // (JI, JiOpcode, "ji", (offset)),
-    // (JNE, JneOpcode, "jne", (lhs, rhs, offset)),
-    // (JNEI, JneiOpcode, "jnei", (lhs, rhs, offset)),
-    // (JNZI, JnziOpcode, "jnzi", (arg, offset)),
-    // (RET, RetOpcode, "ret", (value)),
+    (Jmp, JmpOpcode, "jmp", (offset)),
+    (Ji, JiOpcode, "ji", (offset)),
+    (Jne, JneOpcode, "jne", (lhs, rhs, offset)),
+    (Jnei, JneiOpcode, "jnei", (lhs, rhs, offset)),
+    (Jnzi, JnziOpcode, "jnzi", (arg, offset)),
+    (Ret, RetOpcode, "ret", (value)),
     /* Memory Instructions */
     (Aloc, AlocOpcode, "aloc", (size)),
     (Cfei, CfeiOpcode, "cfei", (size)),
@@ -90,8 +90,8 @@ define_op_codes!(
     (Log, LogOpcode, "log", (reg_a, reg_b, reg_c, reg_d)),
     (Logd, LogdOpcode, "logd", (reg_a, reg_b, addr, size)),
     (Mint, MintOpcode, "mint", (coins)),
-    // (Retd, RetdOpcode, "retd", (addr, size)),
-    // (Rvrt, RvrtOpcode, "rvrt", (value)),
+    (Retd, RetdOpcode, "retd", (addr, size)),
+    (Rvrt, RvrtOpcode, "rvrt", (value)),
     (Smo, SmoOpcode, "smo", (addr, len, output, coins)),
     (Scwq, ScwqOpcode, "scwq", (addr, is_set, len)),
     (Srw, SrwOpcode, "srw", (ret, is_set, state_addr)),

test/src/e2e_vm_tests/test_programs/should_fail/asm_disallowed_opcodes/Forc.lock

@@ -0,0 +1,3 @@
+[[package]]
+name = 'asm_disallowed_opcodes'
+source = 'member'

test/src/e2e_vm_tests/test_programs/should_fail/asm_disallowed_opcodes/Forc.toml

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

test/src/e2e_vm_tests/test_programs/should_fail/asm_disallowed_opcodes/src/main.sw

@@ -0,0 +1,15 @@
+script;
+
+fn main() -> u64 {
+    asm(r1: 0, r2: 0, r3, r4) {
+        jmp r1;
+        ji i5;
+        jne r1 r2 r3;
+        jnei r1 r2 i5;
+        jnzi r1 i5;
+        ret r1;
+        retd r1 r2;
+        rvrt r1;
+    };
+    0
+}

test/src/e2e_vm_tests/test_programs/should_fail/asm_disallowed_opcodes/test.toml

@@ -0,0 +1,32 @@
+category = "fail"
+
+# check: $()jmp r1
+# nextln: $()Control flow VM instructions are not allowed in assembly blocks.
+# nextln: $()ji i5
+
+# check: $()ji i5
+# nextln: $()Control flow VM instructions are not allowed in assembly blocks.
+# nextln: $()jne r1 r2 r3;
+
+# check: $()jne r1 r2 r3;
+# nextln: $()Control flow VM instructions are not allowed in assembly blocks.
+# nextln: $()jnei r1 r2 i5;
+
+# check: $()jnei r1 r2 i5;
+# nextln: $()Control flow VM instructions are not allowed in assembly blocks.
+# nextln: $()jnzi r1 i5;
+
+# check: $()jnzi r1 i5;
+# nextln: $()Control flow VM instructions are not allowed in assembly blocks.
+# nextln: $()ret r1;
+
+# check: $()ret r1;
+# nextln: $()Control flow VM instructions are not allowed in assembly blocks.
+# nextln: $()retd r1 r2;
+
+# check: $()retd r1 r2;
+# nextln: $()Control flow VM instructions are not allowed in assembly blocks.
+# nextln: $()rvrt r1;
+
+# check: $()rvrt r1;
+# nextln: $()Control flow VM instructions are not allowed in assembly blocks.