AA benchmark tests (#577)

* add yarn aabenchmark

* lint/prettier

* fix foundry.toml

* format

* remove libs incompatible with forge coverage

* prettier ignore file w bytecode

* rename import

---------

Co-authored-by: WhiteOakKong <aadam.debevec@outlook.com>
Co-authored-by: WhiteOakKong <92236155+WhiteOakKong@users.noreply.github.com>

Author: nkrishang

.prettierignore

@@ -9,4 +9,5 @@ node_modules/
 typechain/
 
 # files
+src/test/smart-wallet/utils/AABenchmarkArtifacts.sol
 coverage.json

contracts/lib/TWStrings.sol

@@ -64,4 +64,132 @@ library TWStrings {
         require(value == 0, "Strings: hex length insufficient");
         return string(buffer);
     }
+
+    /// @dev Returns the hexadecimal representation of `value`.
+    /// The output is prefixed with "0x", encoded using 2 hexadecimal digits per byte,
+    /// and the alphabets are capitalized conditionally according to
+    /// https://eips.ethereum.org/EIPS/eip-55
+    function toHexStringChecksummed(address value) internal pure returns (string memory str) {
+        str = toHexString(value);
+        /// @solidity memory-safe-assembly
+        assembly {
+            let mask := shl(6, div(not(0), 255)) // `0b010000000100000000 ...`
+            let o := add(str, 0x22)
+            let hashed := and(keccak256(o, 40), mul(34, mask)) // `0b10001000 ... `
+            let t := shl(240, 136) // `0b10001000 << 240`
+            for {
+                let i := 0
+            } 1 {
+
+            } {
+                mstore(add(i, i), mul(t, byte(i, hashed)))
+                i := add(i, 1)
+                if eq(i, 20) {
+                    break
+                }
+            }
+            mstore(o, xor(mload(o), shr(1, and(mload(0x00), and(mload(o), mask)))))
+            o := add(o, 0x20)
+            mstore(o, xor(mload(o), shr(1, and(mload(0x20), and(mload(o), mask)))))
+        }
+    }
+
+    /// @dev Returns the hexadecimal representation of `value`.
+    /// The output is prefixed with "0x" and encoded using 2 hexadecimal digits per byte.
+    function toHexString(address value) internal pure returns (string memory str) {
+        str = toHexStringNoPrefix(value);
+        /// @solidity memory-safe-assembly
+        assembly {
+            let strLength := add(mload(str), 2) // Compute the length.
+            mstore(str, 0x3078) // Write the "0x" prefix.
+            str := sub(str, 2) // Move the pointer.
+            mstore(str, strLength) // Write the length.
+        }
+    }
+
+    /// @dev Returns the hexadecimal representation of `value`.
+    /// The output is encoded using 2 hexadecimal digits per byte.
+    function toHexStringNoPrefix(address value) internal pure returns (string memory str) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            str := mload(0x40)
+
+            // Allocate the memory.
+            // We need 0x20 bytes for the trailing zeros padding, 0x20 bytes for the length,
+            // 0x02 bytes for the prefix, and 0x28 bytes for the digits.
+            // The next multiple of 0x20 above (0x20 + 0x20 + 0x02 + 0x28) is 0x80.
+            mstore(0x40, add(str, 0x80))
+
+            // Store "0123456789abcdef" in scratch space.
+            mstore(0x0f, 0x30313233343536373839616263646566)
+
+            str := add(str, 2)
+            mstore(str, 40)
+
+            let o := add(str, 0x20)
+            mstore(add(o, 40), 0)
+
+            value := shl(96, value)
+
+            // We write the string from rightmost digit to leftmost digit.
+            // The following is essentially a do-while loop that also handles the zero case.
+            for {
+                let i := 0
+            } 1 {
+
+            } {
+                let p := add(o, add(i, i))
+                let temp := byte(i, value)
+                mstore8(add(p, 1), mload(and(temp, 15)))
+                mstore8(p, mload(shr(4, temp)))
+                i := add(i, 1)
+                if eq(i, 20) {
+                    break
+                }
+            }
+        }
+    }
+
+    /// @dev Returns the hex encoded string from the raw bytes.
+    /// The output is encoded using 2 hexadecimal digits per byte.
+    function toHexString(bytes memory raw) internal pure returns (string memory str) {
+        str = toHexStringNoPrefix(raw);
+        /// @solidity memory-safe-assembly
+        assembly {
+            let strLength := add(mload(str), 2) // Compute the length.
+            mstore(str, 0x3078) // Write the "0x" prefix.
+            str := sub(str, 2) // Move the pointer.
+            mstore(str, strLength) // Write the length.
+        }
+    }
+
+    /// @dev Returns the hex encoded string from the raw bytes.
+    /// The output is encoded using 2 hexadecimal digits per byte.
+    function toHexStringNoPrefix(bytes memory raw) internal pure returns (string memory str) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            let length := mload(raw)
+            str := add(mload(0x40), 2) // Skip 2 bytes for the optional prefix.
+            mstore(str, add(length, length)) // Store the length of the output.
+
+            // Store "0123456789abcdef" in scratch space.
+            mstore(0x0f, 0x30313233343536373839616263646566)
+
+            let o := add(str, 0x20)
+            let end := add(raw, length)
+
+            for {
+
+            } iszero(eq(raw, end)) {
+
+            } {
+                raw := add(raw, 1)
+                mstore8(add(o, 1), mload(and(mload(raw), 15)))
+                mstore8(o, mload(and(shr(4, mload(raw)), 15)))
+                o := add(o, 2)
+            }
+            mstore(o, 0) // Zeroize the slot after the string.
+            mstore(0x40, add(o, 0x20)) // Allocate the memory.
+        }
+    }
 }

foundry.toml

@@ -38,8 +38,9 @@ remappings = [
   'erc721a-upgradeable/=lib/ERC721A-Upgradeable/',
   'erc721a/=lib/ERC721A/',
   '@thirdweb-dev/dynamic-contracts/=lib/dynamic-contracts/',
-  'lib/sstore2=lib/dynamic-contracts/lib/sstore2/'
+  'lib/sstore2=lib/dynamic-contracts/lib/sstore2/',
 ]
+fs_permissions = [{ access = "read-write", path = "./src/test/smart-wallet/utils"}]
 src = 'contracts'
 test = 'src/test'
 verbosity = 0

package.json

@@ -59,7 +59,8 @@
     "forge:build": "forge build",
     "forge:test": "forge test",
     "gas": "forge test --mc Benchmark --gas-report > gasreport.txt",
-    "forge:snapshot": "forge snapshot --check"
+    "forge:snapshot": "forge snapshot --check",
+    "aabenchmark": "forge test --mc AABenchmarkPrepare && forge test --mc ProfileThirdwebAccount -vvv"
   },
   "dependencies": {}
 }

src/test/smart-wallet/utils/AABenchmarkArtifacts.sol

@@ -0,0 +1,81 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity ^0.8.0;
+
+// Test utils
+import "../../utils/BaseTest.sol";
+
+// Account Abstraction setup for smart wallets.
+import { IEntryPoint } from "contracts/prebuilts/account/utils/Entrypoint.sol";
+
+import { AccountFactory } from "contracts/prebuilts/account/non-upgradeable/AccountFactory.sol";
+import "contracts/lib/TWStrings.sol";
+
+import "forge-std/Test.sol";
+
+contract AABenchmarkPrepare is BaseTest {
+    AccountFactory private accountFactory;
+
+    function setUp() public override {
+        super.setUp();
+        accountFactory = new AccountFactory(IEntryPoint(payable(address(0x5FF137D4b0FDCD49DcA30c7CF57E578a026d2789))));
+    }
+
+    function test_prepareBenchmarkFile() public {
+        address accountFactoryAddress = address(accountFactory);
+        bytes memory accountFactoryBytecode = accountFactoryAddress.code;
+
+        address accountImplAddress = accountFactory.accountImplementation();
+        bytes memory accountImplBytecode = accountImplAddress.code;
+
+        string memory accountFactoryAddressString = string.concat(
+            "address constant THIRDWEB_ACCOUNT_FACTORY_ADDRESS = ",
+            TWStrings.toHexStringChecksummed(accountFactoryAddress),
+            ";"
+        );
+        string memory accountFactoryBytecodeString = string.concat(
+            'bytes constant THIRDWEB_ACCOUNT_FACTORY_BYTECODE = hex"',
+            TWStrings.toHexStringNoPrefix(accountFactoryBytecode),
+            '"',
+            ";"
+        );
+
+        string memory accountImplAddressString = string.concat(
+            "address constant THIRDWEB_ACCOUNT_IMPL_ADDRESS = ",
+            TWStrings.toHexStringChecksummed(accountImplAddress),
+            ";"
+        );
+        string memory accountImplBytecodeString = string.concat(
+            'bytes constant THIRDWEB_ACCOUNT_IMPL_BYTECODE = hex"',
+            TWStrings.toHexStringNoPrefix(accountImplBytecode),
+            '"',
+            ";"
+        );
+
+        string memory path = "src/test/smart-wallet/utils/AABenchmarkArtifacts.sol";
+
+        vm.removeFile(path);
+
+        vm.writeLine(path, "");
+        vm.writeLine(path, "pragma solidity ^0.8.0;");
+        vm.writeLine(path, "interface ThirdwebAccountFactory {");
+        vm.writeLine(
+            path,
+            "    function createAccount(address _admin, bytes calldata _data) external returns (address);"
+        );
+        vm.writeLine(
+            path,
+            "    function getAddress(address _adminSigner, bytes calldata _data) external view returns (address);"
+        );
+        vm.writeLine(path, "}");
+
+        vm.writeLine(path, "interface ThirdwebAccount {");
+        vm.writeLine(path, "    function execute(address _target, uint256 _value, bytes calldata _calldata) external;");
+        vm.writeLine(path, "}");
+        vm.writeLine(path, accountFactoryAddressString);
+        vm.writeLine(path, accountImplAddressString);
+        vm.writeLine(path, accountFactoryBytecodeString);
+        vm.writeLine(path, accountImplBytecodeString);
+
+        vm.writeLine(path, "");
+    }
+}

src/test/smart-wallet/utils/AABenchmarkPrepare.sol

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity ^0.8.0;
+
+import "./AATestBase.sol";
+import { ThirdwebAccountFactory, ThirdwebAccount, THIRDWEB_ACCOUNT_FACTORY_ADDRESS, THIRDWEB_ACCOUNT_IMPL_ADDRESS, THIRDWEB_ACCOUNT_FACTORY_BYTECODE, THIRDWEB_ACCOUNT_IMPL_BYTECODE } from "./AABenchmarkArtifacts.sol";
+
+contract ProfileThirdwebAccount is AAGasProfileBase {
+    ThirdwebAccountFactory factory;
+
+    function setUp() external {
+        initializeTest("thirdwebAccount");
+        factory = ThirdwebAccountFactory(THIRDWEB_ACCOUNT_FACTORY_ADDRESS);
+        vm.etch(address(factory), THIRDWEB_ACCOUNT_FACTORY_BYTECODE);
+        vm.etch(THIRDWEB_ACCOUNT_IMPL_ADDRESS, THIRDWEB_ACCOUNT_IMPL_BYTECODE);
+        setAccount();
+    }
+
+    function fillData(
+        address _to,
+        uint256 _value,
+        bytes memory _data
+    ) internal view override returns (bytes memory) {
+        return abi.encodeWithSelector(ThirdwebAccount.execute.selector, _to, _value, _data);
+    }
+
+    function getSignature(UserOperation memory _op) internal view override returns (bytes memory) {
+        return signUserOpHash(key, _op);
+    }
+
+    function createAccount(address _owner) internal override {
+        // if (address(account).code.length == 0) {
+        factory.createAccount(_owner, "");
+        // }
+    }
+
+    function getAccountAddr(address _owner) internal view override returns (IAccount) {
+        return IAccount(factory.getAddress(_owner, ""));
+    }
+
+    function getInitCode(address _owner) internal view override returns (bytes memory) {
+        return abi.encodePacked(address(factory), abi.encodeWithSelector(factory.createAccount.selector, _owner, ""));
+    }
+
+    function getDummySig(UserOperation memory _op) internal pure override returns (bytes memory) {
+        return
+            hex"fffffffffffffffffffffffffffffff0000000000000000000000000000000007aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1c";
+    }
+}