feat: local dispatch flag

fp-crypto · Aug 28, 2023 · a33b67c · a33b67c
1 parent 08c19f1
commit a33b67c
Show file tree

Hide file tree

Showing 4 changed files with 91 additions and 20 deletions.
diff --git a/contracts/VM.sol b/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;
@@ -63,20 +65,19 @@ abstract contract VM {
                     )
                 );
             } else if (flags & FLAG_CT_MASK == FLAG_CT_CALL) {
-                address _target = address(uint160(uint256(command)));
-                bytes memory inputs = state.buildInputs(
-                    //selector
-                    bytes4(command),
-                    indices
-                ); 
-                success = false;
-
-                if (_target == address(this)) {
-                    (success, outdata) = dispatch(inputs);
-                }
-
-                if (!success) {
-                    (success, outdata) = _target.call( // target
+                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

diff --git a/contracts/test/TestableVMWithMath.sol b/contracts/test/TestableVMWithMath.sol
@@ -16,6 +16,14 @@ contract TestableVMWithMath is VM {
         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
@@ -37,5 +45,35 @@ contract TestableVMWithMath is VM {
             }
             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);
+        }
     }
 }
diff --git a/tests/test_weiroll_local.py b/tests/test_weiroll_local.py
@@ -7,14 +7,33 @@ def test_vm_with_math(weiroll_vm_with_math):
     whale = accounts.at("0x57757E3D981446D585Af0D9Ae4d7DF6D64647806", force=True)
 
     planner = WeirollPlanner(weiroll_vm)
-    sum = planner.call(weiroll_vm, "sum", 1, 2)
-    sum_2 = planner.call(weiroll_vm, "sum", sum, 3)
-    sum_2 = planner.call(weiroll_vm, "sum", 3, sum_2)
+    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}
     )
 
-    print(weiroll_tx.return_value)
-    #assert False
+    assert False
+
+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}
+    )
+
+    assert False
diff --git a/weiroll.py b/weiroll.py
@@ -167,7 +167,8 @@ class CommandFlags(IntFlag):
     EXTENDED_COMMAND = 0x40
     # Specifies that the return value of this call should be wrapped in a `bytes`. Internal use only.
     TUPLE_RETURN = 0x80
-
+    # Specifies to use fast local dispatcher. Custom flag.
+    LOCAL_DISPATCH = 0x20
 
 class FunctionCall:
     def __init__(self, contract, flags: CommandFlags, fragment: FunctionFragment, args, callvalue=0):
@@ -220,6 +221,18 @@ def staticcall(self):
             self.callvalue,
         )
 
+    def localDispatch(self):
+        """
+        Returns a new [[FunctionCall]] with local dispatch specified 
+        """
+        return self.__class__(
+            self.contract,
+            self.flags | CommandFlags.LOCAL_DISPATCH,
+            self.fragment,
+            self.args,
+            self.callvalue,
+        )
+
 
 def isDynamicType(param) -> bool:
     return eth_abi.grammar.parse(param).is_dynamic