feat: custom local dispatch flag

brownie-config.yaml

@@ -13,7 +13,7 @@ compiler:
     version: 0.8.11
     remappings:
       - "@openzeppelin=OpenZeppelin/openzeppelin-contracts@4.1.0"
-
+    viaIR: true
     optimizer:
       details:
           yul: true

contracts/VM.sol

@@ -15,6 +15,8 @@ abstract contract VM {
     uint256 constant FLAG_EXTENDED_COMMAND = 0x40;
     uint256 constant FLAG_TUPLE_RETURN = 0x80;
 
+    uint256 constant FLAG_LOCAL_DISPATCH = 0x20; // custom local dispatch flag
+
     uint256 constant SHORT_COMMAND_FILL = 0x000000000000FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF;
 
     address immutable self;
@@ -29,6 +31,8 @@ abstract contract VM {
         self = address(this);
     }
 
+    function dispatch(bytes memory inputs) internal virtual returns (bool success, bytes memory ret) {}
+
     function _execute(bytes32[] calldata commands, bytes[] memory state)
       internal returns (bytes[] memory)
     {
@@ -61,14 +65,27 @@ abstract contract VM {
                     )
                 );
             } else if (flags & FLAG_CT_MASK == FLAG_CT_CALL) {
-                (success, outdata) = address(uint160(uint256(command))).call( // target
-                    // inputs
-                    state.buildInputs(
-                        //selector
-                        bytes4(command),
-                        indices
-                    )
-                );
+                if (flags & FLAG_LOCAL_DISPATCH != 0) {
+                    address _target = address(uint160(uint256(command)));
+                    require(_target == address(this), "_execute: local dispatch must target VM.");
+                    (success, outdata) = dispatch(
+                        // inputs
+                        state.buildInputs(
+                            //selector
+                            bytes4(command),
+                            indices
+                        )
+                    );
+                } else {
+                    (success, outdata) = address(uint160(uint256(command))).call( // target
+                        // inputs
+                        state.buildInputs(
+                            //selector
+                            bytes4(command),
+                            indices
+                        )
+                    );
+                }
             } else if (flags & FLAG_CT_MASK == FLAG_CT_STATICCALL) {
                 (success, outdata) = address(uint160(uint256(command))).staticcall( // target
                     // inputs

contracts/test/TestableVMWithMath.sol

@@ -0,0 +1,79 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.11;
+
+import "../VM.sol";
+
+contract TestableVMWithMath is VM {
+    function execute(bytes32[] calldata commands, bytes[] memory state)
+        public
+        payable
+        returns (bytes[] memory)
+    {
+        return _execute(commands, state);
+    }
+
+    function sum(uint256 a, uint256 b) public pure returns (uint256) {
+        return a + b;
+    }
+
+    function sum3(uint256 a, uint256 b, uint256 c) public pure returns (uint256) {
+        return a + b;
+    }
+
+    function sub(uint256 a, uint256 b) public pure returns (uint256) {
+        return a - b;
+    }
+
+
+    function dispatch(bytes memory inputs)
+        internal
+        override
+        returns (bool _success, bytes memory _ret)
+    {
+        bytes4 _selector = bytes4(bytes32(inputs));
+        if (this.sum.selector == _selector) {
+            uint256 a;
+            uint256 b;
+            assembly {
+                a := mload(add(inputs, 36))
+                b := mload(add(inputs, 68))
+            }
+            uint256 res = sum(a, b);
+            _ret = new bytes(32);
+            assembly {
+                mstore(add(_ret, 32), res)
+            }
+            return (true, _ret);
+        }
+        if (this.sub.selector == _selector) {
+            uint256 a;
+            uint256 b;
+            assembly {
+                a := mload(add(inputs, 36))
+                b := mload(add(inputs, 68))
+            }
+            uint256 res = sub(a, b);
+            _ret = new bytes(32);
+            assembly {
+                mstore(add(_ret, 32), res)
+            }
+            return (true, _ret);
+        }
+        if (this.sum3.selector == _selector) {
+            uint256 a;
+            uint256 b;
+            uint256 c;
+            assembly {
+                a := mload(add(inputs, 36))
+                b := mload(add(inputs, 68))
+                c := mload(add(inputs, 100))
+            }
+            uint256 res = sum3(a, b, c);
+            _ret = new bytes(32);
+            assembly {
+                mstore(add(_ret, 32), res)
+            }
+            return (true, _ret);
+        }
+    }
+}

pyproject.toml

@@ -17,3 +17,6 @@ eth-brownie = "^v1.17"
 [build-system]
 requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
+
+[tool.black]
+line-length = 120

tests/conftest.py

@@ -23,6 +23,10 @@ def math(alice, Math):
     math_brownie = alice.deploy(Math)
     yield WeirollContract.createLibrary(math_brownie)
 
+@pytest.fixture(scope="module")
+def weiroll_vm_with_math(alice, TestableVMWithMath):
+    vm = alice.deploy(TestableVMWithMath)
+    yield vm
 
 @pytest.fixture(scope="module")
 def testContract(alice, TestContract):

tests/test_weiroll.py

@@ -123,9 +123,7 @@ def test_weiroll_takes_dynamic_arguments(alice, strings):
     commands, state = planner.plan()
 
     assert len(commands) == 1
-    assert commands[0] == weiroll.hexConcat(
-        "0x367bbd780080ffffffffffff", strings.address
-    )
+    assert commands[0] == weiroll.hexConcat("0x367bbd780080ffffffffffff", strings.address)
 
     print(state)
     assert len(state) == 1
@@ -138,9 +136,7 @@ def test_weiroll_returns_dynamic_arguments(alice, strings):
     commands, state = planner.plan()
 
     assert len(commands) == 1
-    assert commands[0] == weiroll.hexConcat(
-        "0xd824ccf3008081ffffffffff", strings.address
-    )
+    assert commands[0] == weiroll.hexConcat("0xd824ccf3008081ffffffffff", strings.address)
 
     assert len(state) == 2
     assert state[0] == eth_abi.encode_single("string", "Hello, ")
@@ -154,12 +150,8 @@ def test_weiroll_takes_dynamic_argument_from_a_return_value(alice, strings):
     commands, state = planner.plan()
 
     assert len(commands) == 2
-    assert commands[0] == weiroll.hexConcat(
-        "0xd824ccf3008081ffffffff81", strings.address
-    )
-    assert commands[1] == weiroll.hexConcat(
-        "0x367bbd780081ffffffffffff", strings.address
-    )
+    assert commands[0] == weiroll.hexConcat("0xd824ccf3008081ffffffff81", strings.address)
+    assert commands[1] == weiroll.hexConcat("0x367bbd780081ffffffffffff", strings.address)
 
     assert len(state) == 2
     assert state[0] == eth_abi.encode_single("string", "Hello, ")
@@ -179,9 +171,7 @@ def test_weiroll_func_takes_and_replaces_current_state(alice, testContract):
     commands, state = planner.plan()
 
     assert len(commands) == 1
-    assert commands[0] == weiroll.hexConcat(
-        "0x08f389c800fefffffffffffe", testContract.address
-    )
+    assert commands[0] == weiroll.hexConcat("0x08f389c800fefffffffffffe", testContract.address)
 
     assert len(state) == 0
 
@@ -194,9 +184,7 @@ def test_weiroll_supports_subplan(alice, math, subplanContract):
     planner.addSubplan(subplanContract.execute(subplanner, subplanner.state))
 
     commands, state = planner.plan()
-    assert commands == [
-        weiroll.hexConcat("0xde792d5f0082fefffffffffe", subplanContract.address)
-    ]
+    assert commands == [weiroll.hexConcat("0xde792d5f0082fefffffffffe", subplanContract.address)]
 
     assert len(state) == 3
     assert state[0] == eth_abi.encode_single("uint", 1)
@@ -216,9 +204,7 @@ def test_weiroll_subplan_allows_return_in_parent_scope(alice, math, subplanContr
     commands, _ = planner.plan()
     assert len(commands) == 2
     # Invoke subplanner
-    assert commands[0] == weiroll.hexConcat(
-        "0xde792d5f0083fefffffffffe", subplanContract.address
-    )
+    assert commands[0] == weiroll.hexConcat("0xde792d5f0083fefffffffffe", subplanContract.address)
     # sum + 3
     assert commands[1] == weiroll.hexConcat("0x771602f7000102ffffffffff", math.address)
 
@@ -237,12 +223,8 @@ def test_weiroll_return_values_across_scopes(alice, math, subplanContract):
     commands, state = planner.plan()
 
     assert len(commands) == 2
-    assert commands[0] == weiroll.hexConcat(
-        "0xde792d5f0083fefffffffffe", subplanContract.address
-    )
-    assert commands[1] == weiroll.hexConcat(
-        "0xde792d5f0084fefffffffffe", subplanContract.address
-    )
+    assert commands[0] == weiroll.hexConcat("0xde792d5f0083fefffffffffe", subplanContract.address)
+    assert commands[1] == weiroll.hexConcat("0xde792d5f0084fefffffffffe", subplanContract.address)
 
     assert len(state) == 5
     # TODO: javascript tests were more complex than this
@@ -266,43 +248,29 @@ def test_weiroll_add_subplan_needs_args(alice, math, subplanContract):
 
     planner = weiroll.WeirollPlanner(alice)
 
-    with pytest.raises(
-        ValueError, match="Subplans must take planner and state arguments"
-    ):
+    with pytest.raises(ValueError, match="Subplans must take planner and state arguments"):
         planner.addSubplan(subplanContract.execute(subplanner, []))
 
-    with pytest.raises(
-        ValueError, match="Subplans must take planner and state arguments"
-    ):
+    with pytest.raises(ValueError, match="Subplans must take planner and state arguments"):
         planner.addSubplan(subplanContract.execute([], subplanner.state))
 
 
-def test_weiroll_doesnt_allow_multiple_subplans_per_call(
-    alice, math, multiSubplanContract
-):
+def test_weiroll_doesnt_allow_multiple_subplans_per_call(alice, math, multiSubplanContract):
     subplanner = weiroll.WeirollPlanner(alice)
     subplanner.add(math.add(1, 2))
 
     planner = weiroll.WeirollPlanner(alice)
     with pytest.raises(ValueError, match="Subplans can only take one planner argument"):
-        planner.addSubplan(
-            multiSubplanContract.execute(subplanner, subplanner, subplanner.state)
-        )
+        planner.addSubplan(multiSubplanContract.execute(subplanner, subplanner, subplanner.state))
 
 
-def test_weiroll_doesnt_allow_state_array_per_call(
-    alice, math, multiStateSubplanContract
-):
+def test_weiroll_doesnt_allow_state_array_per_call(alice, math, multiStateSubplanContract):
     subplanner = weiroll.WeirollPlanner(alice)
     subplanner.add(math.add(1, 2))
 
     planner = weiroll.WeirollPlanner(alice)
     with pytest.raises(ValueError, match="Subplans can only take one state argument"):
-        planner.addSubplan(
-            multiStateSubplanContract.execute(
-                subplanner, subplanner.state, subplanner.state
-            )
-        )
+        planner.addSubplan(multiStateSubplanContract.execute(subplanner, subplanner.state, subplanner.state))
 
 
 def test_weiroll_subplan_has_correct_return_type(alice, math, badSubplanContract):
@@ -335,9 +303,7 @@ def test_subplans_without_returns(alice, math, readonlySubplanContract):
     commands, _ = planner.plan()
 
     assert len(commands) == 1
-    commands[0] == weiroll.hexConcat(
-        "0xde792d5f0082feffffffffff", readonlySubplanContract.address
-    )
+    commands[0] == weiroll.hexConcat("0xde792d5f0082feffffffffff", readonlySubplanContract.address)
 
 
 def test_read_only_subplans_requirements(alice, math, readonlySubplanContract):
@@ -355,9 +321,7 @@ def test_read_only_subplans_requirements(alice, math, readonlySubplanContract):
 
 @pytest.mark.xfail(reason="need to write this")
 def test_plan_with_loop(alice):
-    target_calldata = (
-        "0xc6b6816900000000000000000000000000000000000000000000054b40b1f852bda0"
-    )
+    target_calldata = "0xc6b6816900000000000000000000000000000000000000000000054b40b1f852bda0"
 
     """
     [

tests/test_weiroll_local.py

@@ -0,0 +1,60 @@
+from brownie import Contract, accounts, Wei, chain, TestableVM
+from weiroll import WeirollContract, WeirollPlanner
+import pytest
+
+
+@pytest.mark.skip()
+def test_vm_with_math(weiroll_vm_with_math):
+    weiroll_vm = weiroll_vm_with_math
+    whale = accounts.at("0x57757E3D981446D585Af0D9Ae4d7DF6D64647806", force=True)
+
+    planner = WeirollPlanner(weiroll_vm)
+    w_math = WeirollContract.createContract(weiroll_vm)
+    sum = planner.add(w_math.sum(1, 2))
+    sum_2 = planner.add(w_math.sum(3, sum))
+    sum_3 = planner.add(w_math.sum3(3, sum_2, 4))
+    planner.add(w_math.sub(sum_3, 3))
+
+    cmds, state = planner.plan()
+    weiroll_tx = weiroll_vm.execute(
+        cmds, state, {"from": whale, "gas_limit": 8_000_000, "gas_price": 0}
+    )
+
+
+# @pytest.mark.skip()
+def test_vm_with_math_local_dispatch(weiroll_vm_with_math):
+    weiroll_vm = weiroll_vm_with_math
+    whale = accounts.at("0x57757E3D981446D585Af0D9Ae4d7DF6D64647806", force=True)
+
+    planner = WeirollPlanner(weiroll_vm)
+    w_math = WeirollContract.createContract(weiroll_vm)
+    sum = planner.add(w_math.sum(1, 2).localDispatch())
+    sum_2 = planner.add(w_math.sum(3, sum).localDispatch())
+    sum_3 = planner.add(w_math.sum3(3, sum_2, 4).localDispatch())
+    planner.add(w_math.sub(sum_3, 3).localDispatch())
+
+    cmds, state = planner.plan()
+    weiroll_tx = weiroll_vm.execute(
+        cmds, state, {"from": whale, "gas_limit": 8_000_000, "gas_price": 0}
+    )
+
+
+# @pytest.mark.skip()
+def test_vm_with_math_auto_local_dispatch(weiroll_vm_with_math):
+    weiroll_vm = weiroll_vm_with_math
+    whale = accounts.at("0x57757E3D981446D585Af0D9Ae4d7DF6D64647806", force=True)
+
+    planner = WeirollPlanner(weiroll_vm, auto_local_dispatch=True)
+    w_math = WeirollContract.createContract(weiroll_vm)
+    sum = planner.add(w_math.sum(1, 2))
+    sum_2 = planner.add(w_math.sum(3, sum))
+    sum_3 = planner.add(w_math.sum3(3, sum_2, 4))
+    planner.add(w_math.sub(sum_3, 3))
+
+    for command in planner.commands:
+        assert(command.call.flags >> 5 & 0x1)
+
+    cmds, state = planner.plan()
+    weiroll_tx = weiroll_vm.execute(
+        cmds, state, {"from": whale, "gas_limit": 8_000_000, "gas_price": 0}
+    )