Disallow assigning to initialized registers in asm blocks (FuelLabs…

…#3239) Closes FuelLabs#3155 This is needed due to `mem2reg` and some unclear semantics around initialized registers being passed as ref/mut ref. Besides, it's a good idea to restrict `asm` blocks as much as possible to avoid weird unwanted behavior. <img width="883" alt="image" src="https://user-images.githubusercontent.com/59666792/199332120-974d1821-680e-43a1-922d-1a32d491b3cd.png"> * I also removed the `disallow_opcodes` check because the parser takes care of that now. The parser only allows legal opcodes in `asm` blocks. * Finally, I update the parser to accept the correct storage opcodes with `fuel-core 0.13` and updates all the tests. Co-authored-by: Vaivaswatha N <[email protected]>
vimday · Nov 2, 2022 · b83ad9b · b83ad9b
1 parent cfecaf4
commit b83ad9b
Show file tree

Hide file tree

Showing 16 changed files with 163 additions and 56 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/sway-ast/src/expr/op_code.rs b/sway-ast/src/expr/op_code.rs
@@ -264,10 +264,30 @@ define_op_codes!(
         "sldc",
         (contract: reg, addr: reg, size: reg)
     ),
-    (Srw, SrwOpcode, "srw", (ret: reg, state_addr: reg)),
-    (Srwq, SrwqOpcode, "srwq", (addr: reg, state_addr: reg)),
-    (Sww, SwwOpcode, "sww", (state_addr: reg, value: reg)),
-    (Swwq, SwwqOpcode, "swwq", (state_addr: reg, addr: reg)),
+    (
+        Srw,
+        SrwOpcode,
+        "srw",
+        (ret: reg, is_set: reg, state_addr: reg)
+    ),
+    (
+        Srwq,
+        SrwqOpcode,
+        "srwq",
+        (addr: reg, is_set: reg, state_addr: reg, count: reg)
+    ),
+    (
+        Sww,
+        SwwOpcode,
+        "sww",
+        (state_addr: reg, is_set: reg, value: reg)
+    ),
+    (
+        Swwq,
+        SwwqOpcode,
+        "swwq",
+        (state_addr: reg, is_set: reg, addr: reg, count: reg)
+    ),
     (Time, TimeOpcode, "time", (ret: reg, height: reg)),
     (Tr, TrOpcode, "tr", (contract: reg, coins: reg, asset: reg)),
     (

diff --git a/sway-core/Cargo.toml b/sway-core/Cargo.toml
@@ -19,6 +19,7 @@ im = "15.0"
 itertools = "0.10"
 lazy_static = "1.4"
 petgraph = "0.6"
+rustc-hash = "1.1.0"
 serde = { version = "1.0", features = ["derive"] }
 sha2 = "0.9"
 smallvec = "1.7"

diff --git a/sway-core/src/semantic_analysis/ast_node/expression/typed_expression.rs b/sway-core/src/semantic_analysis/ast_node/expression/typed_expression.rs
@@ -15,6 +15,7 @@ pub(crate) use self::{
 };
 
 use crate::{
+    asm_lang::virtual_register::VirtualRegister,
     declaration_engine::declaration_engine::*,
     error::*,
     language::{parsed::*, ty, *},
@@ -31,6 +32,8 @@ use sway_error::{
 };
 use sway_types::{integer_bits::IntegerBits, Ident, Span, Spanned};
 
+use rustc_hash::FxHashSet;
+
 use std::collections::{HashMap, VecDeque};
 
 #[allow(clippy::too_many_arguments)]
@@ -670,6 +673,7 @@ impl ty::TyExpression {
         // type check the initializers
         let typed_registers = asm
             .registers
+            .clone()
             .into_iter()
             .map(
                 |AsmRegisterDeclaration { name, initializer }| ty::TyAsmRegisterDeclaration {
@@ -689,10 +693,15 @@ impl ty::TyExpression {
                 },
             )
             .collect();
-        // check for any disallowed opcodes
-        for op in &asm.body {
-            check!(disallow_opcode(&op.op_name), continue, warnings, errors)
-        }
+
+        // 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,
@@ -1920,24 +1929,70 @@ mod tests {
         assert!(comp_res.warnings.is_empty() && comp_res.errors.is_empty());
     }
 }
-fn disallow_opcode(op: &Ident) -> CompileResult<()> {
-    let mut errors = vec![];
 
-    match op.as_str().to_lowercase().as_str() {
-        "ji" => {
-            errors.push(CompileError::DisallowedJi { span: op.span() });
-        }
-        "jnei" => {
-            errors.push(CompileError::DisallowedJnei { span: op.span() });
-        }
-        "jnzi" => {
-            errors.push(CompileError::DisallowedJnzi { span: op.span() });
-        }
-        _ => (),
-    };
-    if errors.is_empty() {
-        ok((), vec![], vec![])
-    } else {
-        err(vec![], errors)
+fn disallow_assigning_initialized_registers(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 {
+        let registers = op
+            .op_args
+            .iter()
+            .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(),
+            ),
+            return err(warnings, errors),
+            warnings,
+            errors
+        ));
     }
+
+    // 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() {
+                acc.insert(u.clone());
+            }
+            acc
+        });
+
+    // 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.
+    errors.extend(
+        asm.registers
+            .iter()
+            .filter(|reg| {
+                initialized_and_assigned_registers
+                    .contains(&VirtualRegister::Virtual(reg.name.to_string()))
+            })
+            .map(|reg| CompileError::InitializedRegisterReassignment {
+                name: reg.name.to_string(),
+                span: reg.name.span(),
+            })
+            .collect::<Vec<_>>(),
+    );
+
+    ok((), vec![], errors)
 }
diff --git a/sway-error/src/error.rs b/sway-error/src/error.rs
@@ -670,6 +670,11 @@ pub enum CompileError {
     RefMutableNotAllowedInMain { param_name: Ident },
     #[error("returning a `raw_ptr` from `main()` is not allowed")]
     PointerReturnNotAllowedInMain { span: Span },
+    #[error(
+        "Register \"{name}\" is initialized and later reassigned which is not allowed. \
+            Consider assigning to a different register inside the ASM block."
+    )]
+    InitializedRegisterReassignment { name: String, span: Span },
 }
 
 impl std::convert::From<TypeError> for CompileError {
@@ -851,6 +856,7 @@ impl Spanned for CompileError {
             ConfigTimeConstantNotALiteral { span } => span.clone(),
             RefMutableNotAllowedInMain { param_name } => param_name.span(),
             PointerReturnNotAllowedInMain { span } => span.clone(),
+            InitializedRegisterReassignment { span, .. } => span.clone(),
         }
     }
 }

diff --git a/sway-lib-core/src/ops.sw b/sway-lib-core/src/ops.sw
@@ -623,8 +623,7 @@ fn rsh_with_carry(word: u64, shift_amount: u64) -> (u64, u64) {
 
 /// Extract a single 64 bit word from a b256 value using the specified offset.
 fn get_word_from_b256(val: b256, offset: u64) -> u64 {
-    let mut empty: u64 = 0;
-    asm(r1: val, offset: offset, r2, res: empty) {
+    asm(r1: val, offset: offset, r2, res) {
         add r2 r1 offset;
         lw res r2 i0;
         res: u64

diff --git a/sway-parse/src/expr/op_code.rs b/sway-parse/src/expr/op_code.rs
@@ -83,10 +83,10 @@ 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)),
-    (Srw, SrwOpcode, "srw", (ret, state_addr)),
-    (Srwq, SrwqOpcode, "srwq", (addr, state_addr)),
-    (Sww, SwwOpcode, "sww", (state_addr, value)),
-    (Swwq, SwwqOpcode, "swwq", (state_addr, addr)),
+    (Srw, SrwOpcode, "srw", (ret, is_set, state_addr)),
+    (Srwq, SrwqOpcode, "srwq", (addr, is_set, state_addr, count)),
+    (Sww, SwwOpcode, "sww", (state_addr, is_set, value)),
+    (Swwq, SwwqOpcode, "swwq", (state_addr, is_set, addr, count)),
     (Time, TimeOpcode, "time", (ret, height)),
     (Tr, TrOpcode, "tr", (contract, coins, asset)),
     (Tro, TroOpcode, "tro", (addr, output, coins, asset)),

diff --git a/test/src/e2e_vm_tests/test_programs/should_fail/asm_assign_to_initialized_reg/Forc.lock b/test/src/e2e_vm_tests/test_programs/should_fail/asm_assign_to_initialized_reg/Forc.lock
@@ -0,0 +1,3 @@
+[[package]]
+name = 'asm_assign_to_initialized_reg'
+source = 'root'
diff --git a/test/src/e2e_vm_tests/test_programs/should_fail/asm_assign_to_initialized_reg/Forc.toml b/test/src/e2e_vm_tests/test_programs/should_fail/asm_assign_to_initialized_reg/Forc.toml
@@ -0,0 +1,6 @@
+[project]
+authors = ["Fuel Labs <[email protected]>"]
+entry = "main.sw"
+implicit-std = false
+license = "Apache-2.0"
+name = "asm_assign_to_initialized_reg"
diff --git a/test/src/e2e_vm_tests/test_programs/should_fail/asm_assign_to_initialized_reg/src/main.sw b/test/src/e2e_vm_tests/test_programs/should_fail/asm_assign_to_initialized_reg/src/main.sw
@@ -0,0 +1,11 @@
+script;
+
+fn main() -> u64 {
+    asm(r1: 0, r2: 0, r3, r4) {
+        move r1 r2;
+        move r2 r1;
+        move r3 r1;
+        move r4 r2;
+    };
+    0
+}
diff --git a/test/src/e2e_vm_tests/test_programs/should_fail/asm_assign_to_initialized_reg/test.toml b/test/src/e2e_vm_tests/test_programs/should_fail/asm_assign_to_initialized_reg/test.toml
@@ -0,0 +1,7 @@
+category = "fail"
+
+# check: $()asm(r1: 0, r2: 0, r3, r4) {
+# nextln: $()Register "r1" is initialized and later reassigned which is not allowed. Consider assigning to a different register inside the ASM block.
+
+# check: $()asm(r1: 0, r2: 0, r3, r4) {
+# nextln: $()Register "r2" is initialized and later reassigned which is not allowed. Consider assigning to a different register inside the ASM block.
diff --git a/test/src/e2e_vm_tests/test_programs/should_fail/storage_conflict/src/main.sw b/test/src/e2e_vm_tests/test_programs/should_fail/storage_conflict/src/main.sw
@@ -101,35 +101,35 @@ fn do_pure_stuff_e() -> bool {
 const KEY: b256 = 0xfefefefefefefefefefefefefefefefefefefefefefefefefefefefefefefefe;
 
 fn read_storage_word() -> u64 {
-    asm (key: KEY, res, a, b, c) {
+    asm (key: KEY, is_set, res, a, b, c) {
         move a b;
-        srw res key;
+        srw res is_set key;
         addi b c i1;
         res: u64
     }
 }
 
 fn read_storage_b256() -> b256 {
     let res = ZERO_B256;
-    asm (key: KEY, buf: res, a, b, c) {
+    asm (key: KEY, is_set, buf: res, count: 1, a, b, c) {
         modi a b i10;
         lt a b c;
-        srwq buf key;
+        srwq buf is_set key count;
         and a b c;
     }
     res
 }
 
 fn write_storage_word() {
-    asm (key: KEY, val: 42) {
-        sww key val;
+    asm (key: KEY, is_set, val: 42) {
+        sww key is_set val;
     }
 }
 
 fn write_storage_b256() {
     let val = 0xbabababababababababababababababababababababababababababababababa;
-    asm (key: KEY, val: val) {
-        swwq key val;
+    asm (key: KEY, is_set, val: val, count: 1) {
+        swwq key is_set val count;
     }
 }
 

diff --git a/test/src/e2e_vm_tests/test_programs/should_fail/storage_ops_in_library/src/do_storage.sw b/test/src/e2e_vm_tests/test_programs/should_fail/storage_ops_in_library/src/do_storage.sw
@@ -4,8 +4,8 @@ const KEY: b256 = 0xfafafafafafafafafafafafafafafafafafafafafafafafafafafafafafa
 
 #[storage(read, write)]
 pub fn side_effects() {
-    asm(key: KEY, v) {
-        srw v key;
-        sww key v;
+    asm(key: KEY, is_set, v) {
+        srw v is_set key;
+        sww key is_set v;
     }
 }
diff --git a/test/src/e2e_vm_tests/test_programs/should_pass/language/asm_without_return/Forc.toml b/test/src/e2e_vm_tests/test_programs/should_pass/language/asm_without_return/Forc.toml
@@ -1,8 +1,8 @@
 [project]
 authors = ["Fuel Labs <[email protected]>"]
+entry = "main.sw"
 license = "Apache-2.0"
 name = "asm_without_return"
-entry = "main.sw"
 
 [dependencies]
 core = { path = "../../../../../../../sway-lib-core" }
diff --git a/test/src/e2e_vm_tests/test_programs/should_pass/language/asm_without_return/src/main.sw b/test/src/e2e_vm_tests/test_programs/should_pass/language/asm_without_return/src/main.sw
@@ -2,12 +2,10 @@ script;
 
 fn main() {
     asm() {
-
     };
 
-    asm(r1: 5, r2: 5) {
-        add r1 r1 r2;
-        add r2 r2 r2;
+    asm(r1: 5, r2: 5, r3, r4) {
+        add r3 r1 r2;
+        add r4 r2 r2;
     };
-
 }
diff --git a/test/src/ir_generation/tests/storage_metadata.sw b/test/src/ir_generation/tests/storage_metadata.sw
@@ -16,27 +16,27 @@ abi Incrementor {
 impl Incrementor for Contract {
     #[storage(write)]
     fn initialize(initial_value: u64) -> u64 {
-        asm(key: KEY, v: initial_value) {
-            sww key v;
+        asm(key: KEY, is_set, v: initial_value) {
+            sww key is_set v;
         }
         initial_value
     }
 
     #[storage(read, write)]
     fn increment(increment_by: u64) -> u64 {
-        let new_val = asm(key: KEY, i: increment_by, res) {
-            srw res key;
+        let new_val = asm(key: KEY, is_set, i: increment_by, res) {
+            srw res is_set key;
             add res res i;
-            sww key res;
+            sww key is_set res;
             res: u64
         };
         new_val
     }
 
     #[storage(read)]
     fn get() -> u64 {
-        asm(key: KEY, res) {
-            srw key res;
+        asm(key: KEY, is_set, res) {
+            srw res is_set key;
             res: u64
         }
     }