Support value and salt in constructors

Author: grandizzy

crates/compilers/src/preprocessor/data.rs

@@ -171,14 +171,12 @@ impl ContractData {
     /// vm.deployCode("artifact path", encodeArgs335(DeployHelper335.ConstructorArgs({name: name, symbol: symbol})))
     /// ```
     pub fn build_helper(&self) -> Option<String> {
-        let Self { contract_id, path, name, constructor_data, artifact } = self;
+        let Self { contract_id, path, name, constructor_data, artifact: _ } = self;
 
         let Some(constructor_details) = constructor_data else { return None };
         let contract_id = contract_id.get();
         let struct_fields = &constructor_details.struct_fields;
         let abi_encode_args = &constructor_details.abi_encode_args;
-        let vm_interface_name = format!("VmContractHelper{contract_id}");
-        let vm = format!("{vm_interface_name}(0x7109709ECfa91a80626fF3989D68f67F5b1DD12D)");
 
         let helper = format!(
             r#"
@@ -195,16 +193,6 @@ abstract contract DeployHelper{contract_id} is {name} {{
 function encodeArgs{contract_id}(DeployHelper{contract_id}.ConstructorArgs memory args) pure returns (bytes memory) {{
     return abi.encode({abi_encode_args});
 }}
-
-function deployCode{contract_id}(DeployHelper{contract_id}.ConstructorArgs memory args) returns({name}) {{
-    return {name}(payable({vm}.deployCode("{artifact}", encodeArgs{contract_id}(args))));
-}}
-
-interface {vm_interface_name} {{
-    function deployCode(string memory _artifact, bytes memory _data) external returns (address);
-    function deployCode(string memory _artifact) external returns (address);
-    function getCode(string memory _artifact) external returns (bytes memory);
-}}
         "#,
             path = path.to_slash_lossy(),
         );

crates/compilers/src/preprocessor/deps.rs

@@ -8,7 +8,7 @@ use crate::{
 use itertools::Itertools;
 use solar_parse::interface::Session;
 use solar_sema::{
-    hir::{ContractId, Expr, ExprKind, Hir, TypeKind, Visit},
+    hir::{ContractId, Expr, ExprKind, Hir, NamedArg, TypeKind, Visit},
     interface::{data_structures::Never, source_map::FileName, SourceMap},
 };
 use std::{
@@ -95,8 +95,9 @@ impl PreprocessorDependencies {
 enum BytecodeDependencyKind {
     /// `type(Contract).creationCode`
     CreationCode,
-    /// `new Contract`. Holds the name of the contract and args length.
-    New(String, usize, usize),
+    /// `new Contract`. Holds the name of the contract, args length and offset, `msg.value` (if
+    /// any) and salt (if any).
+    New(String, usize, usize, Option<String>, Option<String>),
 }
 
 /// Represents a single bytecode dependency.
@@ -193,6 +194,8 @@ impl<'hir> Visit<'hir> for BytecodeDependencyCollector<'hir> {
                                     name.to_string(),
                                     args_len.to_usize(),
                                     offset,
+                                    named_arg(self.src, named_args, "value", self.source_map),
+                                    named_arg(self.src, named_args, "salt", self.source_map),
                                 ),
                                 loc: SourceMapLocation::from_span(self.source_map, ty.span),
                                 referenced_contract: contract_id,
@@ -222,6 +225,21 @@ impl<'hir> Visit<'hir> for BytecodeDependencyCollector<'hir> {
     }
 }
 
+/// Helper function to extract value of a given named arg.
+fn named_arg(
+    src: &str,
+    named_args: &Option<&[NamedArg<'_>]>,
+    arg: &str,
+    source_map: &SourceMap,
+) -> Option<String> {
+    named_args.unwrap_or_default().iter().find(|named_arg| named_arg.name.as_str() == arg).map(
+        |named_arg| {
+            let named_arg_loc = SourceMapLocation::from_span(source_map, named_arg.value.span);
+            src[named_arg_loc.start..named_arg_loc.end].to_string()
+        },
+    )
+}
+
 /// Goes over all test/script files and replaces bytecode dependencies with cheatcode
 /// invocations.
 pub(crate) fn remove_bytecode_dependencies(
@@ -260,39 +278,46 @@ pub(crate) fn remove_bytecode_dependencies(
                         format!("{vm}.getCode(\"{artifact}\")"),
                     ));
                 }
-                BytecodeDependencyKind::New(name, args_length, offset) => {
-                    if constructor_data.is_none() {
-                        // if there's no constructor, we can just call deployCode with one
-                        // argument
-                        updates.insert((
-                            dep.loc.start,
-                            dep.loc.end + args_length,
-                            format!("{name}(payable({vm}.deployCode(\"{artifact}\")))"),
-                        ));
-                    } else {
-                        // if there's a constructor, we use our helper
+                BytecodeDependencyKind::New(name, args_length, offset, value, salt) => {
+                    let mut update = format!("{name}(payable({vm}.deployCode({{");
+                    update.push_str(&format!("_artifact: \"{artifact}\""));
+
+                    if let Some(value) = value {
+                        update.push_str(", ");
+                        update.push_str(&format!("_value: {value}"));
+                    }
+
+                    if let Some(salt) = salt {
+                        update.push_str(", ");
+                        update.push_str(&format!("_salt: {salt}"));
+                    }
+
+                    if constructor_data.is_some() {
+                        // Insert our helper
                         used_helpers.insert(dep.referenced_contract);
-                        updates.insert((
-                            dep.loc.start,
-                            dep.loc.end + offset,
-                            format!(
-                                "deployCode{id}(DeployHelper{id}.ConstructorArgs",
-                                id = dep.referenced_contract.get()
-                            ),
+
+                        update.push_str(", ");
+                        update.push_str(&format!(
+                            "_args: encodeArgs{id}(DeployHelper{id}.ConstructorArgs",
+                            id = dep.referenced_contract.get()
                         ));
+                        updates.insert((dep.loc.start, dep.loc.end + offset, update));
                         updates.insert((
                             dep.loc.end + args_length,
                             dep.loc.end + args_length,
-                            ")".to_string(),
+                            ")})))".to_string(),
                         ));
+                    } else {
+                        update.push_str("})))");
+                        updates.insert((dep.loc.start, dep.loc.end + args_length, update));
                     }
                 }
             };
         }
         let helper_imports = used_helpers.into_iter().map(|id| {
             let id = id.get();
             format!(
-                "import {{DeployHelper{id}, encodeArgs{id}, deployCode{id}}} from \"foundry-pp/DeployHelper{id}.sol\";",
+                "import {{DeployHelper{id}, encodeArgs{id}}} from \"foundry-pp/DeployHelper{id}.sol\";",
             )
         }).join("\n");
         updates.insert((
@@ -303,8 +328,14 @@ pub(crate) fn remove_bytecode_dependencies(
 {helper_imports}
 
 interface {vm_interface_name} {{
-    function deployCode(string memory _artifact, bytes memory _data) external returns (address);
     function deployCode(string memory _artifact) external returns (address);
+    function deployCode(string memory _artifact, bytes32 _salt) external returns (address);
+    function deployCode(string memory _artifact, bytes memory _args) external returns (address);
+    function deployCode(string memory _artifact, bytes memory _args, bytes32 _salt) external returns (address);
+    function deployCode(string memory _artifact, uint256 _value) external returns (address);
+    function deployCode(string memory _artifact, uint256 _value, bytes32 _salt) external returns (address);
+    function deployCode(string memory _artifact, bytes memory _args, uint256 _value) external returns (address);
+    function deployCode(string memory _artifact, bytes memory _args, uint256 _value, bytes32 _salt) external returns (address);
     function getCode(string memory _artifact) external returns (bytes memory);
 }}"#
             ),

test-data/preprocessor/src/CounterE.sol

@@ -0,0 +1,13 @@
+// Contracts with payable constructor
+contract CounterE {
+    constructor() payable {}
+}
+
+contract CounterF {
+    constructor(uint256 x) payable {}
+}
+
+contract CounterG {
+    constructor(address) payable {}
+}
+

test-data/preprocessor/test/CounterTest.sol

@@ -3,6 +3,7 @@ import {Counter as CounterV1} from "src/v1/Counter.sol";
 import "src/CounterB.sol";
 import "src/CounterC.sol";
 import "src/CounterD.sol";
+import "src/CounterE.sol";
 
 contract CounterTest {
     Counter public counter;
@@ -12,6 +13,10 @@ contract CounterTest {
     CounterV1 public counterv1;
     Counter public counter5 = new Counter{salt: bytes32("123")}();
     CounterB public counter6 = new CounterB {salt: bytes32("123")}   (address(this), 44, true, address(this));
+    CounterE public counter7 = new CounterE{  value: 111,   salt: bytes32("123")}();
+    CounterF public counter8 = new CounterF{value: 222, salt: bytes32("123")}(11);
+    CounterG public counter9 = new CounterG      {       value: 333, salt: bytes32("123")         }    (address(this));
+    CounterG public counter10 = new CounterG{  value: 333 }(address(this));
 
     function setUp() public {
         counter = new Counter();