feat(world-modules): add SystemSwitch util (#1665)

.changeset/modern-stingrays-kneel.md

@@ -0,0 +1,19 @@
+---
+"@latticexyz/world-modules": minor
+---
+
+Since [#1564](https://github.com/latticexyz/mud/pull/1564) the World can no longer call itself via an external call.
+This made the developer experience of calling other systems via root systems worse, since calls from root systems are executed from the context of the World.
+The recommended approach is to use `delegatecall` to the system if in the context of a root system, and an external call via the World if in the context of a non-root system.
+To bring back the developer experience of calling systems from other sysyems without caring about the context in which the call is executed, we added the `SystemSwitch` util.
+
+```diff
+- // Instead of calling the system via an external call to world...
+- uint256 value = IBaseWorld(_world()).callMySystem();
+
++ // ...you can now use the `SystemSwitch` util.
++ // This works independent of whether used in a root system or non-root system.
++ uint256 value = abi.decode(SystemSwitch.call(abi.encodeCall(IBaseWorld.callMySystem, ()), (uint256));
+```
+
+Note that if you already know your system is always executed as non-root system, you can continue to use the approach of calling other systems via the `IBaseWorld(world)`.

packages/world-modules/src/modules/uniqueentity/getUniqueEntity.sol

@@ -1,10 +1,13 @@
 // SPDX-License-Identifier: MIT
 pragma solidity >=0.8.21;
 
-import { StoreSwitch } from "@latticexyz/store/src/StoreSwitch.sol";
 import { IBaseWorld } from "@latticexyz/world/src/codegen/interfaces/IBaseWorld.sol";
 
 import { IUniqueEntitySystem } from "../../interfaces/IUniqueEntitySystem.sol";
+import { UniqueEntitySystem } from "./UniqueEntitySystem.sol";
+
+import { SystemSwitch } from "../../utils/SystemSwitch.sol";
+import { SYSTEM_ID } from "./constants.sol";
 
 /**
  * Increment and get an entity nonce.
@@ -13,8 +16,7 @@ import { IUniqueEntitySystem } from "../../interfaces/IUniqueEntitySystem.sol";
  * For usage outside of a World, use the overload that takes an explicit store argument.
  */
 function getUniqueEntity() returns (bytes32 uniqueEntity) {
-  address world = StoreSwitch.getStoreAddress();
-  return IUniqueEntitySystem(world).uniqueEntity_system_getUniqueEntity();
+  return abi.decode(SystemSwitch.call(SYSTEM_ID, abi.encodeCall(UniqueEntitySystem.getUniqueEntity, ())), (bytes32));
 }
 
 /**

packages/world-modules/src/utils/SystemSwitch.sol

@@ -0,0 +1,82 @@
+// SPDX-License-Identifier: MIT
+pragma solidity >=0.8.21;
+
+import { Hook } from "@latticexyz/store/src/Hook.sol";
+import { Bytes } from "@latticexyz/store/src/Bytes.sol";
+
+import { IWorldKernel } from "@latticexyz/world/src/IWorldKernel.sol";
+import { ResourceId, WorldResourceIdInstance } from "@latticexyz/world/src/WorldResourceId.sol";
+import { WorldContextProviderLib, WorldContextConsumerLib } from "@latticexyz/world/src/WorldContext.sol";
+import { AccessControl } from "@latticexyz/world/src/AccessControl.sol";
+import { ROOT_NAMESPACE } from "@latticexyz/world/src/constants.sol";
+import { revertWithBytes } from "@latticexyz/world/src/revertWithBytes.sol";
+import { BEFORE_CALL_SYSTEM, AFTER_CALL_SYSTEM } from "@latticexyz/world/src/systemHookTypes.sol";
+
+import { IWorldErrors } from "@latticexyz/world/src/IWorldErrors.sol";
+import { ISystemHook } from "@latticexyz/world/src/ISystemHook.sol";
+
+import { FunctionSelectors } from "@latticexyz/world/src/codegen/tables/FunctionSelectors.sol";
+import { Systems } from "@latticexyz/world/src/codegen/tables/Systems.sol";
+import { SystemHooks } from "@latticexyz/world/src/codegen/tables/SystemHooks.sol";
+import { Balances } from "@latticexyz/world/src/codegen/tables/Balances.sol";
+
+/**
+ * @title SystemSwitch
+ * @dev The SystemSwitch library provides functions for interacting with systems from other systems.
+ */
+library SystemSwitch {
+  using WorldResourceIdInstance for ResourceId;
+
+  /**
+   * @notice Calls a system identified by its Resource ID.
+   * @dev Reverts if the system is not found, or if the system call reverts.
+   * If the call is executed from the root context, the system is called directly via delegatecall.
+   * Otherwise, the call is executed via an external call to the World contract.
+   * @param systemId The unique Resource ID of the system being called.
+   * @param callData The calldata to be executed in the system.
+   * @return returnData The return data from the system call.
+   */
+  function call(ResourceId systemId, bytes memory callData) internal returns (bytes memory returnData) {
+    address worldAddress = WorldContextConsumerLib._world();
+
+    // If we're in the World context, call the system directly via delegatecall
+    if (address(this) == worldAddress) {
+      (address systemAddress, ) = Systems.get(systemId);
+      // Check if the system exists
+      if (systemAddress == address(0)) revert IWorldErrors.World_ResourceNotFound(systemId, systemId.toString());
+
+      bool success;
+      (success, returnData) = WorldContextProviderLib.delegatecallWithContext({
+        msgSender: WorldContextConsumerLib._msgSender(),
+        msgValue: WorldContextConsumerLib._msgValue(),
+        target: systemAddress,
+        callData: callData
+      });
+
+      if (!success) revertWithBytes(returnData);
+      return returnData;
+    }
+
+    // Otherwise, call the system via world.call
+    returnData = IWorldKernel(worldAddress).call(systemId, callData);
+  }
+
+  /**
+   * @notice Calls a system via the function selector registered for it in the World contract.
+   * @dev Reverts if the system is not found, or if the system call reverts.
+   * If the call is executed from the root context, the system is called directly via delegatecall.
+   * Otherwise, the call is executed via an external call to the World contract.
+   * @param callData The world function selector, and call data to be forwarded to the system.
+   * @return returnData The return data from the system call.
+   */
+  function call(bytes memory callData) internal returns (bytes memory returnData) {
+    // Get the systemAddress and systemFunctionSelector from the worldFunctionSelector encoded in the calldata
+    (ResourceId systemId, bytes4 systemFunctionSelector) = FunctionSelectors.get(bytes4(callData));
+
+    // Revert if the function selector is not found
+    if (ResourceId.unwrap(systemId) == 0) revert IWorldErrors.World_FunctionSelectorNotFound(msg.sig);
+
+    // Replace function selector in the calldata with the system function selector, and call the system
+    return call({ systemId: systemId, callData: Bytes.setBytes4(callData, 0, systemFunctionSelector) });
+  }
+}

packages/world-modules/test/SystemSwitch.t.sol

@@ -0,0 +1,225 @@
+// SPDX-License-Identifier: MIT
+pragma solidity >=0.8.21;
+
+import { Test } from "forge-std/Test.sol";
+import { GasReporter } from "@latticexyz/gas-report/src/GasReporter.sol";
+
+import { System } from "@latticexyz/world/src/System.sol";
+import { IBaseWorld } from "@latticexyz/world/src/codegen/interfaces/IBaseWorld.sol";
+import { World } from "@latticexyz/world/src/World.sol";
+import { CoreModule } from "@latticexyz/world/src/modules/core/CoreModule.sol";
+import { ResourceId, WorldResourceIdLib } from "@latticexyz/world/src/WorldResourceId.sol";
+import { RESOURCE_SYSTEM } from "@latticexyz/world/src/worldResourceTypes.sol";
+import { ROOT_NAMESPACE } from "@latticexyz/world/src/constants.sol";
+import { SystemSwitch } from "../src/utils/SystemSwitch.sol";
+
+contract EchoSystem is System {
+  function msgSender() public view returns (address) {
+    return _msgSender();
+  }
+
+  function world() public view returns (address) {
+    return _world();
+  }
+
+  function echo(string memory message) public view returns (string memory) {
+    return message;
+  }
+
+  function call(ResourceId systemId, bytes memory callData) public returns (bytes memory) {
+    return SystemSwitch.call(systemId, callData);
+  }
+
+  function callViaSelector(bytes memory callData) public returns (bytes memory) {
+    return SystemSwitch.call(callData);
+  }
+}
+
+address constant caller = address(4232);
+
+contract SystemSwitchTest is Test, GasReporter {
+  IBaseWorld world;
+
+  EchoSystem systemA;
+  EchoSystem systemB;
+  EchoSystem rootSystemA;
+  EchoSystem rootSystemB;
+
+  ResourceId systemAId;
+  ResourceId systemBId;
+  ResourceId rootSystemAId;
+  ResourceId rootSystemBId;
+
+  function setUp() public {
+    // Deploy world
+    World _world = new World();
+    _world.initialize(new CoreModule());
+    world = IBaseWorld(address(_world));
+
+    // Deploy systems
+    systemA = new EchoSystem();
+    systemB = new EchoSystem();
+    rootSystemA = new EchoSystem();
+    rootSystemB = new EchoSystem();
+
+    // Encode system IDs
+    systemAId = WorldResourceIdLib.encode({ typeId: RESOURCE_SYSTEM, namespace: "namespaceA", name: "systemA" });
+    systemBId = WorldResourceIdLib.encode({ typeId: RESOURCE_SYSTEM, namespace: "namespaceB", name: "systemB" });
+    rootSystemAId = WorldResourceIdLib.encode({ typeId: RESOURCE_SYSTEM, namespace: ROOT_NAMESPACE, name: "systemA" });
+    rootSystemBId = WorldResourceIdLib.encode({ typeId: RESOURCE_SYSTEM, namespace: ROOT_NAMESPACE, name: "systemB" });
+
+    // Register systems
+    world.registerSystem(systemAId, systemA, true);
+    world.registerSystem(systemBId, systemB, true);
+    world.registerSystem(rootSystemAId, rootSystemA, true);
+    world.registerSystem(rootSystemBId, rootSystemB, true);
+  }
+
+  function _executeFromRootSystemA(ResourceId systemId, bytes memory callData) public returns (bytes memory) {
+    return abi.decode(world.call(rootSystemAId, abi.encodeCall(EchoSystem.call, (systemId, callData))), (bytes));
+  }
+
+  function _executeFromSystemA(ResourceId systemId, bytes memory callData) public returns (bytes memory) {
+    return abi.decode(world.call(systemAId, abi.encodeCall(EchoSystem.call, (systemId, callData))), (bytes));
+  }
+
+  function _executeFromRootSystemA(bytes memory callData) public returns (bytes memory) {
+    return abi.decode(world.call(rootSystemAId, abi.encodeCall(EchoSystem.callViaSelector, (callData))), (bytes));
+  }
+
+  function _executeFromSystemA(bytes memory callData) public returns (bytes memory) {
+    return abi.decode(world.call(systemAId, abi.encodeCall(EchoSystem.callViaSelector, (callData))), (bytes));
+  }
+
+  // - ROOT FROM ROOT ---------------------------------------------------------------------------- //
+
+  function testCallRootFromRootMsgSender() public {
+    vm.prank(caller);
+    bytes memory returnData = _executeFromRootSystemA(rootSystemBId, abi.encodeCall(EchoSystem.msgSender, ()));
+    assertEq(abi.decode(returnData, (address)), caller);
+  }
+
+  function testCallRootFromRootWorld() public {
+    vm.prank(caller);
+    bytes memory returnData = _executeFromRootSystemA(rootSystemBId, abi.encodeCall(EchoSystem.world, ()));
+    assertEq(abi.decode(returnData, (address)), address(world));
+  }
+
+  function testCallRootFromRootEcho() public {
+    vm.prank(caller);
+    bytes memory returnData = _executeFromRootSystemA(rootSystemBId, abi.encodeCall(EchoSystem.echo, ("hello")));
+    assertEq(abi.decode(returnData, (string)), "hello");
+  }
+
+  function testCallRootFromRootWorldSelector() public {
+    bytes4 worldFunctionSelector = world.registerRootFunctionSelector(
+      rootSystemBId,
+      "echo(string)",
+      EchoSystem.echo.selector
+    );
+    bytes memory callData = abi.encodeWithSelector(worldFunctionSelector, "hello");
+
+    vm.prank(caller);
+    bytes memory returnData = _executeFromRootSystemA(callData);
+    assertEq(abi.decode(returnData, (string)), "hello");
+  }
+
+  // - ROOT FROM NON ROOT ---------------------------------------------------------------------------- //
+
+  function testCallRootFromNonRootMsgSender() public {
+    vm.prank(caller);
+    bytes memory returnData = _executeFromSystemA(rootSystemBId, abi.encodeCall(EchoSystem.msgSender, ()));
+    assertEq(abi.decode(returnData, (address)), address(systemA));
+  }
+
+  function testCallRootFromNonRootWorld() public {
+    vm.prank(caller);
+    bytes memory returnData = _executeFromSystemA(rootSystemBId, abi.encodeCall(EchoSystem.world, ()));
+    assertEq(abi.decode(returnData, (address)), address(world));
+  }
+
+  function testCallRootFromNonRootEcho() public {
+    vm.prank(caller);
+    bytes memory returnData = _executeFromSystemA(rootSystemBId, abi.encodeCall(EchoSystem.echo, ("hello")));
+    assertEq(abi.decode(returnData, (string)), "hello");
+  }
+
+  function testCallRootFromNonRootWorldSelector() public {
+    bytes4 worldFunctionSelector = world.registerRootFunctionSelector(
+      rootSystemBId,
+      "echo(string)",
+      EchoSystem.echo.selector
+    );
+    bytes memory callData = abi.encodeWithSelector(worldFunctionSelector, "hello");
+
+    vm.prank(caller);
+    bytes memory returnData = _executeFromSystemA(callData);
+    assertEq(abi.decode(returnData, (string)), "hello");
+  }
+
+  // - NON ROOT FROM ROOT ---------------------------------------------------------------------------- //
+
+  function testCallNonRootFromRootMsgSender() public {
+    vm.prank(caller);
+    bytes memory returnData = _executeFromRootSystemA(systemBId, abi.encodeCall(EchoSystem.msgSender, ()));
+    assertEq(abi.decode(returnData, (address)), caller);
+  }
+
+  function testNonCallRootFromRootWorld() public {
+    vm.prank(caller);
+    bytes memory returnData = _executeFromRootSystemA(systemBId, abi.encodeCall(EchoSystem.world, ()));
+    assertEq(abi.decode(returnData, (address)), address(world));
+  }
+
+  function testNonCallRootFromRootEcho() public {
+    vm.prank(caller);
+    bytes memory returnData = _executeFromRootSystemA(systemBId, abi.encodeCall(EchoSystem.echo, ("hello")));
+    assertEq(abi.decode(returnData, (string)), "hello");
+  }
+
+  function testNonCallRootFromRootWorldSelector() public {
+    bytes4 worldFunctionSelector = world.registerRootFunctionSelector(
+      systemBId,
+      "echo(string)",
+      EchoSystem.echo.selector
+    );
+    bytes memory callData = abi.encodeWithSelector(worldFunctionSelector, "hello");
+
+    vm.prank(caller);
+    bytes memory returnData = _executeFromRootSystemA(callData);
+    assertEq(abi.decode(returnData, (string)), "hello");
+  }
+
+  // - NON ROOT FROM NON ROOT ---------------------------------------------------------------------------- //
+
+  function testCallNonRootFromNonRootMsgSender() public {
+    vm.prank(caller);
+    bytes memory returnData = _executeFromSystemA(systemBId, abi.encodeCall(EchoSystem.msgSender, ()));
+    assertEq(abi.decode(returnData, (address)), address(systemA));
+  }
+
+  function testNonCallRootFromNonRootWorld() public {
+    vm.prank(caller);
+    bytes memory returnData = _executeFromSystemA(systemBId, abi.encodeCall(EchoSystem.world, ()));
+    assertEq(abi.decode(returnData, (address)), address(world));
+  }
+
+  function testNonCallRootFromNonRootEcho() public {
+    vm.prank(caller);
+    bytes memory returnData = _executeFromSystemA(systemBId, abi.encodeCall(EchoSystem.echo, ("hello")));
+    assertEq(abi.decode(returnData, (string)), "hello");
+  }
+
+  function testNonCallRootFromNonRootWorldSelector() public {
+    bytes4 worldFunctionSelector = world.registerRootFunctionSelector(
+      systemBId,
+      "echo(string)",
+      EchoSystem.echo.selector
+    );
+    bytes memory callData = abi.encodeWithSelector(worldFunctionSelector, "hello");
+
+    vm.prank(caller);
+    bytes memory returnData = _executeFromSystemA(callData);
+    assertEq(abi.decode(returnData, (string)), "hello");
+  }
+}