Allow script main fns to have args (FuelLabs#2592)

* Allow script and predicate main fns to receive args

* Update sway-core/src/error.rs

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

* error fix

* add test for various types

* move failing test

* fix path

* remove length check

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

sway-core/src/asm_generation/from_ir.rs

@@ -7,6 +7,7 @@
 // But this is not ideal and needs to be refactored:
 // - AsmNamespace is tied to data structures from other stages like Ident and Literal.
 
+use fuel_asm::GTFArgs;
 use fuel_crypto::Hasher;
 use std::{collections::HashMap, sync::Arc};
 
@@ -237,10 +238,8 @@ impl<'ir> AsmBuilder<'ir> {
 
     // Handle loading the arguments of a contract call
     fn compile_fn_args(&mut self, function: Function) {
-        // Do this only for contract methods. Contract methods have selectors.
-        if !function.has_selector(self.context) {
-            return;
-        }
+        // We treat contract methods differently. Contract methods have selectors.
+        let is_contract_method = function.has_selector(self.context);
 
         match function.args_iter(self.context).count() {
             // Nothing to do if there are no arguments
@@ -251,13 +250,35 @@ impl<'ir> AsmBuilder<'ir> {
             1 => {
                 let (_, val) = function.args_iter(self.context).next().unwrap();
                 let single_arg_reg = self.value_to_register(val);
-                self.read_args_value_from_frame(&single_arg_reg);
+
+                if is_contract_method {
+                    self.read_args_value_from_frame(&single_arg_reg);
+                } else {
+                    self.read_args_value_from_script_data(&single_arg_reg);
+
+                    if val.get_type(self.context).unwrap().is_copy_type() {
+                        self.bytecode.push(Op {
+                            opcode: either::Either::Left(VirtualOp::LW(
+                                single_arg_reg.clone(),
+                                single_arg_reg.clone(),
+                                VirtualImmediate12 { value: 0 },
+                            )),
+                            comment: "Load main fn parameter".into(),
+                            owning_span: None,
+                        });
+                    }
+                }
             }
 
             // Otherwise, the args are bundled together and pointed to by the base register.
             _ => {
                 let args_base_reg = self.reg_seqr.next();
-                self.read_args_value_from_frame(&args_base_reg);
+
+                if is_contract_method {
+                    self.read_args_value_from_frame(&args_base_reg);
+                } else {
+                    self.read_args_value_from_script_data(&args_base_reg);
+                }
 
                 // Successively load each argument. The asm generated depends on the arg type size
                 // and whether the offset fits in a 12-bit immediate.
@@ -346,6 +367,21 @@ impl<'ir> AsmBuilder<'ir> {
         });
     }
 
+    // Read the argument(s) base from the script data.
+    fn read_args_value_from_script_data(&mut self, reg: &VirtualRegister) {
+        self.bytecode.push(Op {
+            opcode: either::Either::Left(VirtualOp::GTF(
+                reg.clone(),
+                VirtualRegister::Constant(ConstantRegister::Zero),
+                VirtualImmediate12 {
+                    value: GTFArgs::ScriptData as u16,
+                },
+            )),
+            comment: "Base register for main fn parameter".into(),
+            owning_span: None,
+        });
+    }
+
     fn add_locals(&mut self, function: Function) {
         // If they're immutable and have a constant initialiser then they go in the data section.
         // Otherwise they go in runtime allocated space, either a register or on the stack.

sway-core/src/error.rs

@@ -1066,12 +1066,12 @@ pub enum CompileError {
     BreakOutsideLoop { span: Span },
     #[error("\"continue\" used outside of a loop")]
     ContinueOutsideLoop { span: Span },
-    #[error("arguments to \"main()\" are not yet supported. See the issue here: github.com/FuelLabs/sway/issues/845")]
-    MainArgsNotYetSupported { span: Span },
     #[error("Configuration-time constant value is not a constant item.")]
     ConfigTimeConstantNotAConstDecl { span: Span },
     #[error("Configuration-time constant value is not a literal.")]
     ConfigTimeConstantNotALiteral { span: Span },
+    #[error("ref mut parameter not allowed for main()")]
+    RefMutableNotAllowedInMain { param_name: Ident },
 }
 
 impl std::convert::From<TypeError> for CompileError {
@@ -1243,9 +1243,9 @@ impl Spanned for CompileError {
             IntrinsicIncorrectNumTArgs { span, .. } => span.clone(),
             BreakOutsideLoop { span } => span.clone(),
             ContinueOutsideLoop { span } => span.clone(),
-            MainArgsNotYetSupported { span } => span.clone(),
             ConfigTimeConstantNotAConstDecl { span } => span.clone(),
             ConfigTimeConstantNotALiteral { span } => span.clone(),
+            RefMutableNotAllowedInMain { param_name } => param_name.span(),
         }
     }
 }

sway-core/src/semantic_analysis/program.rs

@@ -196,14 +196,16 @@ impl TypedProgram {
                 }
             }
         };
-        // check if no arguments passed to a `main()` in a `script` or `predicate`.
+        // check if no ref mut arguments passed to a `main()` in a `script` or `predicate`.
         match &typed_program_kind {
             TypedProgramKind::Script { main_function, .. }
             | TypedProgramKind::Predicate { main_function, .. } => {
-                if !main_function.parameters.is_empty() {
-                    errors.push(CompileError::MainArgsNotYetSupported {
-                        span: main_function.span.clone(),
-                    })
+                for param in &main_function.parameters {
+                    if param.is_reference && param.is_mutable {
+                        errors.push(CompileError::RefMutableNotAllowedInMain {
+                            param_name: param.name.clone(),
+                        })
+                    }
                 }
             }
             _ => (),
@@ -362,7 +364,8 @@ impl TypedProgramKind {
                     functions: result,
                 }
             }
-            TypedProgramKind::Script { main_function, .. } => {
+            TypedProgramKind::Script { main_function, .. }
+            | TypedProgramKind::Predicate { main_function, .. } => {
                 let result = vec![main_function.generate_json_abi_function(types)];
                 JsonABIProgram {
                     types: types.to_vec(),

test/Cargo.toml

@@ -17,6 +17,7 @@ fuel-asm = "0.8"
 fuel-tx = { version = "0.18", features = ["builder"]}
 fuel-vm = { version = "0.15", features = ["random"] }
 gag = "1.0"
+hex = "0.4.3"
 prettydiff = "0.6"
 rand = "0.8"
 regex = "1"

test/src/e2e_vm_tests/harness.rs

@@ -71,23 +71,29 @@ pub(crate) fn runs_on_node(
 
 /// Very basic check that code does indeed run in the VM.
 /// `true` if it does, `false` if not.
-pub(crate) fn runs_in_vm(file_name: &str, locked: bool) -> (ProgramState, Compiled) {
+pub(crate) fn runs_in_vm(
+    file_name: &str,
+    script_data: Option<Vec<u8>>,
+    locked: bool,
+) -> (ProgramState, Compiled) {
     let storage = MemoryStorage::default();
 
     let rng = &mut StdRng::seed_from_u64(2322u64);
     let script = compile_to_bytes(file_name, locked).unwrap();
     let maturity = 1;
+    let script_data = script_data.unwrap_or_default();
     let block_height = (u32::MAX >> 1) as u64;
     let params = &ConsensusParameters::DEFAULT;
 
-    let tx = TransactionBuilder::script(script.bytecode.clone(), Default::default())
+    let tx = TransactionBuilder::script(script.bytecode.clone(), script_data)
         .add_unsigned_coin_input(rng.gen(), rng.gen(), 1, Default::default(), rng.gen(), 0)
         .gas_limit(fuel_tx::ConsensusParameters::DEFAULT.max_gas_per_tx)
         .maturity(maturity)
         .finalize_checked(block_height as Word, params);
 
     let mut i = Interpreter::with_storage(storage, Default::default());
-    (*i.transact(tx).unwrap().state(), script)
+    let transition = i.transact(tx).unwrap();
+    (*transition.state(), script)
 }
 
 /// Compiles the code and captures the output of forc and the compilation.

test/src/e2e_vm_tests/mod.rs

@@ -31,6 +31,7 @@ enum TestResult {
 struct TestDescription {
     name: String,
     category: TestCategory,
+    script_data: Option<Vec<u8>>,
     expected_result: Option<TestResult>,
     contract_paths: Vec<String>,
     validate_abi: bool,
@@ -54,6 +55,7 @@ pub fn run(locked: bool, filter_regex: Option<&regex::Regex>) {
     for TestDescription {
         name,
         category,
+        script_data,
         expected_result,
         contract_paths,
         validate_abi,
@@ -81,7 +83,7 @@ pub fn run(locked: bool, filter_regex: Option<&regex::Regex>) {
                     ),
                 };
 
-                let result = crate::e2e_vm_tests::harness::runs_in_vm(&name, locked);
+                let result = crate::e2e_vm_tests::harness::runs_in_vm(&name, script_data, locked);
                 assert_eq!(result.0, res);
                 if validate_abi {
                     assert!(crate::e2e_vm_tests::harness::test_json_abi(&name, &result.1).is_ok());
@@ -282,6 +284,23 @@ fn parse_test_toml(path: &Path) -> Result<TestDescription, String> {
         return Err("'fail' tests must contain some FileCheck verification directives.".to_owned());
     }
 
+    let script_data = match &category {
+        TestCategory::Runs | TestCategory::RunsWithContract => {
+            match toml_content.get("script_data") {
+                Some(toml::Value::String(v)) => {
+                    let decoded = hex::decode(v)
+                        .map_err(|e| format!("Invalid hex value for 'script_data': {}", e))?;
+                    Some(decoded)
+                }
+                Some(_) => {
+                    return Err("Expected 'script_data' to be a hex string.".to_owned());
+                }
+                _ => None,
+            }
+        }
+        TestCategory::Compiles | TestCategory::FailsToCompile | TestCategory::Disabled => None,
+    };
+
     let expected_result = match &category {
         TestCategory::Runs | TestCategory::RunsWithContract => {
             Some(get_expected_result(&toml_content)?)
@@ -333,6 +352,7 @@ fn parse_test_toml(path: &Path) -> Result<TestDescription, String> {
     Ok(TestDescription {
         name,
         category,
+        script_data,
         expected_result,
         contract_paths,
         validate_abi,

test/src/e2e_vm_tests/test_programs/should_fail/main_args/main_args_mutation/Forc.lock

@@ -0,0 +1,9 @@
+[[package]]
+name = 'core'
+source = 'path+from-root-2F8090B18C264A91'
+dependencies = []
+
+[[package]]
+name = 'main_args_mutation'
+source = 'root'
+dependencies = ['core']

test/src/e2e_vm_tests/test_programs/should_fail/main_args/main_args_mutation/Forc.toml

@@ -0,0 +1,9 @@
+[project]
+authors = ["Fuel Labs <[email protected]>"]
+entry = "main.sw"
+license = "Apache-2.0"
+name = "main_args_mutation"
+implicit-std = false
+
+[dependencies]
+core = { path = "../../../../../../../sway-lib-core" }

test/src/e2e_vm_tests/test_programs/should_fail/main_args/main_args_mutation/json_abi_oracle.json

@@ -0,0 +1,42 @@
+[
+  {
+    "inputs": [
+      {
+        "components": [
+          {
+            "components": null,
+            "name": "val",
+            "type": "u64",
+            "typeArguments": null
+          }
+        ],
+        "name": "baba",
+        "type": "struct TestStruct",
+        "typeArguments": null
+      },
+      {
+        "components": [
+          {
+            "components": null,
+            "name": "val",
+            "type": "u64",
+            "typeArguments": null
+          }
+        ],
+        "name": "keke",
+        "type": "struct TestStruct",
+        "typeArguments": null
+      }
+    ],
+    "name": "main",
+    "outputs": [
+      {
+        "components": null,
+        "name": "",
+        "type": "u64",
+        "typeArguments": null
+      }
+    ],
+    "type": "function"
+  }
+]

test/src/e2e_vm_tests/test_programs/should_fail/main_args/main_args_mutation/src/main.sw

@@ -0,0 +1,10 @@
+script;
+
+struct TestStruct {
+    val: u64,
+}
+
+fn main(ref mut baba: TestStruct) -> u64 {
+    baba.val += 1;
+    baba.val
+}

test/src/e2e_vm_tests/test_programs/should_fail/main_args/main_args_mutation/test.toml

@@ -0,0 +1,4 @@
+category = "fail"
+
+# check: fn main(ref mut baba: TestStruct) -> u64 {
+# nextln: $()ref mut parameter not allowed for main()

test/src/e2e_vm_tests/test_programs/should_pass/language/main_args/main_args_copy/Forc.lock

@@ -0,0 +1,9 @@
+[[package]]
+name = 'core'
+source = 'path+from-root-EE00857339AC041F'
+dependencies = []
+
+[[package]]
+name = 'main_args_copy'
+source = 'root'
+dependencies = ['core']

test/src/e2e_vm_tests/test_programs/should_pass/language/main_args/main_args_copy/Forc.toml

@@ -0,0 +1,9 @@
+[project]
+authors = ["Fuel Labs <[email protected]>"]
+entry = "main.sw"
+license = "Apache-2.0"
+name = "main_args_copy"
+implicit-std = false
+
+[dependencies]
+core = { path = "../../../../../../../../sway-lib-core" }