├── .github └── workflows │ ├── lint.yaml │ └── test.yaml ├── .gitignore ├── LICENSE ├── README.md ├── brownie-config.yaml ├── contracts ├── CommandBuilder.sol ├── Helpers │ └── TupleHelper.sol ├── Libraries │ ├── Events.sol │ ├── Math.sol │ ├── Receiver.sol │ ├── Strings.sol │ └── Tupler.sol ├── VM.sol └── test │ ├── CommandBuilderHarness.sol │ ├── MultiReturn.sol │ ├── Revert.sol │ ├── Sender.sol │ ├── SimpleToken.sol │ ├── StateTest.sol │ ├── SubPlanTests.sol │ ├── TestContract.sol │ └── TestableVM.sol ├── poetry.lock ├── pyproject.toml ├── tests ├── conftest.py ├── test_chaining_actions.py ├── test_curve_add_liquidity.py ├── test_dyn.py ├── test_helpers.py ├── test_one_inch.py ├── test_swaps.py └── test_weiroll.py └── weiroll.py /.github/workflows/lint.yaml: -------------------------------------------------------------------------------- 1 | name: Lint 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | - develop 8 | pull_request: 9 | 10 | jobs: 11 | 12 | solidity: 13 | runs-on: ubuntu-latest 14 | 15 | steps: 16 | - name: Check out github repository 17 | uses: actions/checkout@v2 18 | with: 19 | fetch-depth: 1 20 | 21 | - name: Setup node.js 22 | uses: actions/setup-node@v1 23 | with: 24 | node-version: '12.x' 25 | 26 | - name: Set yarn cache directory path 27 | id: yarn-cache-dir-path 28 | run: echo "::set-output name=dir::$(yarn cache dir)" 29 | 30 | - name: Restore yarn cache 31 | uses: actions/cache@v2 32 | id: yarn-cache 33 | with: 34 | path: | 35 | ${{ steps.yarn-cache-dir-path.outputs.dir }} 36 | key: ${{ runner.os }}-yarn-${{ hashFiles('**/yarn.lock') }} 37 | restore-keys: | 38 | ${{ runner.os }}-yarn-${{ hashFiles('**/yarn.lock') }} 39 | ${{ runner.os }}-yarn- 40 | 41 | - name: Install node.js dependencies 42 | run: yarn --frozen-lockfile 43 | 44 | - name: Run linter on *.sol and *.json 45 | run: yarn lint:check 46 | 47 | commits: 48 | runs-on: ubuntu-latest 49 | 50 | steps: 51 | - name: Check out github repository 52 | uses: actions/checkout@v2 53 | with: 54 | fetch-depth: 0 55 | 56 | - name: Run commitlint 57 | uses: wagoid/commitlint-github-action@v2 58 | 59 | brownie: 60 | runs-on: ubuntu-latest 61 | 62 | steps: 63 | - name: Check out github repository 64 | uses: actions/checkout@v2 65 | with: 66 | fetch-depth: 1 67 | 68 | - name: Set up python 3.9 69 | uses: actions/setup-python@v2 70 | with: 71 | python-version: 3.9 72 | 73 | - name: Poetry 74 | uses: abatilo/actions-poetry@v2.0.0 75 | with: 76 | poetry-version: 1.1.13 77 | 78 | - name: Install dependencies 79 | run: poetry install 80 | 81 | - name: Run black 82 | run: poetry run black --check --include "(tests|scripts)" . 83 | 84 | # TODO: Add Slither Static Analyzer 85 | -------------------------------------------------------------------------------- /.github/workflows/test.yaml: -------------------------------------------------------------------------------- 1 | name: Test 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | - develop 8 | pull_request: 9 | 10 | jobs: 11 | test: 12 | runs-on: ubuntu-latest 13 | 14 | steps: 15 | - uses: actions/checkout@v1 16 | 17 | - name: Cache compiler installations 18 | uses: actions/cache@v2 19 | with: 20 | path: | 21 | ~/.solcx 22 | ~/.vvm 23 | key: ${{ runner.os }}-compiler-cache 24 | 25 | - name: Setup node.js 26 | uses: actions/setup-node@v1 27 | with: 28 | node-version: '12.x' 29 | 30 | - name: Install ganache 31 | run: npm install -g ganache-cli@6.12.1 32 | 33 | - name: Set up python 3.9 34 | uses: actions/setup-python@v2 35 | with: 36 | python-version: 3.9 37 | 38 | - name: Poetry 39 | uses: abatilo/actions-poetry@v2.0.0 40 | with: 41 | poetry-version: 1.1.13 42 | 43 | - name: Install dependencies 44 | run: poetry install 45 | 46 | - name: Compile Code 47 | run: poetry run brownie compile --size 48 | 49 | - name: Run Tests 50 | env: 51 | ETHERSCAN_TOKEN: MW5CQA6QK5YMJXP2WP3RA36HM5A7RA1IHA 52 | WEB3_INFURA_PROJECT_ID: b7821200399e4be2b4e5dbdf06fbe85b 53 | run: poetry run brownie test 54 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .hypothesis 2 | __pycache__ 3 | build 4 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 weiroll 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # weiroll-py 2 | 3 | weiroll-py is a planner for the operation-chaining/scripting language [weiroll](https://github.com/weiroll/weiroll). 4 | weiroll-py is inspired by [weiroll.js](https://github.com/weiroll/weiroll.js). 5 | 6 | It provides an easy-to-use API for generating weiroll programs that can then be passed to any compatible implementation. 7 | 8 | ## Installation 9 | 10 | ``` 11 | pip install weiroll-py==0.2.1 12 | ``` 13 | 14 | where `0.2.1` is the latest version. 15 | 16 | ## Usage 17 | 18 | ### Wrapping contracts 19 | Weiroll programs consist of a sequence of calls to functions in external contracts. These calls can either be delegate calls to dedicated library contracts, or standard/static calls to external contracts. Before you can start creating a weiroll program, you will need to create interfaces for at least one contract you intend to use. 20 | 21 | The easiest way to do this is by wrapping brownie contract instances: 22 | 23 | ```python 24 | brownie_contract = brownie.Contract(address) 25 | contract = weiroll.WeirollContract( 26 | brownie_contract 27 | ) 28 | ``` 29 | 30 | This will produce a contract object that generates delegate calls to the brownie contract in `WeirollContract`. 31 | 32 | To create delegate to an external contract, use `createLibrary`: 33 | 34 | ```python 35 | brownie_contract = brownie.Contract(address) 36 | # Makes calls using CALL 37 | contract = weiroll.WeirollContract.createContract(brownie_contract) 38 | # Makes calls using STATICCALL 39 | contract = weiroll.WeirollContract.createContract(brownie_contract, weiroll.CommandFlags.STATICCALL) 40 | ``` 41 | 42 | You can repeat this for each contract you wish to use. A weiroll `WeirollContract` object can be reused across as many planner instances as you wish; there is no need to construct them again for each new program. 43 | 44 | ### Planning programs 45 | 46 | First, instantiate a planner: 47 | 48 | ```python 49 | planner = weiroll.WeirollPlanner() 50 | ``` 51 | 52 | Next, add one or more commands to execute: 53 | 54 | ```python 55 | ret = planner.add(contract.func(a, b)) 56 | ``` 57 | 58 | Return values from one invocation can be used in another one: 59 | 60 | ```python 61 | planner.add(contract.func2(ret)) 62 | ``` 63 | 64 | Remember to wrap each call to a contract in `planner.add`. Attempting to pass the result of one contract function directly to another will not work - each one needs to be added to the planner! 65 | 66 | For calls to external contracts, you can also pass a value in ether to send: 67 | 68 | ```python 69 | planner.add(contract.func(a, b).withValue(c)) 70 | ``` 71 | 72 | `withValue` takes the same argument types as contract functions, so you can pass the return value of another function, or a literal value. You cannot combine `withValue` with delegate calls (eg, calls to a library created with `Contract.newLibrary`) or static calls. 73 | 74 | Likewise, if you want to make a particular call static, you can use `.staticcall()`: 75 | 76 | ```python 77 | result = planner.add(contract.func(a, b).staticcall()) 78 | ``` 79 | 80 | Weiroll only supports functions that return a single value by default. If your function returns multiple values, though, you can instruct weiroll to wrap it in a `bytes`, which subsequent commands can decode and work with: 81 | 82 | ```python 83 | ret = planner.add(contract.func(a, b).rawValue()) 84 | ``` 85 | 86 | Once you are done planning operations, generate the program: 87 | 88 | ```python 89 | commands, state = planner.plan() 90 | ``` 91 | 92 | ### Subplans 93 | In some cases it may be useful to be able to instantiate nested instances of the weiroll VM - for example, when using flash loans, or other systems that function by making a callback to your code. The weiroll planner supports this via 'subplans'. 94 | 95 | To make a subplan, construct the operations that should take place inside the nested instance normally, then pass the planner object to a contract function that executes the subplan, and pass that to the outer planner's `.addSubplan()` function instead of `.add()`. 96 | 97 | For example, suppose you want to call a nested instance to do some math: 98 | 99 | ```python 100 | subplanner = WeirollPlanner() 101 | sum = subplanner.add(Math.add(1, 2)) 102 | 103 | planner = WeirollPlanner() 104 | planner.addSubplan(Weiroll.execute(subplanner, subplanner.state)) 105 | planner.add(events.logUint(sum)) 106 | 107 | commands, state = planner.plan() 108 | ``` 109 | 110 | Subplan functions must specify which argument receives the current state using the special variable `Planner.state`, and must take exactly one subplanner and one state argument. Subplan functions must either return an updated state or nothing. 111 | 112 | If a subplan returns updated state, return values created in a subplanner, such as `sum` above, can be referenced in the outer scope, and even in other subplans, as long as they are referenced after the command that produces them. Subplans that do not return updated state are read-only, and return values defined inside them cannot be referenced outside them. 113 | 114 | ## More examples 115 | 116 | Review [tests](/tests) for more examples. 117 | 118 | ## Credits 119 | 120 | - [@WyseNynja](https://github.com/WyseNynja) for the original implementation 121 | -------------------------------------------------------------------------------- /brownie-config.yaml: -------------------------------------------------------------------------------- 1 | networks: 2 | default: mainnet-fork 3 | 4 | autofetch_sources: true 5 | 6 | # require OpenZepplin Contracts 7 | dependencies: 8 | - OpenZeppelin/openzeppelin-contracts@4.1.0 9 | 10 | # path remapping to support imports from GitHub/NPM 11 | compiler: 12 | solc: 13 | version: 0.8.11 14 | remappings: 15 | - "@openzeppelin=OpenZeppelin/openzeppelin-contracts@4.1.0" 16 | 17 | optimizer: 18 | details: 19 | yul: true 20 | -------------------------------------------------------------------------------- /contracts/CommandBuilder.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | 3 | pragma solidity ^0.8.11; 4 | 5 | library CommandBuilder { 6 | 7 | uint256 constant IDX_VARIABLE_LENGTH = 0x80; 8 | uint256 constant IDX_VALUE_MASK = 0x7f; 9 | uint256 constant IDX_END_OF_ARGS = 0xff; 10 | uint256 constant IDX_USE_STATE = 0xfe; 11 | 12 | function buildInputs( 13 | bytes[] memory state, 14 | bytes4 selector, 15 | bytes32 indices 16 | ) internal view returns (bytes memory ret) { 17 | uint256 free; // Pointer to first free byte in tail part of message 18 | uint256 idx; 19 | 20 | // Determine the length of the encoded data 21 | for (uint256 i; i < 32;) { 22 | idx = uint8(indices[i]); 23 | if (idx == IDX_END_OF_ARGS) break; 24 | unchecked{free += 32;} 25 | unchecked{++i;} 26 | } 27 | 28 | // Encode it 29 | uint256 bytesWritten; 30 | assembly { 31 | ret := mload(0x40) 32 | bytesWritten := add(bytesWritten, 4) 33 | mstore(0x40, add(ret, and(add(add(bytesWritten, 0x20), 0x1f), not(0x1f)))) 34 | mstore(add(ret, 32), selector) 35 | } 36 | uint256 count = 0; 37 | bytes memory stateData; // Optionally encode the current state if the call requires it 38 | for (uint256 i; i < 32;) { 39 | idx = uint8(indices[i]); 40 | if (idx == IDX_END_OF_ARGS) break; 41 | 42 | if (idx & IDX_VARIABLE_LENGTH != 0) { 43 | if (idx == IDX_USE_STATE) { 44 | assembly { 45 | bytesWritten := add(bytesWritten, 32) 46 | mstore(0x40, add(ret, and(add(add(bytesWritten, 0x20), 0x1f), not(0x1f)))) 47 | mstore(add(add(ret, 36), count), free) 48 | } 49 | if (stateData.length == 0) { 50 | stateData = abi.encode(state); 51 | } 52 | assembly { 53 | bytesWritten := add(bytesWritten, mload(stateData)) 54 | mstore(0x40, add(ret, and(add(add(bytesWritten, 0x20), 0x1f), not(0x1f)))) 55 | } 56 | memcpy(stateData, 32, ret, free + 4, stateData.length - 32); 57 | free += stateData.length - 32; 58 | } else { 59 | bytes memory stateVar = state[idx & IDX_VALUE_MASK]; 60 | uint256 arglen = stateVar.length; 61 | 62 | // Variable length data; put a pointer in the slot and write the data at the end 63 | assembly { 64 | bytesWritten := add(bytesWritten, 32) 65 | mstore(0x40, add(ret, and(add(add(bytesWritten, 0x20), 0x1f), not(0x1f)))) 66 | mstore(add(add(ret, 36), count), free) 67 | } 68 | assembly { 69 | bytesWritten := add(bytesWritten, arglen) 70 | mstore(0x40, add(ret, and(add(add(bytesWritten, 0x20), 0x1f), not(0x1f)))) 71 | } 72 | memcpy( 73 | stateVar, 74 | 0, 75 | ret, 76 | free + 4, 77 | arglen 78 | ); 79 | free += arglen; 80 | } 81 | } else { 82 | // Fixed length data; write it directly 83 | bytes memory stateVar = state[idx & IDX_VALUE_MASK]; 84 | assembly { 85 | bytesWritten := add(bytesWritten, mload(stateVar)) 86 | mstore(0x40, add(ret, and(add(add(bytesWritten, 0x20), 0x1f), not(0x1f)))) 87 | mstore(add(add(ret, 36), count), mload(add(stateVar, 32))) 88 | } 89 | } 90 | unchecked{count += 32;} 91 | unchecked{++i;} 92 | } 93 | assembly { 94 | mstore(ret, bytesWritten) 95 | } 96 | } 97 | 98 | function writeOutputs( 99 | bytes[] memory state, 100 | bytes1 index, 101 | bytes memory output 102 | ) internal pure returns (bytes[] memory) { 103 | uint256 idx = uint8(index); 104 | if (idx == IDX_END_OF_ARGS) return state; 105 | 106 | if (idx & IDX_VARIABLE_LENGTH != 0) { 107 | if (idx == IDX_USE_STATE) { 108 | state = abi.decode(output, (bytes[])); 109 | } else { 110 | // Check the first field is 0x20 (because we have only a single return value) 111 | uint256 argptr; 112 | assembly { 113 | argptr := mload(add(output, 32)) 114 | } 115 | require( 116 | argptr == 32, 117 | "Only one return value permitted (variable)" 118 | ); 119 | 120 | assembly { 121 | // Overwrite the first word of the return data with the length - 32 122 | mstore(add(output, 32), sub(mload(output), 32)) 123 | // Insert a pointer to the return data, starting at the second word, into state 124 | mstore( 125 | add(add(state, 32), mul(and(idx, IDX_VALUE_MASK), 32)), 126 | add(output, 32) 127 | ) 128 | } 129 | } 130 | } else { 131 | // Single word 132 | require( 133 | output.length == 32, 134 | "Only one return value permitted (static)" 135 | ); 136 | 137 | state[idx & IDX_VALUE_MASK] = output; 138 | } 139 | 140 | return state; 141 | } 142 | 143 | function writeTuple( 144 | bytes[] memory state, 145 | bytes1 index, 146 | bytes memory output 147 | ) internal view { 148 | uint256 idx = uint256(uint8(index)); 149 | if (idx == IDX_END_OF_ARGS) return; 150 | 151 | bytes memory entry = state[idx & IDX_VALUE_MASK] = new bytes(output.length + 32); 152 | 153 | memcpy(output, 0, entry, 32, output.length); 154 | assembly { 155 | let l := mload(output) 156 | mstore(add(entry, 32), l) 157 | } 158 | } 159 | 160 | function memcpy( 161 | bytes memory src, 162 | uint256 srcidx, 163 | bytes memory dest, 164 | uint256 destidx, 165 | uint256 len 166 | ) internal view { 167 | assembly { 168 | pop( 169 | staticcall( 170 | gas(), 171 | 4, 172 | add(add(src, 32), srcidx), 173 | len, 174 | add(add(dest, 32), destidx), 175 | len 176 | ) 177 | ) 178 | } 179 | } 180 | } 181 | -------------------------------------------------------------------------------- /contracts/Helpers/TupleHelper.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | 5 | contract TupleHelper { 6 | function insertElement(bytes calldata tuple, uint256 index, bytes32 element, bool returnRaw) 7 | public 8 | pure 9 | returns (bytes memory newTuple) 10 | { 11 | uint256 byteIndex; 12 | unchecked { byteIndex = index * 32; } 13 | require(byteIndex <= tuple.length); 14 | newTuple = bytes.concat(tuple[:byteIndex], element, tuple[byteIndex:]); 15 | if (returnRaw) { 16 | assembly { 17 | return(add(newTuple, 32), tuple.length) 18 | } 19 | } 20 | } 21 | 22 | function replaceElement(bytes calldata tuple, uint256 index, bytes32 element, bool returnRaw) 23 | public 24 | pure 25 | returns (bytes memory newTuple) 26 | { 27 | uint256 byteIndex; 28 | unchecked { 29 | byteIndex = index * 32; 30 | require(tuple.length >= 32 && byteIndex <= tuple.length - 32); 31 | newTuple = bytes.concat(tuple[:byteIndex], element, tuple[byteIndex+32:]); 32 | } 33 | if (returnRaw) { 34 | assembly { 35 | return(add(newTuple, 32), tuple.length) 36 | } 37 | } 38 | } 39 | 40 | function getElement(bytes memory tuple, uint256 index) 41 | public 42 | pure 43 | returns (bytes32) 44 | { 45 | assembly { 46 | return(add(tuple, mul(add(index, 1), 32)), 32) 47 | } 48 | } 49 | } 50 | -------------------------------------------------------------------------------- /contracts/Libraries/Events.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | contract Events { 5 | event LogBytes(bytes message); 6 | event LogAddress(address message); 7 | event LogString(string message); 8 | event LogBytes32(bytes32 message); 9 | event LogUint(uint256 message); 10 | 11 | function logBytes(bytes calldata message) external { 12 | emit LogBytes(message); 13 | } 14 | 15 | function logAddress(address message) external { 16 | emit LogAddress(message); 17 | } 18 | 19 | function logString(string calldata message) external { 20 | emit LogString(message); 21 | } 22 | 23 | function logBytes32(bytes32 message) external { 24 | emit LogBytes32(message); 25 | } 26 | 27 | function logUint(uint256 message) external { 28 | emit LogUint(message); 29 | } 30 | } 31 | -------------------------------------------------------------------------------- /contracts/Libraries/Math.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | contract Math { 5 | function add(uint256 a, uint256 b) external pure returns (uint256) { 6 | return a + b; 7 | } 8 | 9 | function sub(uint256 a, uint256 b) external pure returns (uint256) { 10 | return a - b; 11 | } 12 | 13 | function mul(uint256 a, uint256 b) external pure returns (uint256) { 14 | return a * b; 15 | } 16 | 17 | function sum(uint256[] calldata values) 18 | external 19 | pure 20 | returns (uint256 ret) 21 | { 22 | uint256 valuesLength = values.length; 23 | for (uint256 i; i < valuesLength; ++i) { 24 | ret += values[i]; 25 | } 26 | } 27 | 28 | function sumExtended( 29 | uint256 a, 30 | uint256 b, 31 | uint256 c, 32 | uint256 d, 33 | uint256 e, 34 | uint256 f, 35 | uint256 g // 7 args will trigger extended 36 | 37 | ) 38 | external 39 | pure 40 | returns (uint256 ret) 41 | { 42 | return a + b + c + d + e + f + g; 43 | } 44 | } 45 | -------------------------------------------------------------------------------- /contracts/Libraries/Receiver.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | contract Receiver { 5 | function receive(uint256 value) external payable returns (uint256) { 6 | require(value == msg.value, "receive: msg.value not sent."); 7 | } 8 | } 9 | -------------------------------------------------------------------------------- /contracts/Libraries/Strings.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | contract Strings { 5 | function strlen(string calldata x) external pure returns (uint256) { 6 | return bytes(x).length; 7 | } 8 | 9 | function strcat(string calldata a, string calldata b) 10 | external 11 | pure 12 | returns (string memory) 13 | { 14 | return string(abi.encodePacked(a, b)); 15 | } 16 | } 17 | -------------------------------------------------------------------------------- /contracts/Libraries/Tupler.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | contract LibTupler { 5 | function extractElement(bytes memory tuple, uint256 index) 6 | public 7 | pure 8 | returns (bytes32) 9 | { 10 | assembly { 11 | return(add(tuple, mul(add(index, 1), 32)), 32) 12 | } 13 | } 14 | } 15 | -------------------------------------------------------------------------------- /contracts/VM.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | 3 | pragma solidity ^0.8.11; 4 | 5 | import "./CommandBuilder.sol"; 6 | 7 | abstract contract VM { 8 | using CommandBuilder for bytes[]; 9 | 10 | uint256 constant FLAG_CT_DELEGATECALL = 0x00; 11 | uint256 constant FLAG_CT_CALL = 0x01; 12 | uint256 constant FLAG_CT_STATICCALL = 0x02; 13 | uint256 constant FLAG_CT_VALUECALL = 0x03; 14 | uint256 constant FLAG_CT_MASK = 0x03; 15 | uint256 constant FLAG_EXTENDED_COMMAND = 0x40; 16 | uint256 constant FLAG_TUPLE_RETURN = 0x80; 17 | 18 | uint256 constant SHORT_COMMAND_FILL = 0x000000000000FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF; 19 | 20 | address immutable self; 21 | 22 | error ExecutionFailed( 23 | uint256 command_index, 24 | address target, 25 | string message 26 | ); 27 | 28 | constructor() { 29 | self = address(this); 30 | } 31 | 32 | function _execute(bytes32[] calldata commands, bytes[] memory state) 33 | internal returns (bytes[] memory) 34 | { 35 | bytes32 command; 36 | uint256 flags; 37 | bytes32 indices; 38 | 39 | bool success; 40 | bytes memory outdata; 41 | 42 | uint256 commandsLength = commands.length; 43 | for (uint256 i; i < commandsLength;) { 44 | command = commands[i]; 45 | flags = uint256(command >> 216) & 0xFF; // more efficient 46 | // flags = uint256(uint8(bytes1(command << 32))); // more readable 47 | 48 | if (flags & FLAG_EXTENDED_COMMAND != 0) { 49 | indices = commands[++i]; 50 | } else { 51 | indices = bytes32(uint256(command << 40) | SHORT_COMMAND_FILL); 52 | } 53 | 54 | if (flags & FLAG_CT_MASK == FLAG_CT_DELEGATECALL) { 55 | (success, outdata) = address(uint160(uint256(command))).delegatecall( // target 56 | // inputs 57 | state.buildInputs( 58 | //selector 59 | bytes4(command), 60 | indices 61 | ) 62 | ); 63 | } else if (flags & FLAG_CT_MASK == FLAG_CT_CALL) { 64 | (success, outdata) = address(uint160(uint256(command))).call( // target 65 | // inputs 66 | state.buildInputs( 67 | //selector 68 | bytes4(command), 69 | indices 70 | ) 71 | ); 72 | } else if (flags & FLAG_CT_MASK == FLAG_CT_STATICCALL) { 73 | (success, outdata) = address(uint160(uint256(command))).staticcall( // target 74 | // inputs 75 | state.buildInputs( 76 | //selector 77 | bytes4(command), 78 | indices 79 | ) 80 | ); 81 | } else if (flags & FLAG_CT_MASK == FLAG_CT_VALUECALL) { 82 | uint256 callEth; 83 | bytes memory v = state[uint8(bytes1(indices))]; 84 | require(v.length == 32, "_execute: value call has no value indicated."); 85 | assembly { 86 | callEth := mload(add(v, 0x20)) 87 | } 88 | (success, outdata) = address(uint160(uint256(command))).call{ // target 89 | value: callEth 90 | }( 91 | // inputs 92 | state.buildInputs( 93 | //selector 94 | bytes4(command), 95 | bytes32(uint256(indices << 8) | CommandBuilder.IDX_END_OF_ARGS) 96 | ) 97 | ); 98 | } else { 99 | revert("Invalid calltype"); 100 | } 101 | 102 | if (!success) { 103 | if (outdata.length > 0) { 104 | assembly { 105 | outdata := add(outdata, 68) 106 | } 107 | } 108 | revert ExecutionFailed({ 109 | command_index: 0, 110 | target: address(uint160(uint256(command))), 111 | message: outdata.length > 0 ? string(outdata) : "Unknown" 112 | }); 113 | } 114 | 115 | if (flags & FLAG_TUPLE_RETURN != 0) { 116 | state.writeTuple(bytes1(command << 88), outdata); 117 | } else { 118 | state = state.writeOutputs(bytes1(command << 88), outdata); 119 | } 120 | unchecked{++i;} 121 | } 122 | return state; 123 | } 124 | } 125 | -------------------------------------------------------------------------------- /contracts/test/CommandBuilderHarness.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | import "../CommandBuilder.sol"; 5 | 6 | contract CommandBuilderHarness { 7 | using CommandBuilder for bytes[]; 8 | 9 | function basecall() public pure {} 10 | 11 | function testBuildInputsBaseGas( 12 | bytes[] memory state, 13 | bytes4 selector, 14 | bytes32 indices 15 | ) public view returns (bytes memory out) {} 16 | 17 | function testWriteOutputsBaseGas( 18 | bytes[] memory state, 19 | bytes1 index, 20 | bytes memory output 21 | ) public pure returns (bytes[] memory, bytes memory) { 22 | (index, output); // shh compiler 23 | return (state, new bytes(32)); 24 | } 25 | 26 | function testBuildInputs( 27 | bytes[] memory state, 28 | bytes4 selector, 29 | bytes32 indices 30 | ) public view returns (bytes memory) { 31 | bytes memory input = state.buildInputs(selector, indices); 32 | 33 | return input; 34 | } 35 | 36 | function testWriteOutputs( 37 | bytes[] memory state, 38 | bytes1 index, 39 | bytes memory output 40 | ) public pure returns (bytes[] memory, bytes memory) { 41 | state = state.writeOutputs(index, output); 42 | 43 | return (state, output); 44 | } 45 | } 46 | -------------------------------------------------------------------------------- /contracts/test/MultiReturn.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | contract MultiReturn { 5 | event Calculated(uint256 j); 6 | 7 | function intTuple() 8 | public 9 | pure 10 | returns ( 11 | uint256, 12 | uint256, 13 | uint256 14 | ) 15 | { 16 | return (0xbad, 0xdeed, 0xcafe); 17 | } 18 | 19 | function tupleConsumer(uint256 arg) public { 20 | emit Calculated(arg); 21 | } 22 | } 23 | -------------------------------------------------------------------------------- /contracts/test/Revert.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | contract Revert { 5 | function fail() public pure { 6 | require(false, "Hello World!"); 7 | } 8 | } 9 | -------------------------------------------------------------------------------- /contracts/test/Sender.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | contract Sender { 5 | function sender() public view returns (address) { 6 | return msg.sender; 7 | } 8 | } 9 | -------------------------------------------------------------------------------- /contracts/test/SimpleToken.sol: -------------------------------------------------------------------------------- 1 | // contracts/GLDToken.sol 2 | // SPDX-License-Identifier: MIT 3 | pragma solidity ^0.8.11; 4 | 5 | import "@openzeppelin/contracts/token/ERC20/ERC20.sol"; 6 | 7 | contract ExecutorToken is ERC20 { 8 | constructor(uint256 initialSupply) ERC20("Exec", "EXE") { 9 | _mint(msg.sender, initialSupply); 10 | } 11 | } 12 | -------------------------------------------------------------------------------- /contracts/test/StateTest.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | contract StateTest { 5 | function addSlots( 6 | uint256 dest, 7 | uint256 src, 8 | uint256 src2, 9 | bytes[] memory state 10 | ) public pure returns (bytes[] memory) { 11 | state[dest] = abi.encode( 12 | abi.decode(state[src], (uint256)) + 13 | abi.decode(state[src2], (uint256)) 14 | ); 15 | return state; 16 | } 17 | } 18 | -------------------------------------------------------------------------------- /contracts/test/SubPlanTests.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | import "../VM.sol"; 5 | 6 | contract TestSubplan is VM { 7 | function execute(bytes32[] calldata commands, bytes[] memory state) 8 | public 9 | payable 10 | returns (bytes[] memory) 11 | { 12 | return _execute(commands, state); 13 | } 14 | } 15 | 16 | contract TestReadonlySubplan is VM { 17 | function execute(bytes32[] calldata commands, bytes[] memory state) 18 | public 19 | payable 20 | { 21 | _execute(commands, state); 22 | } 23 | } 24 | 25 | contract TestMultiSubplan is VM { 26 | function execute( 27 | bytes32[] calldata commands, 28 | bytes32[] calldata commands2, 29 | bytes[] memory state 30 | ) public payable returns (bytes[] memory) { 31 | state = _execute(commands, state); 32 | state = _execute(commands2, state); 33 | return state; 34 | } 35 | } 36 | 37 | contract TestMultiStateSubplan is VM { 38 | function execute( 39 | bytes32[] calldata commands, 40 | bytes[] memory state, 41 | bytes[] memory state2 42 | ) public payable returns (bytes[] memory) { 43 | _execute(commands, state); 44 | return _execute(commands, state2); 45 | } 46 | } 47 | 48 | contract TestBadSubplan is VM { 49 | function execute(bytes32[] calldata commands, bytes[] memory state) 50 | public 51 | payable 52 | returns (int256) 53 | { 54 | _execute(commands, state); 55 | return 0; 56 | } 57 | } 58 | -------------------------------------------------------------------------------- /contracts/test/TestContract.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | contract TestContract { 5 | function useState(bytes[] memory state) public returns (bytes[] memory) { 6 | return state; 7 | } 8 | } 9 | -------------------------------------------------------------------------------- /contracts/test/TestableVM.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity ^0.8.11; 3 | 4 | import "../VM.sol"; 5 | 6 | contract TestableVM is VM { 7 | function execute(bytes32[] calldata commands, bytes[] memory state) 8 | public payable returns (bytes[] memory) { 9 | return _execute(commands, state); 10 | } 11 | } 12 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "weiroll-py" 3 | version = "0.2.2" 4 | description = "Build weiroll transactions with brownie" 5 | authors = ["FP "] 6 | license = "MIT" 7 | repository = "https://github.com/fp-crypto/weiroll-py" 8 | packages = [ 9 | { include = "./weiroll.py" }, 10 | ] 11 | 12 | [tool.poetry.dependencies] 13 | python = ">=3.10,<3.11" 14 | eth-brownie = ">=v1.20" 15 | 16 | [tool.poetry.dev-dependencies] 17 | 18 | [build-system] 19 | requires = ["poetry-core>=1.0.0"] 20 | build-backend = "poetry.core.masonry.api" 21 | -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from weiroll import WeirollContract 3 | 4 | 5 | @pytest.fixture(scope="function", autouse=True) 6 | def shared_setup(fn_isolation): 7 | pass 8 | 9 | 10 | @pytest.fixture(scope="session") 11 | def alice(accounts): 12 | yield accounts[0] 13 | 14 | 15 | @pytest.fixture(scope="session") 16 | def weiroll_vm(alice, TestableVM): 17 | vm = alice.deploy(TestableVM) 18 | yield vm 19 | 20 | 21 | @pytest.fixture(scope="session") 22 | def math(alice, Math): 23 | math_brownie = alice.deploy(Math) 24 | yield WeirollContract.createLibrary(math_brownie) 25 | 26 | 27 | @pytest.fixture(scope="session") 28 | def testContract(alice, TestContract): 29 | brownie_contract = alice.deploy(TestContract) 30 | yield WeirollContract.createLibrary(brownie_contract) 31 | 32 | 33 | @pytest.fixture(scope="session") 34 | def strings(alice, Strings): 35 | strings_brownie = alice.deploy(Strings) 36 | yield WeirollContract.createLibrary(strings_brownie) 37 | 38 | 39 | @pytest.fixture(scope="session") 40 | def subplanContract(alice, TestSubplan): 41 | brownie_contract = alice.deploy(TestSubplan) 42 | yield WeirollContract.createLibrary(brownie_contract) 43 | 44 | 45 | @pytest.fixture(scope="session") 46 | def multiSubplanContract(alice, TestMultiSubplan): 47 | brownie_contract = alice.deploy(TestMultiSubplan) 48 | yield WeirollContract.createLibrary(brownie_contract) 49 | 50 | 51 | @pytest.fixture(scope="session") 52 | def badSubplanContract(alice, TestBadSubplan): 53 | brownie_contract = alice.deploy(TestBadSubplan) 54 | yield WeirollContract.createLibrary(brownie_contract) 55 | 56 | 57 | @pytest.fixture(scope="session") 58 | def multiStateSubplanContract(alice, TestMultiStateSubplan): 59 | brownie_contract = alice.deploy(TestMultiStateSubplan) 60 | yield WeirollContract.createLibrary(brownie_contract) 61 | 62 | 63 | @pytest.fixture(scope="session") 64 | def readonlySubplanContract(alice, TestReadonlySubplan): 65 | brownie_contract = alice.deploy(TestReadonlySubplan) 66 | yield WeirollContract.createLibrary(brownie_contract) 67 | 68 | 69 | @pytest.fixture(scope="session") 70 | def tuple_helper(alice, TupleHelper): 71 | yield alice.deploy(TupleHelper) 72 | 73 | 74 | @pytest.fixture(scope="session") 75 | def tuple_helper_yul(alice, TupleHelperYul): 76 | yield alice.deploy(TupleHelperYul) 77 | 78 | 79 | @pytest.fixture(scope="session") 80 | def tuple_helper_vy(alice, TupleHelperVy): 81 | yield alice.deploy(TupleHelperVy) 82 | -------------------------------------------------------------------------------- /tests/test_chaining_actions.py: -------------------------------------------------------------------------------- 1 | from brownie import Contract, accounts, Wei, chain, TestableVM 2 | from weiroll import WeirollContract, WeirollPlanner, ReturnValue 3 | import requests 4 | 5 | 6 | def test_chaining_action(weiroll_vm, tuple_helper): 7 | whale = accounts.at("0x57757E3D981446D585Af0D9Ae4d7DF6D64647806", force=True) 8 | weth = Contract("0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2") 9 | yfi = Contract("0x0bc529c00C6401aEF6D220BE8C6Ea1667F6Ad93e") 10 | one_inch = Contract("0x1111111254fb6c44bAC0beD2854e76F90643097d") 11 | crv_yfi_weth = Contract("0x29059568bB40344487d62f7450E78b8E6C74e0e5") 12 | curve_swap = Contract("0xC26b89A667578ec7b3f11b2F98d6Fd15C07C54ba") 13 | 14 | # Check initial setup and send 10 eth to start the process 15 | assert weth.balanceOf(weiroll_vm.address) == 0 16 | assert yfi.balanceOf(weiroll_vm.address) == 0 17 | assert crv_yfi_weth.balanceOf(weiroll_vm.address) == 0 18 | weth.transfer(weiroll_vm.address, Wei("10 ether"), {"from": whale}) 19 | 20 | # Planner and all weiroll contracts 21 | planner = WeirollPlanner(weiroll_vm) 22 | w_one_inch = WeirollContract.createContract(one_inch) 23 | w_weth = WeirollContract.createContract(weth) 24 | w_yfi = WeirollContract.createContract(yfi) 25 | w_tuple_helper = WeirollContract.createContract(tuple_helper) 26 | w_curve_swap = WeirollContract.createContract(curve_swap) 27 | w_crv_yfi_weth = WeirollContract.createContract(crv_yfi_weth) 28 | 29 | # One inch section, eth->yfi 30 | planner.add(w_weth.approve(one_inch.address, 2**256-1)) 31 | swap_url = "https://api.1inch.io/v4.0/1/swap" 32 | r = requests.get( 33 | swap_url, 34 | params={ 35 | "fromTokenAddress": weth.address, 36 | "toTokenAddress": yfi.address, 37 | "amount": Wei("5 ether"), 38 | "fromAddress": weiroll_vm.address, 39 | "slippage": 5, 40 | "disableEstimate": "true", 41 | "allowPartialFill": "false", 42 | }, 43 | ) 44 | 45 | assert r.ok and r.status_code == 200 46 | tx = r.json()["tx"] 47 | 48 | decoded = one_inch.decode_input(tx["data"]) 49 | func_name = decoded[0] 50 | params = decoded[1] 51 | 52 | one_inch_ret = planner.add( 53 | w_one_inch.swap(*params).rawValue() 54 | ) 55 | 56 | # Since one inch's swap returns a tuple, we need to do an additional 57 | # action with the tuple helper contract to extract the amount value 58 | # in index 0. 59 | one_inch_amount = planner.add(w_tuple_helper.getElement(one_inch_ret, 0)) 60 | yfi_int_amount = ReturnValue('uint256', one_inch_amount.command) 61 | 62 | # Now that we have the yfi amount, let's do curve logic 63 | planner.add(w_weth.approve(w_curve_swap.address, 2**256-1)) 64 | planner.add(w_yfi.approve(w_curve_swap.address, 2**256-1)) 65 | 66 | curve_ret = planner.add( 67 | w_curve_swap.add_liquidity([Wei("5 ether"), yfi_int_amount], 0) 68 | ) 69 | 70 | planner.add(w_crv_yfi_weth.transfer(w_tuple_helper.address, curve_ret)) 71 | 72 | cmds, state = planner.plan() 73 | weiroll_tx = weiroll_vm.execute( 74 | cmds, state, {"from": whale, "gas_limit": 8_000_000, "gas_price": 0} 75 | ) 76 | 77 | assert crv_yfi_weth.balanceOf(w_tuple_helper.address) > 0 78 | -------------------------------------------------------------------------------- /tests/test_curve_add_liquidity.py: -------------------------------------------------------------------------------- 1 | from brownie import Contract, accounts, Wei, chain, TestableVM 2 | from weiroll import WeirollContract, WeirollPlanner 3 | 4 | 5 | def test_curve_add_liquidity(weiroll_vm): 6 | whale = accounts.at("0x5d3a536E4D6DbD6114cc1Ead35777bAB948E3643", force=True) 7 | dai = Contract("0x6B175474E89094C44Da98b954EedeAC495271d0F") 8 | curve_pool = Contract("0xbebc44782c7db0a1a60cb6fe97d0b483032ff1c7") 9 | three_crv = Contract("0x6c3F90f043a72FA612cbac8115EE7e52BDe6E490") 10 | 11 | # regular way 12 | assert three_crv.balanceOf(whale) == 0 13 | dai.approve(curve_pool.address, 2 ** 256 - 1, {"from": whale}) 14 | curve_pool.add_liquidity([Wei("10 ether"), 0, 0], 0, {"from": whale}) 15 | assert three_crv.balanceOf(whale) > 0 16 | 17 | dai.transfer(weiroll_vm.address, Wei("10 ether"), {"from": whale}) 18 | 19 | # Weiroll version 20 | planner = WeirollPlanner(weiroll_vm) 21 | 22 | w_dai = WeirollContract.createContract(dai) 23 | w_curve_pool = WeirollContract.createContract(curve_pool) 24 | 25 | planner.add(w_dai.approve(w_curve_pool.address, 2 ** 256 - 1)) 26 | w_dai_balance = planner.add(w_dai.balanceOf(weiroll_vm.address)) 27 | planner.add(w_curve_pool.add_liquidity([w_dai_balance, 0, 0], 0)) 28 | 29 | cmds, state = planner.plan() 30 | weiroll_tx = weiroll_vm.execute( 31 | cmds, state, {"from": whale, "gas_limit": 8_000_000, "gas_price": 0} 32 | ) 33 | 34 | assert three_crv.balanceOf(weiroll_vm.address) > 0 35 | 36 | 37 | def test_curve_add_liquidity_with_call(weiroll_vm): 38 | whale = accounts.at("0x5d3a536E4D6DbD6114cc1Ead35777bAB948E3643", force=True) 39 | dai = Contract("0x6B175474E89094C44Da98b954EedeAC495271d0F") 40 | curve_pool = Contract("0xbebc44782c7db0a1a60cb6fe97d0b483032ff1c7") 41 | three_crv = Contract("0x6c3F90f043a72FA612cbac8115EE7e52BDe6E490") 42 | 43 | dai.transfer(weiroll_vm.address, Wei("10 ether"), {"from": whale}) 44 | 45 | planner = WeirollPlanner(weiroll_vm) 46 | w_dai = WeirollContract.createContract(dai) 47 | w_curve_pool = WeirollContract.createContract(curve_pool) 48 | dai_balance = dai.balanceOf(weiroll_vm.address) 49 | 50 | planner.add(w_dai.approve(w_curve_pool.address, 2 ** 256 - 1)) 51 | planner.call(curve_pool, "add_liquidity", [dai_balance, 0, 0], 0) 52 | 53 | cmds, state = planner.plan() 54 | weiroll_tx = weiroll_vm.execute( 55 | cmds, state, {"from": whale, "gas_limit": 8_000_000, "gas_price": 0} 56 | ) 57 | 58 | assert three_crv.balanceOf(weiroll_vm.address) > 0 59 | -------------------------------------------------------------------------------- /tests/test_dyn.py: -------------------------------------------------------------------------------- 1 | from brownie import Contract, accounts, Wei, chain, TestableVM 2 | from weiroll import WeirollContract, WeirollPlanner, ReturnValue 3 | import requests 4 | 5 | 6 | def test_chaining_action(): 7 | 8 | settlement_address = "0x9008D19f58AAbD9eD0D60971565AA8510560ab41" 9 | c = Contract.from_explorer(settlement_address) 10 | wc = WeirollContract(c) 11 | -------------------------------------------------------------------------------- /tests/test_helpers.py: -------------------------------------------------------------------------------- 1 | from brownie.convert import to_bytes 2 | from hexbytes import HexBytes 3 | import random 4 | from brownie import reverts 5 | 6 | b2 = to_bytes(2) 7 | b1 = to_bytes(1) 8 | b4 = to_bytes(4) 9 | 10 | 11 | def test_insert(tuple_helper): 12 | 13 | assert ( 14 | HexBytes( 15 | tuple_helper.insertElement.transact(b2 + b1, 0, b4, False).return_value 16 | ) 17 | == b4 + b2 + b1 18 | ) 19 | assert ( 20 | HexBytes( 21 | tuple_helper.insertElement.transact(b2 + b1, 1, b4, False).return_value 22 | ) 23 | == b2 + b4 + b1 24 | ) 25 | assert ( 26 | HexBytes( 27 | tuple_helper.insertElement.transact(b2 + b1, 2, b4, False).return_value 28 | ) 29 | == b2 + b1 + b4 30 | ) 31 | 32 | with reverts(): 33 | tuple_helper.insertElement.transact(b2 + b1, 3, b4, False) 34 | 35 | rands = HexBytes( 36 | b"".join([to_bytes(random.randint(0, 2 ** 256 - 1)) for _ in range(100)]) 37 | ) 38 | 39 | for i in range(101): 40 | r = HexBytes( 41 | tuple_helper.insertElement.transact(rands, i, b4, False).return_value 42 | ) 43 | inserted = HexBytes(rands[: i * 32] + HexBytes(b4) + rands[i * 32 :]) 44 | assert r == inserted 45 | 46 | 47 | def test_replace(tuple_helper): 48 | 49 | assert ( 50 | HexBytes( 51 | tuple_helper.replaceElement.transact(b2 + b1, 0, b4, False).return_value 52 | ) 53 | == b4 + b1 54 | ) 55 | assert ( 56 | HexBytes( 57 | tuple_helper.replaceElement.transact(b2 + b1, 1, b4, False).return_value 58 | ) 59 | == b2 + b4 60 | ) 61 | 62 | with reverts(): 63 | tuple_helper.replaceElement.transact(b2 + b1, 2, b4, False) 64 | 65 | rands = HexBytes( 66 | b"".join([to_bytes(random.randint(0, 2 ** 256 - 1)) for _ in range(100)]) 67 | ) 68 | 69 | for i in range(100): 70 | r = HexBytes( 71 | tuple_helper.replaceElement.transact(rands, i, b4, False).return_value 72 | ) 73 | inserted = HexBytes(rands[: i * 32] + HexBytes(b4) + rands[(i + 1) * 32 :]) 74 | assert r == inserted 75 | -------------------------------------------------------------------------------- /tests/test_one_inch.py: -------------------------------------------------------------------------------- 1 | from brownie import Contract, accounts, Wei, chain, TestableVM 2 | from weiroll import WeirollContract, WeirollPlanner 3 | import requests 4 | 5 | 6 | def test_one_inch(weiroll_vm): 7 | whale = accounts.at("0x57757E3D981446D585Af0D9Ae4d7DF6D64647806", force=True) 8 | weth = Contract("0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2") 9 | crv = Contract("0xD533a949740bb3306d119CC777fa900bA034cd52") 10 | one_inch = Contract("0x1111111254fb6c44bAC0beD2854e76F90643097d") 11 | 12 | weth.transfer(weiroll_vm.address, Wei("10 ether"), {"from": whale}) 13 | 14 | swap_url = "https://api.1inch.io/v4.0/1/swap" 15 | r = requests.get( 16 | swap_url, 17 | params={ 18 | "fromTokenAddress": weth.address, 19 | "toTokenAddress": crv.address, 20 | "amount": Wei("10 ether"), 21 | "fromAddress": weiroll_vm.address, 22 | "slippage": 5, 23 | "disableEstimate": "true", 24 | "allowPartialFill": "false", 25 | }, 26 | ) 27 | 28 | assert r.ok and r.status_code == 200 29 | tx = r.json()["tx"] 30 | 31 | weth.approve(one_inch, 2 ** 256 - 1, {"from": weiroll_vm, "gas_price": 0}) 32 | 33 | decoded = one_inch.decode_input(tx["data"]) 34 | func_name = decoded[0] 35 | params = decoded[1] 36 | 37 | planner = WeirollPlanner(weiroll_vm) 38 | planner.call(one_inch, func_name, *params) 39 | 40 | cmds, state = planner.plan() 41 | weiroll_tx = weiroll_vm.execute( 42 | cmds, state, {"from": whale, "gas_limit": 8_000_000, "gas_price": 0} 43 | ) 44 | 45 | assert crv.balanceOf(weiroll_vm) > 0 46 | -------------------------------------------------------------------------------- /tests/test_swaps.py: -------------------------------------------------------------------------------- 1 | from brownie import TestableVM, Contract, convert 2 | from weiroll import WeirollContract, WeirollPlanner 3 | import weiroll 4 | from web3 import Web3 5 | import random 6 | import eth_abi 7 | import pytest 8 | 9 | 10 | def test_swaps(accounts, weiroll_vm): 11 | whale = accounts.at("0xF5BCE5077908a1b7370B9ae04AdC565EBd643966", force=True) 12 | 13 | weth = Contract("0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2") 14 | crvseth = Contract("0xc5424B857f758E906013F3555Dad202e4bdB4567") 15 | susd = Contract("0x57Ab1ec28D129707052df4dF418D58a2D46d5f51") 16 | 17 | weiroll_vm = accounts[0].deploy(TestableVM) 18 | planner = WeirollPlanner(whale) 19 | yvweth = WeirollContract.createContract( 20 | Contract("0xa258C4606Ca8206D8aA700cE2143D7db854D168c") 21 | ) 22 | weth = WeirollContract.createContract( 23 | Contract("0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2") 24 | ) 25 | susd = WeirollContract.createContract( 26 | Contract("0x57Ab1ec28D129707052df4dF418D58a2D46d5f51") 27 | ) 28 | seth = WeirollContract.createContract(Contract(crvseth.coins(1))) 29 | 30 | sushi_router_w = WeirollContract.createContract( 31 | Contract("0xd9e1cE17f2641f24aE83637ab66a2cca9C378B9F") 32 | ) 33 | univ3_router_w = WeirollContract.createContract( 34 | Contract("0xE592427A0AEce92De3Edee1F18E0157C05861564") 35 | ) 36 | 37 | yvweth.brownieContract.transfer(weiroll_vm, 2e18, {"from": whale}) 38 | weth.brownieContract.transfer(weiroll_vm, 1.118383e18, {"from": whale}) 39 | 40 | planner.call(yvweth.brownieContract, "withdraw(uint256)", int(1e18)) 41 | 42 | weth_bal = planner.add(weth.balanceOf(weiroll_vm.address)) 43 | 44 | planner.add(weth.approve(sushi_router_w.address, weth_bal)) 45 | planner.add( 46 | sushi_router_w.swapExactTokensForTokens( 47 | weth_bal, 0, [weth.address, susd.address], weiroll_vm.address, 2 ** 256 - 1 48 | ) 49 | ) 50 | 51 | susd_bal = planner.add(susd.balanceOf(weiroll_vm.address)) 52 | planner.add(susd.approve(sushi_router_w.address, susd_bal)) 53 | planner.add( 54 | sushi_router_w.swapExactTokensForTokens( 55 | susd_bal, 56 | 0, 57 | [susd.address, weth.address, seth.address], 58 | weiroll_vm.address, 59 | 2 ** 256 - 1, 60 | ) 61 | ) 62 | 63 | seth_bal = planner.add(seth.balanceOf(weiroll_vm.address)) 64 | planner.add(seth.approve(univ3_router_w.address, seth_bal)) 65 | planner.add( 66 | univ3_router_w.exactInputSingle( 67 | ( 68 | seth.address, 69 | weth.address, 70 | 500, 71 | weiroll_vm.address, 72 | 2 ** 256 - 1, 73 | seth_bal, 74 | 0, 75 | 0, 76 | ) 77 | ) 78 | ) 79 | 80 | cmds, state = planner.plan() 81 | weiroll_tx = weiroll_vm.execute( 82 | cmds, state, {"from": weiroll_vm, "gas_limit": 8_000_000, "gas_price": 0} 83 | ) 84 | 85 | 86 | @pytest.mark.skip("broken") 87 | def test_balancer_swap(accounts, weiroll_vm, tuple_helper): 88 | 89 | bal_whale = accounts.at("0xF977814e90dA44bFA03b6295A0616a897441aceC", force=True) 90 | bal_amount = random.randint(1000, 50000) * 10 ** 18 91 | 92 | planner = WeirollPlanner(weiroll_vm) 93 | 94 | bal = Contract("0xba100000625a3754423978a60c9317c58a424e3D") 95 | weth = Contract("0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2") 96 | balancer_vault = Contract("0xBA12222222228d8Ba445958a75a0704d566BF2C8") 97 | 98 | bal.transfer(weiroll_vm, bal_amount, {"from": bal_whale}) 99 | 100 | w_bal = WeirollContract.createContract(bal) 101 | w_balancer_vault = WeirollContract.createContract(balancer_vault) 102 | w_tuple_helper = WeirollContract(tuple_helper) 103 | 104 | w_bal_balance = planner.add(w_bal.balanceOf(weiroll_vm.address)) 105 | 106 | planner.add(w_bal.approve(w_balancer_vault.address, w_bal_balance)) 107 | 108 | bal_weth_pool_id = convert.to_bytes( 109 | "0x5c6ee304399dbdb9c8ef030ab642b10820db8f56000200000000000000000014", "bytes32" 110 | ) 111 | deadline = int(999999999999999999) 112 | 113 | min_out_weth_bal = int(bal_amount * 0.001) 114 | 115 | fund_settings = { 116 | "sender": weiroll_vm.address, 117 | "recipient": weiroll_vm.address, 118 | "fromInternalBalance": False, 119 | "toInternalBalance": False, 120 | } 121 | 122 | swap = { 123 | "poolId": bal_weth_pool_id, 124 | "assetIn": bal.address, 125 | "assetOut": weth.address, 126 | "amount": w_bal_balance, 127 | } 128 | swap_kind = int(0) # GIVEN_IN 129 | 130 | user_data = convert.to_bytes(bal_weth_pool_id, "bytes") 131 | 132 | swap_struct = ( 133 | swap["poolId"], 134 | swap_kind, 135 | Web3.toChecksumAddress(swap["assetIn"]), 136 | Web3.toChecksumAddress(swap["assetOut"]), 137 | 0, # replace with w_bal_balance, 138 | user_data, 139 | ) 140 | 141 | w_bal_balance = weiroll.ReturnValue("bytes32", w_bal_balance.command) 142 | swap_struct_layout = "(bytes32,uint8,address,address,uint256,bytes)" 143 | 144 | w_swap_struct = planner.add( 145 | w_tuple_helper.replaceElement( 146 | eth_abi.encode_single(swap_struct_layout, swap_struct), 147 | 4, 148 | w_bal_balance, 149 | True, 150 | ).rawValue() 151 | ) 152 | w_swap_struct = weiroll.ReturnValue(swap_struct_layout, w_swap_struct.command) 153 | 154 | fund_struct = ( 155 | Web3.toChecksumAddress(fund_settings["sender"]), 156 | fund_settings["fromInternalBalance"], 157 | Web3.toChecksumAddress(fund_settings["recipient"]), 158 | fund_settings["toInternalBalance"], 159 | ) 160 | 161 | planner.add( 162 | w_balancer_vault.swap(w_swap_struct, fund_struct, min_out_weth_bal, deadline) 163 | ) 164 | 165 | cmds, state = planner.plan() 166 | 167 | assert bal.balanceOf(weiroll_vm) > 0 168 | assert weth.balanceOf(weiroll_vm) == 0 169 | 170 | weiroll_tx = weiroll_vm.execute(cmds, state) 171 | weiroll_tx.call_trace(True) 172 | 173 | assert bal.balanceOf(weiroll_vm) == 0 174 | assert weth.balanceOf(weiroll_vm) > min_out_weth_bal 175 | -------------------------------------------------------------------------------- /tests/test_weiroll.py: -------------------------------------------------------------------------------- 1 | import brownie 2 | import eth_abi 3 | import pytest 4 | from hexbytes import HexBytes 5 | import json 6 | 7 | import weiroll 8 | from weiroll import eth_abi_encode_single 9 | 10 | 11 | def test_weiroll_contract(math): 12 | assert hasattr(math, "add") 13 | 14 | result = math.add(1, 2) 15 | 16 | assert result.contract == math 17 | # assert result.fragment.function.signature == math.add.signature 18 | assert result.fragment.inputs == ["uint256", "uint256"] 19 | assert result.fragment.name == "add" 20 | assert result.fragment.outputs == ["uint256"] 21 | assert result.fragment.signature == "0x771602f7" 22 | assert result.callvalue == 0 23 | assert result.flags == weiroll.CommandFlags.DELEGATECALL 24 | 25 | args = result.args 26 | assert len(args) == 2 27 | assert args[0].param == "uint256" 28 | assert args[0].value == eth_abi_encode_single("uint256", 1) 29 | assert args[1].param == "uint256" 30 | assert args[1].value == eth_abi_encode_single("uint256", 2) 31 | 32 | 33 | def test_weiroll_planner_adds(alice, math): 34 | planner = weiroll.WeirollPlanner(alice) 35 | sum1 = planner.add(math.add(1, 2)) 36 | sum2 = planner.add(math.add(3, 4)) 37 | planner.add(math.add(sum1, sum2)) 38 | 39 | assert len(planner.commands) == 3 40 | 41 | 42 | def test_weiroll_planner_simple_program(alice, math): 43 | planner = weiroll.WeirollPlanner(alice) 44 | planner.add(math.add(1, 2)) 45 | 46 | commands, state = planner.plan() 47 | 48 | assert len(commands) == 1 49 | # TODO: hexconcat? 50 | assert commands[0] == weiroll.hexConcat("0x771602f7000001ffffffffff", math.address) 51 | 52 | assert len(state) == 2 53 | assert state[0] == eth_abi_encode_single("uint", 1) 54 | assert state[1] == eth_abi_encode_single("uint", 2) 55 | 56 | 57 | def test_weiroll_deduplicates_identical_literals(alice, math): 58 | planner = weiroll.WeirollPlanner(alice) 59 | planner.add(math.add(1, 1)) 60 | commands, state = planner.plan() 61 | assert len(commands) == 1 62 | assert len(state) == 1 63 | assert state[0] == eth_abi_encode_single("uint", 1) 64 | 65 | 66 | def test_weiroll_with_return_value(alice, math): 67 | planner = weiroll.WeirollPlanner(alice) 68 | 69 | sum1 = planner.add(math.add(1, 2)) 70 | planner.add(math.add(sum1, 3)) 71 | commands, state = planner.plan() 72 | 73 | assert len(commands) == 2 74 | assert commands[0] == weiroll.hexConcat("0x771602f7000001ffffffff01", math.address) 75 | assert commands[1] == weiroll.hexConcat("0x771602f7000102ffffffffff", math.address) 76 | 77 | assert len(state) == 3 78 | assert state[0] == eth_abi_encode_single("uint", 1) 79 | assert state[1] == eth_abi_encode_single("uint", 2) 80 | assert state[2] == eth_abi_encode_single("uint", 3) 81 | 82 | 83 | def test_weiroll_with_state_slots_for_intermediate_values(alice, math): 84 | planner = weiroll.WeirollPlanner(alice) 85 | sum1 = planner.add(math.add(1, 1)) 86 | planner.add(math.add(1, sum1)) 87 | 88 | commands, state = planner.plan() 89 | 90 | assert len(commands) == 2 91 | assert commands[0] == weiroll.hexConcat("0x771602f7000000ffffffff01", math.address) 92 | assert commands[1] == weiroll.hexConcat("0x771602f7000001ffffffffff", math.address) 93 | 94 | assert len(state) == 2 95 | assert state[0] == eth_abi_encode_single("uint", 1) 96 | assert state[1] == b"" 97 | 98 | 99 | @pytest.mark.parametrize( 100 | "param,value,expected", 101 | [ 102 | ( 103 | "string", 104 | "Hello, world!", 105 | "0x000000000000000000000000000000000000000000000000000000000000000d48656c6c6f2c20776f726c642100000000000000000000000000000000000000", 106 | ), 107 | ], 108 | ) 109 | def test_weiroll_abi_encode_single(param, value, expected): 110 | expected = HexBytes(expected) 111 | print("expected:", expected) 112 | 113 | literalValue = HexBytes(eth_abi_encode_single(param, value)) 114 | print("literalValue:", literalValue) 115 | 116 | assert literalValue == expected 117 | 118 | 119 | def test_weiroll_takes_dynamic_arguments(alice, strings): 120 | test_str = "Hello, world!" 121 | 122 | planner = weiroll.WeirollPlanner(alice) 123 | planner.add(strings.strlen(test_str)) 124 | commands, state = planner.plan() 125 | 126 | assert len(commands) == 1 127 | assert commands[0] == weiroll.hexConcat( 128 | "0x367bbd780080ffffffffffff", strings.address 129 | ) 130 | 131 | print(state) 132 | assert len(state) == 1 133 | assert state[0] == eth_abi_encode_single("string", test_str) 134 | 135 | 136 | def test_weiroll_returns_dynamic_arguments(alice, strings): 137 | planner = weiroll.WeirollPlanner(alice) 138 | planner.add(strings.strcat("Hello, ", "world!")) 139 | commands, state = planner.plan() 140 | 141 | assert len(commands) == 1 142 | assert commands[0] == weiroll.hexConcat( 143 | "0xd824ccf3008081ffffffffff", strings.address 144 | ) 145 | 146 | assert len(state) == 2 147 | assert state[0] == eth_abi_encode_single("string", "Hello, ") 148 | assert state[1] == eth_abi_encode_single("string", "world!") 149 | 150 | 151 | def test_weiroll_takes_dynamic_argument_from_a_return_value(alice, strings): 152 | planner = weiroll.WeirollPlanner(alice) 153 | test_str = planner.add(strings.strcat("Hello, ", "world!")) 154 | planner.add(strings.strlen(test_str)) 155 | commands, state = planner.plan() 156 | 157 | assert len(commands) == 2 158 | assert commands[0] == weiroll.hexConcat( 159 | "0xd824ccf3008081ffffffff81", strings.address 160 | ) 161 | assert commands[1] == weiroll.hexConcat( 162 | "0x367bbd780081ffffffffffff", strings.address 163 | ) 164 | 165 | assert len(state) == 2 166 | assert state[0] == eth_abi_encode_single("string", "Hello, ") 167 | assert state[1] == eth_abi_encode_single("string", "world!") 168 | 169 | 170 | def test_weiroll_argument_counts_match(math): 171 | with pytest.raises(ValueError): 172 | math.add(1) 173 | 174 | 175 | def test_weiroll_func_takes_and_replaces_current_state(alice, testContract): 176 | planner = weiroll.WeirollPlanner(alice) 177 | 178 | planner.replaceState(testContract.useState(planner.state)) 179 | 180 | commands, state = planner.plan() 181 | 182 | assert len(commands) == 1 183 | assert commands[0] == weiroll.hexConcat( 184 | "0x08f389c800fefffffffffffe", testContract.address 185 | ) 186 | 187 | assert len(state) == 0 188 | 189 | 190 | def test_weiroll_supports_subplan(alice, math, subplanContract): 191 | subplanner = weiroll.WeirollPlanner(alice) 192 | subplanner.add(math.add(1, 2)) 193 | 194 | planner = weiroll.WeirollPlanner(alice) 195 | planner.addSubplan(subplanContract.execute(subplanner, subplanner.state)) 196 | 197 | commands, state = planner.plan() 198 | assert commands == [ 199 | weiroll.hexConcat("0xde792d5f0082fefffffffffe", subplanContract.address) 200 | ] 201 | 202 | assert len(state) == 3 203 | assert state[0] == eth_abi_encode_single("uint", 1) 204 | assert state[1] == eth_abi_encode_single("uint", 2) 205 | # TODO: javascript test is more complicated than this. but i think this is fine? 206 | assert state[2] == weiroll.hexConcat("0x771602f7000001ffffffffff", math.address) 207 | 208 | 209 | def test_weiroll_subplan_allows_return_in_parent_scope(alice, math, subplanContract): 210 | subplanner = weiroll.WeirollPlanner(alice) 211 | sum = subplanner.add(math.add(1, 2)) 212 | 213 | planner = weiroll.WeirollPlanner(alice) 214 | planner.addSubplan(subplanContract.execute(subplanner, subplanner.state)) 215 | planner.add(math.add(sum, 3)) 216 | 217 | commands, _ = planner.plan() 218 | assert len(commands) == 2 219 | # Invoke subplanner 220 | assert commands[0] == weiroll.hexConcat( 221 | "0xde792d5f0083fefffffffffe", subplanContract.address 222 | ) 223 | # sum + 3 224 | assert commands[1] == weiroll.hexConcat("0x771602f7000102ffffffffff", math.address) 225 | 226 | 227 | def test_weiroll_return_values_across_scopes(alice, math, subplanContract): 228 | subplanner1 = weiroll.WeirollPlanner(alice) 229 | sum = subplanner1.add(math.add(1, 2)) 230 | 231 | subplanner2 = weiroll.WeirollPlanner(alice) 232 | subplanner2.add(math.add(sum, 3)) 233 | 234 | planner = weiroll.WeirollPlanner(alice) 235 | planner.addSubplan(subplanContract.execute(subplanner1, subplanner1.state)) 236 | planner.addSubplan(subplanContract.execute(subplanner2, subplanner2.state)) 237 | 238 | commands, state = planner.plan() 239 | 240 | assert len(commands) == 2 241 | assert commands[0] == weiroll.hexConcat( 242 | "0xde792d5f0083fefffffffffe", subplanContract.address 243 | ) 244 | assert commands[1] == weiroll.hexConcat( 245 | "0xde792d5f0084fefffffffffe", subplanContract.address 246 | ) 247 | 248 | assert len(state) == 5 249 | # TODO: javascript tests were more complex than this 250 | assert state[4] == weiroll.hexConcat("0x771602f7000102ffffffffff", math.address) 251 | 252 | 253 | def test_weiroll_return_values_must_be_defined(alice, math): 254 | subplanner = weiroll.WeirollPlanner(alice) 255 | sum = subplanner.add(math.add(1, 2)) 256 | 257 | planner = weiroll.WeirollPlanner(alice) 258 | planner.add(math.add(sum, 3)) 259 | 260 | with pytest.raises(ValueError, match="Return value from 'add' is not visible here"): 261 | planner.plan() 262 | 263 | 264 | def test_weiroll_add_subplan_needs_args(alice, math, subplanContract): 265 | subplanner = weiroll.WeirollPlanner(alice) 266 | subplanner.add(math.add(1, 2)) 267 | 268 | planner = weiroll.WeirollPlanner(alice) 269 | 270 | with pytest.raises( 271 | ValueError, match="Subplans must take planner and state arguments" 272 | ): 273 | planner.addSubplan(subplanContract.execute(subplanner, [])) 274 | 275 | with pytest.raises( 276 | ValueError, match="Subplans must take planner and state arguments" 277 | ): 278 | planner.addSubplan(subplanContract.execute([], subplanner.state)) 279 | 280 | 281 | def test_weiroll_doesnt_allow_multiple_subplans_per_call( 282 | alice, math, multiSubplanContract 283 | ): 284 | subplanner = weiroll.WeirollPlanner(alice) 285 | subplanner.add(math.add(1, 2)) 286 | 287 | planner = weiroll.WeirollPlanner(alice) 288 | with pytest.raises(ValueError, match="Subplans can only take one planner argument"): 289 | planner.addSubplan( 290 | multiSubplanContract.execute(subplanner, subplanner, subplanner.state) 291 | ) 292 | 293 | 294 | def test_weiroll_doesnt_allow_state_array_per_call( 295 | alice, math, multiStateSubplanContract 296 | ): 297 | subplanner = weiroll.WeirollPlanner(alice) 298 | subplanner.add(math.add(1, 2)) 299 | 300 | planner = weiroll.WeirollPlanner(alice) 301 | with pytest.raises(ValueError, match="Subplans can only take one state argument"): 302 | planner.addSubplan( 303 | multiStateSubplanContract.execute( 304 | subplanner, subplanner.state, subplanner.state 305 | ) 306 | ) 307 | 308 | 309 | def test_weiroll_subplan_has_correct_return_type(alice, math, badSubplanContract): 310 | subplanner = weiroll.WeirollPlanner(alice) 311 | subplanner.add(math.add(1, 2)) 312 | 313 | planner = weiroll.WeirollPlanner(alice) 314 | with pytest.raises( 315 | ValueError, 316 | match=r"Subplans must return a bytes\[\] replacement state or nothing", 317 | ): 318 | planner.addSubplan(badSubplanContract.execute(subplanner, subplanner.state)) 319 | 320 | 321 | def test_forbid_infinite_loops(alice, subplanContract): 322 | planner = weiroll.WeirollPlanner(alice) 323 | planner.addSubplan(subplanContract.execute(planner, planner.state)) 324 | 325 | with pytest.raises(ValueError, match="A planner cannot contain itself"): 326 | planner.plan() 327 | 328 | 329 | def test_subplans_without_returns(alice, math, readonlySubplanContract): 330 | subplanner = weiroll.WeirollPlanner(alice) 331 | subplanner.add(math.add(1, 2)) 332 | 333 | planner = weiroll.WeirollPlanner(alice) 334 | planner.addSubplan(readonlySubplanContract.execute(subplanner, subplanner.state)) 335 | 336 | commands, _ = planner.plan() 337 | 338 | assert len(commands) == 1 339 | commands[0] == weiroll.hexConcat( 340 | "0xde792d5f0082feffffffffff", readonlySubplanContract.address 341 | ) 342 | 343 | 344 | def test_read_only_subplans_requirements(alice, math, readonlySubplanContract): 345 | """it does not allow return values from inside read-only subplans to be used outside them""" 346 | subplanner = weiroll.WeirollPlanner(alice) 347 | sum = subplanner.add(math.add(1, 2)) 348 | 349 | planner = weiroll.WeirollPlanner(alice) 350 | planner.addSubplan(readonlySubplanContract.execute(subplanner, subplanner.state)) 351 | planner.add(math.add(sum, 3)) 352 | 353 | with pytest.raises(ValueError, match="Return value from 'add' is not visible here"): 354 | planner.plan() 355 | 356 | 357 | @pytest.mark.xfail(reason="need to write this") 358 | def test_plan_with_loop(alice): 359 | target_calldata = ( 360 | "0xc6b6816900000000000000000000000000000000000000000000054b40b1f852bda0" 361 | ) 362 | 363 | """ 364 | [ 365 | '0x0000000000000000000000000000000000000000000000000000000000000005', 366 | '0x000000000000000000000000cecad69d7d4ed6d52efcfa028af8732f27e08f70', 367 | '0x0000000000000000000000000000000000000000000000000000000000000022c6b6816900000000000000000000000000000000000000000000054b40b1f852bda0000000000000000000000000000000000000000000000000000000000000' 368 | ] 369 | """ 370 | planner = weiroll.WeirollPlanner(alice) 371 | 372 | raise NotImplementedError 373 | 374 | 375 | def _test_more(math): 376 | 377 | # TODO: test for curve add_liquidity encoding 378 | """ 379 | 380 | expect(() => planner.plan()).to.throw( 381 | 'Return value from "add" is not visible here' 382 | ); 383 | }); 384 | 385 | it('plans CALLs', () => { 386 | let Math = Contract.createContract( 387 | new ethers.Contract(SAMPLE_ADDRESS, mathABI.abi) 388 | ); 389 | 390 | const planner = new Planner(); 391 | planner.add(Math.add(1, 2)); 392 | const { commands } = planner.plan(); 393 | 394 | expect(commands.length).to.equal(1); 395 | expect(commands[0]).to.equal( 396 | '0x771602f7010001ffffffffffeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee' 397 | ); 398 | }); 399 | 400 | it('plans STATICCALLs', () => { 401 | let Math = Contract.createContract( 402 | new ethers.Contract(SAMPLE_ADDRESS, mathABI.abi), 403 | CommandFlags.STATICCALL 404 | ); 405 | 406 | const planner = new Planner(); 407 | planner.add(Math.add(1, 2)); 408 | const { commands } = planner.plan(); 409 | 410 | expect(commands.length).to.equal(1); 411 | expect(commands[0]).to.equal( 412 | '0x771602f7020001ffffffffffeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee' 413 | ); 414 | }); 415 | 416 | it('plans STATICCALLs via .staticcall()', () => { 417 | let Math = Contract.createContract( 418 | new ethers.Contract(SAMPLE_ADDRESS, mathABI.abi) 419 | ); 420 | 421 | const planner = new Planner(); 422 | planner.add(Math.add(1, 2).staticcall()); 423 | const { commands } = planner.plan(); 424 | 425 | expect(commands.length).to.equal(1); 426 | expect(commands[0]).to.equal( 427 | '0x771602f7020001ffffffffffeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee' 428 | ); 429 | }); 430 | 431 | it('plans CALLs with value', () => { 432 | const Test = Contract.createContract( 433 | new ethers.Contract(SAMPLE_ADDRESS, ['function deposit(uint x) payable']) 434 | ); 435 | 436 | const planner = new Planner(); 437 | planner.add(Test.deposit(123).withValue(456)); 438 | 439 | const { commands } = planner.plan(); 440 | expect(commands.length).to.equal(1); 441 | expect(commands[0]).to.equal( 442 | '0xb6b55f25030001ffffffffffeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee' 443 | ); 444 | }); 445 | 446 | it('allows returns from other calls to be used for the value parameter', () => { 447 | const Test = Contract.createContract( 448 | new ethers.Contract(SAMPLE_ADDRESS, ['function deposit(uint x) payable']) 449 | ); 450 | 451 | const planner = new Planner(); 452 | const sum = planner.add(Math.add(1, 2)); 453 | planner.add(Test.deposit(123).withValue(sum)); 454 | 455 | const { commands } = planner.plan(); 456 | expect(commands.length).to.equal(2); 457 | expect(commands).to.deep.equal([ 458 | '0x771602f7000001ffffffff01eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee', 459 | '0xb6b55f25030102ffffffffffeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee', 460 | ]); 461 | }); 462 | 463 | it('does not allow value-calls for DELEGATECALL or STATICCALL', () => { 464 | expect(() => Math.add(1, 2).withValue(3)).to.throw( 465 | 'Only CALL operations can send value' 466 | ); 467 | 468 | const StaticMath = Contract.createContract( 469 | new ethers.Contract(SAMPLE_ADDRESS, mathABI.abi), 470 | CommandFlags.STATICCALL 471 | ); 472 | expect(() => StaticMath.add(1, 2).withValue(3)).to.throw( 473 | 'Only CALL operations can send value' 474 | ); 475 | }); 476 | 477 | it('does not allow making DELEGATECALL static', () => { 478 | expect(() => Math.add(1, 2).staticcall()).to.throw( 479 | 'Only CALL operations can be made static' 480 | ); 481 | }); 482 | 483 | it('uses extended commands where necessary', () => { 484 | const Test = Contract.createLibrary( 485 | new ethers.Contract(SAMPLE_ADDRESS, [ 486 | 'function test(uint a, uint b, uint c, uint d, uint e, uint f, uint g) returns(uint)', 487 | ]) 488 | ); 489 | 490 | const planner = new Planner(); 491 | planner.add(Test.test(1, 2, 3, 4, 5, 6, 7)); 492 | const { commands } = planner.plan(); 493 | expect(commands.length).to.equal(2); 494 | expect(commands[0]).to.equal( 495 | '0xe473580d40000000000000ffeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee' 496 | ); 497 | expect(commands[1]).to.equal( 498 | '0x00010203040506ffffffffffffffffffffffffffffffffffffffffffffffffff' 499 | ); 500 | }); 501 | 502 | it('supports capturing the whole return value as a bytes', () => { 503 | const Test = Contract.createLibrary( 504 | new ethers.Contract(SAMPLE_ADDRESS, [ 505 | 'function returnsTuple() returns(uint a, bytes32[] b)', 506 | 'function acceptsBytes(bytes raw)', 507 | ]) 508 | ); 509 | 510 | const planner = new Planner(); 511 | const ret = planner.add(Test.returnsTuple().rawValue()); 512 | planner.add(Test.acceptsBytes(ret)); 513 | const { commands } = planner.plan(); 514 | expect(commands).to.deep.equal([ 515 | '0x61a7e05e80ffffffffffff80eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee', 516 | '0x3e9ef66a0080ffffffffffffeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee', 517 | ]); 518 | }); 519 | """ 520 | -------------------------------------------------------------------------------- /weiroll.py: -------------------------------------------------------------------------------- 1 | import re 2 | from collections import defaultdict, namedtuple 3 | from enum import IntEnum, IntFlag 4 | from functools import cache 5 | from typing import Optional, Any 6 | from eth_typing.abi import TypeStr 7 | 8 | import brownie 9 | import eth_abi.codec 10 | import eth_abi.grammar 11 | import eth_abi.packed 12 | from eth_abi.registry import registry 13 | from brownie.convert.utils import get_type_strings 14 | from brownie.network.contract import OverloadedMethod 15 | from hexbytes import HexBytes 16 | 17 | MAX_UINT256 = 2**256 - 1 18 | 19 | # TODO: real types? 20 | Value = namedtuple("Value", "param") 21 | LiteralValue = namedtuple("LiteralValue", "param,value") 22 | ReturnValue = namedtuple("ReturnValue", "param,command") 23 | 24 | 25 | def simple_type_strings(inputs) -> tuple[Optional[list[str]], Optional[list[int]]]: 26 | """cut state variables that are too long into 32 byte chunks. 27 | 28 | related: https://github.com/weiroll/weiroll.js/pull/34 29 | """ 30 | 31 | if not inputs: 32 | return None, None 33 | 34 | simple_inputs = [] 35 | simple_sizes = [] 36 | for i in inputs: 37 | if i.endswith("]") and not i.endswith("[]") and not isDynamicType(i): 38 | # fixed size array. cut it up 39 | m = re.match(r"([a-z0-9]+)\[([0-9]+)\]", i) 40 | 41 | size = int(m.group(2)) 42 | 43 | simple_inputs.extend([m.group(1)] * size) 44 | simple_sizes.append(size) 45 | elif i.startswith("(") and i.endswith(")") and not isDynamicType(i): 46 | types = i[1:-1].split(",") 47 | 48 | simple_inputs.extend(types) 49 | simple_sizes.append(len(types)) 50 | else: 51 | simple_inputs.append(i) 52 | simple_sizes.append(1) 53 | 54 | if all([s == 1 for s in simple_sizes]): 55 | # if no inputs or all the inputs are easily handled sizes, we don't need to simplify them 56 | # we don't clear simple_inputs because its simpler for that to just be a copy of self.inputs 57 | simple_sizes = None 58 | 59 | return simple_inputs, simple_sizes 60 | 61 | 62 | def simple_args(simple_sizes, args): 63 | """split up complex types into 32 byte chunks that weiroll state can handle.""" 64 | if not simple_sizes: 65 | # no need to handle anything specially 66 | return args 67 | 68 | simplified = [] 69 | for i, size in enumerate(simple_sizes): 70 | if size == 1: 71 | # no need to do anything fancy 72 | simplified.append(args[i]) 73 | else: 74 | simplified.extend(args[i]) 75 | 76 | return simplified 77 | 78 | 79 | # TODO: not sure about this class. its mostly here because this is how the javascript sdk works. now that this works, i think we can start refactoring to use brownie more directly 80 | class FunctionFragment: 81 | def __init__(self, brownieContract: brownie.Contract, selector): 82 | function_name = brownieContract.selectors[selector] 83 | 84 | function = getattr(brownieContract, function_name) 85 | 86 | if isinstance(function, OverloadedMethod): 87 | overloaded_func = None 88 | for m in function.methods.values(): 89 | # TODO: everyone is inconsistent about signature vs selector vs name 90 | if m.signature == selector: 91 | overloaded_func = m 92 | break 93 | 94 | assert overloaded_func 95 | function = overloaded_func 96 | 97 | self.function = function 98 | self.name = function_name 99 | self.signature = function.signature 100 | self.inputs = get_type_strings(function.abi["inputs"]) 101 | 102 | # look at the inputs that aren't dynamic types but also aren't 32 bytes long and cut them up 103 | self.simple_inputs, self.simple_sizes = simple_type_strings(self.inputs) 104 | 105 | self.outputs = get_type_strings(function.abi["outputs"]) 106 | # TODO: do something to handle outputs of uncommon types? 107 | 108 | def encode_args(self, *args): 109 | if len(args) != len(self.inputs): 110 | raise ValueError( 111 | f"Function {self.name} has {len(self.inputs)} arguments but {len(self.args)} provided" 112 | ) 113 | 114 | # split up complex types into 32 byte chunks that weiroll state can handle 115 | args = simple_args(self.simple_sizes, args) 116 | 117 | return [encodeArg(arg, self.simple_inputs[i]) for (i, arg) in enumerate(args)] 118 | 119 | 120 | class StateValue: 121 | def __init__(self): 122 | self.param = "bytes[]" 123 | 124 | 125 | class SubplanValue: 126 | def __init__(self, planner): 127 | self.param = "bytes[]" 128 | self.planner = planner 129 | 130 | 131 | # TODO: use python ABC or something like that? 132 | def isValue(arg): 133 | if isinstance(arg, Value): 134 | return True 135 | if isinstance(arg, LiteralValue): 136 | return True 137 | if isinstance(arg, ReturnValue): 138 | return True 139 | if isinstance(arg, StateValue): 140 | return True 141 | if isinstance(arg, SubplanValue): 142 | return True 143 | return False 144 | 145 | 146 | # TODO: this needs tests! I'm 90% sure this is wrong for lists 147 | # TODO: does eth_utils not already have this? it seems like other people should have written something like this 148 | def hexConcat(*items) -> HexBytes: 149 | result = b"" 150 | for item in items: 151 | if isinstance(item, list): 152 | item = hexConcat(*item) 153 | else: 154 | item = HexBytes(item) 155 | result += bytes(item) 156 | return HexBytes(result) 157 | 158 | 159 | class CommandFlags(IntFlag): 160 | # Specifies that a call should be made using the DELEGATECALL opcode 161 | DELEGATECALL = 0x00 162 | # Specifies that a call should be made using the CALL opcode 163 | CALL = 0x01 164 | # Specifies that a call should be made using the STATICCALL opcode 165 | STATICCALL = 0x02 166 | # Specifies that a call should be made using the CALL opcode, and that the first argument will be the value to send 167 | CALL_WITH_VALUE = 0x03 168 | # A bitmask that selects calltype flags 169 | CALLTYPE_MASK = 0x03 170 | # Specifies that this is an extended command, with an additional command word for indices. Internal use only. 171 | EXTENDED_COMMAND = 0x40 172 | # Specifies that the return value of this call should be wrapped in a `bytes`. Internal use only. 173 | TUPLE_RETURN = 0x80 174 | 175 | 176 | class FunctionCall: 177 | def __init__( 178 | self, 179 | contract, 180 | flags: CommandFlags, 181 | fragment: FunctionFragment, 182 | args, 183 | callvalue=0, 184 | ): 185 | self.contract = contract 186 | self.flags = flags 187 | self.fragment = fragment 188 | self.args = args 189 | self.callvalue = callvalue 190 | 191 | def withValue(self, value): 192 | """ 193 | Returns a new [[FunctionCall]] that sends value with the call. 194 | @param value The value (in wei) to send with the call 195 | """ 196 | if (self.flags & CommandFlags.CALLTYPE_MASK) != CommandFlags.CALL and ( 197 | self.flags & CommandFlags.CALLTYPE_MASK 198 | ) != CommandFlags.CALL_WITH_VALUE: 199 | raise ValueError("Only CALL operations can send value") 200 | return self.__class__( 201 | self.contract, 202 | (self.flags & ~CommandFlags.CALLTYPE_MASK) | CommandFlags.CALL_WITH_VALUE, 203 | self.fragment, 204 | self.args, 205 | eth_abi_encode_single("uint", value), 206 | ) 207 | 208 | def rawValue(self): 209 | """ 210 | Returns a new [[FunctionCall]] whose return value will be wrapped as a `bytes`. 211 | This permits capturing the return values of functions with multiple return parameters, 212 | which weiroll does not otherwise support. 213 | """ 214 | return self.__class__( 215 | self.contract, 216 | self.flags | CommandFlags.TUPLE_RETURN, 217 | self.fragment, 218 | self.args, 219 | self.callvalue, 220 | ) 221 | 222 | def staticcall(self): 223 | """Returns a new [[FunctionCall]] that executes a STATICCALL instead of a regular CALL.""" 224 | if (self.flags & CommandFlags.CALLTYPE_MASK) != CommandFlags.CALL: 225 | raise ValueError("Only CALL operations can be made static") 226 | return self.__class__( 227 | self.contract, 228 | (self.flags & ~CommandFlags.CALLTYPE_MASK) | CommandFlags.STATICCALL, 229 | self.fragment, 230 | self.args, 231 | self.callvalue, 232 | ) 233 | 234 | 235 | def isDynamicType(param) -> bool: 236 | return eth_abi.grammar.parse(param).is_dynamic 237 | 238 | 239 | def encodeArg(arg, param): 240 | if isValue(arg): 241 | if arg.param != param: 242 | raise ValueError( 243 | f"Cannot pass value of type ${arg.param} to input of type ${param}" 244 | ) 245 | return arg 246 | if isinstance(arg, WeirollPlanner): 247 | return SubplanValue(arg) 248 | return LiteralValue(param, eth_abi_encode_single(param, arg)) 249 | 250 | 251 | class WeirollContract: 252 | """ 253 | * Provides a dynamically created interface to interact with Ethereum contracts via weiroll. 254 | * 255 | * Once created using the constructor or the [[Contract.createContract]] or [[Contract.createLibrary]] 256 | * functions, the returned object is automatically populated with methods that match those on the 257 | * supplied contract. For instance, if your contract has a method `add(uint, uint)`, you can call it on the 258 | * [[Contract]] object: 259 | * ```typescript 260 | * // Assumes `Math` is an ethers.js Contract instance. 261 | * const math = Contract.createLibrary(Math); 262 | * const result = math.add(1, 2); 263 | * ``` 264 | * 265 | * Calling a contract function returns a [[FunctionCall]] object, which you can pass to [[Planner.add]], 266 | * [[Planner.addSubplan]], or [[Planner.replaceState]] to add to the sequence of calls to plan. 267 | """ 268 | 269 | def __init__( 270 | self, brownieContract: brownie.Contract, commandFlags: CommandFlags = 0 271 | ): 272 | self.brownieContract = brownieContract 273 | self.address = brownieContract.address 274 | 275 | self.commandFlags = commandFlags 276 | self.functions = {} # aka functionsBySelector 277 | self.functionsBySignature = {} 278 | self.fragmentsBySelector = {} 279 | 280 | selectorsByName = defaultdict(list) 281 | 282 | for selector, name in self.brownieContract.selectors.items(): 283 | fragment = FunctionFragment(self.brownieContract, selector) 284 | 285 | # Check that the signature is unique; if not the ABI generation has 286 | # not been cleaned or may be incorrectly generated 287 | if selector in self.functions: 288 | raise ValueError(f"Duplicate ABI entry for selector: {selector}") 289 | 290 | self.fragmentsBySelector[selector] = fragment 291 | 292 | plan_fn = buildCall(self, fragment) 293 | 294 | # save this plan helper function fragment in self.functions 295 | self.functions[selector] = plan_fn 296 | 297 | # make the plan helper function available on self by selector 298 | setattr(self, selector, plan_fn) 299 | 300 | # Track unique names; we only expose bare named functions if they are ambiguous 301 | selectorsByName[name].append(selector) 302 | 303 | self.functionsByUniqueName = {} 304 | 305 | for name, selectors in selectorsByName.items(): 306 | # Ambiguous names to not get attached as bare names 307 | if len(selectors) == 1: 308 | if hasattr(self, name): 309 | # TODO: i think this is impossible 310 | raise ValueError("duplicate name!") 311 | 312 | plan_fn = self.functions[selectors[0]] 313 | 314 | # make the plan helper function available on self 315 | setattr(self, name, plan_fn) 316 | self.functionsByUniqueName[name] = plan_fn 317 | else: 318 | # define a new function which will use brownie' get_fn_from_args 319 | # to decide which plan_fn to route to 320 | def _overload(*args, fn_name=name): 321 | overload_method = self.brownieContract.__getattribute__(fn_name) 322 | method = overload_method._get_fn_from_args(args) 323 | signature = method.signature 324 | plan_fn = self.functions[signature] 325 | return plan_fn(*args) 326 | 327 | setattr(self, name, _overload) 328 | 329 | # attach full signatures (for methods with duplicate names) 330 | for selector in selectors: 331 | fragment = self.fragmentsBySelector[selector] 332 | 333 | signature = name + "(" + ",".join(fragment.inputs) + ")" 334 | 335 | plan_fn = self.functions[selector] 336 | 337 | self.functionsBySignature[signature] = plan_fn 338 | 339 | @classmethod 340 | @cache 341 | def createContract( 342 | cls, 343 | contract: brownie.Contract, 344 | commandflags=CommandFlags.CALL, 345 | ): 346 | """ 347 | Creates a [[Contract]] object from an ethers.js contract. 348 | All calls on the returned object will default to being standard CALL operations. 349 | Use this when you want your weiroll script to call a standard external contract. 350 | @param contract The ethers.js Contract object to wrap. 351 | @param commandflags Optionally specifies a non-default call type to use, such as 352 | [[CommandFlags.STATICCALL]]. 353 | """ 354 | assert commandflags != CommandFlags.DELEGATECALL 355 | return cls( 356 | contract, 357 | commandflags, 358 | ) 359 | 360 | @classmethod 361 | @cache 362 | def createLibrary( 363 | cls, 364 | contract: brownie.Contract, 365 | ): 366 | """ 367 | * Creates a [[Contract]] object from an ethers.js contract. 368 | * All calls on the returned object will default to being DELEGATECALL operations. 369 | * Use this when you want your weiroll script to call a library specifically designed 370 | * for use with weiroll. 371 | * @param contract The ethers.js Contract object to wrap. 372 | """ 373 | return cls(contract, CommandFlags.DELEGATECALL) 374 | 375 | # TODO: port getInterface? 376 | 377 | 378 | # TODO: not sure about this one. this was just how the javascript code worked, but can probably be refactored 379 | def buildCall(contract: WeirollContract, fragment: FunctionFragment): 380 | def _call(*args) -> FunctionCall: 381 | if len(args) != len(fragment.inputs): 382 | raise ValueError( 383 | f"Function {fragment.name} has {len(fragment.inputs)} arguments but {len(args)} provided" 384 | ) 385 | 386 | # TODO: maybe this should just be fragment.encode_args() 387 | encodedArgs = fragment.encode_args(*args) 388 | 389 | return FunctionCall( 390 | contract, 391 | contract.commandFlags, 392 | fragment, 393 | encodedArgs, 394 | ) 395 | 396 | return _call 397 | 398 | 399 | class CommandType(IntEnum): 400 | CALL = 1 401 | RAWCALL = 2 402 | SUBPLAN = 3 403 | 404 | 405 | Command = namedtuple("Command", "call,type") 406 | 407 | 408 | # returnSlotMap: Maps from a command to the slot used for its return value 409 | # literalSlotMap: Maps from a literal to the slot used to store it 410 | # freeSlots: An array of unused state slots 411 | # stateExpirations: Maps from a command to the slots that expire when it's executed 412 | # commandVisibility: Maps from a command to the last command that consumes its output 413 | # state: The initial state array 414 | PlannerState = namedtuple( 415 | "PlannerState", 416 | "returnSlotMap, literalSlotMap, freeSlots, stateExpirations, commandVisibility, state", 417 | ) 418 | 419 | 420 | def padArray(a, length, padValue) -> list: 421 | return a + [padValue] * (length - len(a)) 422 | 423 | 424 | class WeirollPlanner: 425 | def __init__(self, clone): 426 | self.state = StateValue() 427 | self.commands: list[Command] = [] 428 | self.unlimited_approvals = set() 429 | 430 | self.clone = clone 431 | 432 | def approve( 433 | self, token: brownie.Contract, spender: str, wei_needed, approve_wei=None 434 | ) -> Optional[ReturnValue]: 435 | key = (token, self.clone, spender) 436 | 437 | if approve_wei is None: 438 | approve_wei = MAX_UINT256 439 | 440 | if key in self.unlimited_approvals and approve_wei != 0: 441 | # we already planned an infinite approval for this token (and we aren't trying to set the approval to 0) 442 | return 443 | 444 | # check current allowance 445 | if token.allowance(self.clone, spender) >= wei_needed: 446 | return 447 | 448 | if approve_wei == MAX_UINT256: 449 | self.unlimited_approvals.add(key) 450 | 451 | return self.call(token, "approve", spender, approve_wei) 452 | 453 | def call(self, brownieContract: brownie.Contract, func_name, *args): 454 | """func_name can be just the name, or it can be the full signature. 455 | 456 | If there are multiple functions with the same name, you must use the signature. 457 | 458 | TODO: brownie has some logic for figuring out which overloaded method to use. we should use that here 459 | """ 460 | weirollContract = WeirollContract.createContract(brownieContract) 461 | 462 | if func_name.endswith(")"): 463 | # TODO: would be interesting to look at args and do this automatically 464 | func = weirollContract.functionsBySignature[func_name] 465 | else: 466 | func = weirollContract.functionsByUniqueName[func_name] 467 | 468 | return self.add(func(*args)) 469 | 470 | def delegatecall(self, brownieContract: brownie.Contract, func_name, *args): 471 | contract = WeirollContract.createLibrary(brownieContract) 472 | 473 | if func_name in contract.functionsByUniqueName: 474 | func = contract.functionsByUniqueName[func_name] 475 | elif func_name in contract.functionsBySignature: 476 | func = contract.functionsBySignature[func_name] 477 | else: 478 | # print("func_name:", func_name) 479 | # print("functionsByUniqueName:", contract.functionsByUniqueName) 480 | # print("functionsBySignature:", contract.functionsBySignature) 481 | raise ValueError(f"Unknown func_name ({func_name}) on {brownieContract}") 482 | 483 | return self.add(func(*args)) 484 | 485 | def add(self, call: FunctionCall) -> Optional[ReturnValue]: 486 | """ 487 | * Adds a new function call to the planner. Function calls are executed in the order they are added. 488 | * 489 | * If the function call has a return value, `add` returns an object representing that value, which you 490 | * can pass to subsequent function calls. For example: 491 | * ```typescript 492 | * const math = Contract.createLibrary(Math); // Assumes `Math` is an ethers.js contract object 493 | * const events = Contract.createLibrary(Events); // Assumes `Events` is an ethers.js contract object 494 | * const planner = new Planner(); 495 | * const sum = planner.add(math.add(21, 21)); 496 | * planner.add(events.logUint(sum)); 497 | * ``` 498 | * @param call The [[FunctionCall]] to add to the planner 499 | * @returns An object representing the return value of the call, or null if it does not return a value. 500 | """ 501 | command = Command(call, CommandType.CALL) 502 | self.commands.append(command) 503 | 504 | for arg in call.args: 505 | if isinstance(arg, SubplanValue): 506 | raise ValueError( 507 | "Only subplans can have arguments of type SubplanValue" 508 | ) 509 | 510 | if call.flags & CommandFlags.TUPLE_RETURN: 511 | return ReturnValue("bytes", command) 512 | 513 | # TODO: test this more 514 | if len(call.fragment.outputs) != 1: 515 | return None 516 | 517 | # print("call fragment outputs", call.fragment.outputs) 518 | 519 | return ReturnValue(call.fragment.outputs[0], command) 520 | 521 | def subcall(self, brownieContract: brownie.Contract, func_name, *args): 522 | """ 523 | * Adds a call to a subplan. This has the effect of instantiating a nested instance of the weiroll 524 | * interpreter, and is commonly used for functionality such as flashloans, control flow, or anywhere 525 | * else you may need to execute logic inside a callback. 526 | * 527 | * A [[FunctionCall]] passed to [[Planner.addSubplan]] must take another [[Planner]] object as one 528 | * argument, and a placeholder representing the planner state, accessible as [[Planner.state]], as 529 | * another. Exactly one of each argument must be provided. 530 | * 531 | * At runtime, the subplan is replaced by a list of commands for the subplanner (type `bytes32[]`), 532 | * and `planner.state` is replaced by the current state of the parent planner instance (type `bytes[]`). 533 | * 534 | * If the `call` returns a `bytes[]`, this will be used to replace the parent planner's state after 535 | * the call to the subplanner completes. Return values defined inside a subplan may be used outside that 536 | * subplan - both in the parent planner and in subsequent subplans - only if the `call` returns the 537 | * updated planner state. 538 | * 539 | * Example usage: 540 | * ``` 541 | * const exchange = Contract.createLibrary(Exchange); // Assumes `Exchange` is an ethers.js contract 542 | * const events = Contract.createLibrary(Events); // Assumes `Events` is an ethers.js contract 543 | * const subplanner = new Planner(); 544 | * const outqty = subplanner.add(exchange.swap(tokenb, tokena, qty)); 545 | * 546 | * const planner = new Planner(); 547 | * planner.addSubplan(exchange.flashswap(tokena, tokenb, qty, subplanner, planner.state)); 548 | * planner.add(events.logUint(outqty)); // Only works if `exchange.flashswap` returns updated state 549 | * ``` 550 | * @param call The [[FunctionCall]] to add to the planner. 551 | """ 552 | contract = WeirollContract.createContract(brownieContract) 553 | func = getattr(contract, func_name) 554 | func_call = func(*args) 555 | return self.addSubplan(func_call) 556 | 557 | def subdelegatecall(self, brownieContract: brownie.Contract, func_name, *args): 558 | contract = WeirollContract.createLibrary(brownieContract) 559 | func = getattr(contract, func_name) 560 | func_call = func(*args) 561 | return self.addSubplan(func_call) 562 | 563 | def addSubplan(self, call: FunctionCall): 564 | hasSubplan = False 565 | hasState = False 566 | 567 | for arg in call.args: 568 | if isinstance(arg, SubplanValue): 569 | if hasSubplan: 570 | raise ValueError("Subplans can only take one planner argument") 571 | hasSubplan = True 572 | elif isinstance(arg, StateValue): 573 | if hasState: 574 | raise ValueError("Subplans can only take one state argument") 575 | hasState = True 576 | if not hasSubplan or not hasState: 577 | raise ValueError("Subplans must take planner and state arguments") 578 | if ( 579 | call.fragment.outputs 580 | and len(call.fragment.outputs) == 1 581 | and call.fragment.outputs[0] != "bytes[]" 582 | ): 583 | raise ValueError( 584 | "Subplans must return a bytes[] replacement state or nothing" 585 | ) 586 | 587 | self.commands.append(Command(call, CommandType.SUBPLAN)) 588 | 589 | def replaceState(self, call: FunctionCall): 590 | """ 591 | * Executes a [[FunctionCall]], and replaces the planner state with the value it 592 | * returns. This can be used to execute functions that make arbitrary changes to 593 | * the planner state. Note that the planner library is not aware of these changes - 594 | * so it may produce invalid plans if you don't know what you're doing. 595 | * @param call The [[FunctionCall]] to execute 596 | """ 597 | if ( 598 | call.fragment.outputs and len(call.fragment.outputs) != 1 599 | ) or call.fragment.outputs[0] != "bytes[]": 600 | raise ValueError("Function replacing state must return a bytes[]") 601 | self.commands.append(Command(call, CommandType.RAWCALL)) 602 | 603 | def _preplan(self, commandVisibility, literalVisibility, seen=None, planners=None): 604 | if seen is None: 605 | seen: set[Command] = set() 606 | if planners is None: 607 | planners: set[WeirollPlanner] = set() 608 | 609 | if self in planners: 610 | raise ValueError("A planner cannot contain itself") 611 | planners.add(self) 612 | 613 | # Build visibility maps 614 | for command in self.commands: 615 | inargs = command.call.args 616 | if ( 617 | command.call.flags & CommandFlags.CALLTYPE_MASK 618 | == CommandFlags.CALL_WITH_VALUE 619 | ): 620 | if not command.call.callvalue: 621 | raise ValueError("Call with value must have a value parameter") 622 | inargs = [command.call.callvalue] + inargs 623 | 624 | for arg in inargs: 625 | if isinstance(arg, ReturnValue): 626 | if not arg.command in seen: 627 | raise ValueError( 628 | f"Return value from '{arg.command.call.fragment.name}' is not visible here" 629 | ) 630 | commandVisibility[arg.command] = command 631 | elif isinstance(arg, LiteralValue): 632 | literalVisibility[arg.value] = command 633 | elif isinstance(arg, SubplanValue): 634 | subplanSeen = seen # do not copy 635 | if not command.call.fragment.outputs: 636 | # Read-only subplan; return values aren't visible externally 637 | subplanSeen = set(seen) 638 | arg.planner._preplan( 639 | commandVisibility, literalVisibility, subplanSeen, planners 640 | ) 641 | elif not isinstance(arg, StateValue): 642 | raise ValueError(f"Unknown function argument type '{arg}'") 643 | 644 | seen.add(command) 645 | 646 | return commandVisibility, literalVisibility 647 | 648 | def _buildCommandArgs(self, command: Command, returnSlotMap, literalSlotMap, state): 649 | # Build a list of argument value indexes 650 | inargs = command.call.args 651 | if ( 652 | command.call.flags & CommandFlags.CALLTYPE_MASK 653 | == CommandFlags.CALL_WITH_VALUE 654 | ): 655 | if not command.call.callvalue: 656 | raise ValueError("Call with value must have a value parameter") 657 | inargs = [command.call.callvalue] + inargs 658 | 659 | args: list[int] = [] 660 | for arg in inargs: 661 | if isinstance(arg, ReturnValue): 662 | slot = returnSlotMap[arg.command] 663 | elif isinstance(arg, LiteralValue): 664 | slot = literalSlotMap[arg.value] 665 | elif isinstance(arg, StateValue): 666 | slot = 0xFE 667 | elif isinstance(arg, SubplanValue): 668 | # buildCommands has already built the subplan and put it in the last state slot 669 | slot = len(state) - 1 670 | else: 671 | raise ValueError(f"Unknown function argument type {arg}") 672 | if isDynamicType(arg.param): 673 | slot |= 0x80 674 | args.append(slot) 675 | 676 | return args 677 | 678 | def _buildCommands(self, ps: PlannerState) -> list[str]: 679 | encodedCommands = [] 680 | for command in self.commands: 681 | if command.type == CommandType.SUBPLAN: 682 | # find the subplan 683 | subplanner = next( 684 | arg for arg in command.call.args if isinstance(arg, SubplanValue) 685 | ).planner 686 | subcommands = subplanner._buildCommands(ps) 687 | ps.state.append( 688 | HexBytes(eth_abi_encode_single("bytes32[]", subcommands))[32:] 689 | ) 690 | # The slot is no longer needed after this command 691 | ps.freeSlots.append(len(ps.state) - 1) 692 | 693 | flags = command.call.flags 694 | 695 | args = self._buildCommandArgs( 696 | command, ps.returnSlotMap, ps.literalSlotMap, ps.state 697 | ) 698 | 699 | if len(args) > 6: 700 | flags |= CommandFlags.EXTENDED_COMMAND 701 | 702 | # Add any newly unused state slots to the list 703 | ps.freeSlots.extend(ps.stateExpirations[command]) 704 | 705 | ret = 0xFF 706 | if command in ps.commandVisibility: 707 | if command.type in [CommandType.RAWCALL, CommandType.SUBPLAN]: 708 | raise ValueError( 709 | f"Return value of {command.call.fragment.name} cannot be used to replace state and in another function" 710 | ) 711 | ret = len(ps.state) 712 | 713 | if len(ps.freeSlots) > 0: 714 | ret = ps.freeSlots.pop() 715 | 716 | # store the slot mapping 717 | ps.returnSlotMap[command] = ret 718 | 719 | # make the slot available when it's not needed 720 | expiryCommand = ps.commandVisibility[command] 721 | ps.stateExpirations[expiryCommand].append(ret) 722 | 723 | if ret == len(ps.state): 724 | ps.state.append(b"") 725 | 726 | if ( 727 | command.call.fragment.outputs 728 | and isDynamicType(command.call.fragment.outputs[0]) 729 | ) or command.call.flags & CommandFlags.TUPLE_RETURN != 0: 730 | ret |= 0x80 731 | elif command.type in [CommandType.RAWCALL, CommandType.SUBPLAN]: 732 | if ( 733 | command.call.fragment.outputs 734 | and len(command.call.fragment.outputs) == 1 735 | ): 736 | ret = 0xFE 737 | 738 | if flags & CommandFlags.EXTENDED_COMMAND == CommandFlags.EXTENDED_COMMAND: 739 | # extended command 740 | encodedCommands.extend( 741 | [ 742 | hexConcat( 743 | command.call.fragment.signature, 744 | flags, 745 | [0xFF] * 6, 746 | ret, 747 | command.call.contract.address, 748 | ), 749 | hexConcat(padArray(args, 32, 0xFF)), 750 | ] 751 | ) 752 | else: 753 | # standard command 754 | encodedCommands.append( 755 | hexConcat( 756 | command.call.fragment.signature, 757 | flags, 758 | padArray(args, 6, 0xFF), 759 | ret, 760 | command.call.contract.address, 761 | ) 762 | ) 763 | return encodedCommands 764 | 765 | def plan(self) -> tuple[list[str], list[str]]: 766 | # Tracks the last time a literal is used in the program 767 | literalVisibility: dict[str, Command] = {} 768 | # Tracks the last time a command's output is used in the program 769 | commandVisibility: dict[Command, Command] = {} 770 | 771 | self._preplan(commandVisibility, literalVisibility) 772 | 773 | # Maps from commands to the slots that expire on execution (if any) 774 | stateExpirations: dict[Command, list[int]] = defaultdict(list) 775 | 776 | # Tracks the state slot each literal is stored in 777 | literalSlotMap: dict[str, int] = {} 778 | 779 | state: list[str] = [] 780 | 781 | # Prepopulate the state and state expirations with literals 782 | for literal, lastCommand in literalVisibility.items(): 783 | slot = len(state) 784 | state.append(literal) 785 | literalSlotMap[literal] = slot 786 | stateExpirations[lastCommand].append(slot) 787 | 788 | ps: PlannerState = PlannerState( 789 | returnSlotMap={}, 790 | literalSlotMap=literalSlotMap, 791 | freeSlots=[], 792 | stateExpirations=stateExpirations, 793 | commandVisibility=commandVisibility, 794 | state=state, 795 | ) 796 | 797 | encodedCommands = self._buildCommands(ps) 798 | 799 | return encodedCommands, state 800 | 801 | 802 | class _ABIEncoder(eth_abi.codec.BaseABICoder): 803 | """ 804 | Wraps a registry to provide last-mile encoding functionality. 805 | """ 806 | 807 | def encode_single(self, typ: TypeStr, arg: Any) -> bytes: 808 | """ 809 | Encodes the python value ``arg`` as a binary value of the ABI type 810 | ``typ``. 811 | 812 | :param typ: The string representation of the ABI type that will be used 813 | for encoding e.g. ``'uint256'``, ``'bytes[]'``, ``'(int,int)'``, 814 | etc. 815 | :param arg: The python value to be encoded. 816 | 817 | :returns: The binary representation of the python value ``arg`` as a 818 | value of the ABI type ``typ``. 819 | """ 820 | encoder = self._registry.get_encoder(typ) 821 | 822 | return encoder(arg) 823 | 824 | 825 | eth_abi_encode_single = _ABIEncoder(registry).encode_single 826 | --------------------------------------------------------------------------------