Plugin refactor (#302)

* refactor plugin extensions

* tests

* forge updates

* Update comments for IMap

* Fix hardhat config

* Update Plugin struct variable names

* docs update

* fix tests

* rename Map -> PluginMap

* revert addPlugin if default plugin exists in PluginMap

* include PluginMap in scope of updatePlugin

* Reorder Plugin struct vars

* rename test getAllRegistered -> getAllPlugins

* cleanup Router

* update Router tests

Co-authored-by: Krishang <[email protected]>

Author: kumaryash90

contracts/extension/interface/plugin/IMap.sol

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity ^0.8.11;
+
+interface IPluginMap {
+    /**
+     *  @notice An interface to describe a plug-in.
+     *
+     *  @param functionSelector     4-byte function selector.
+     *  @param functionSignature    Function representation as a string. E.g. "transfer(address,address,uint256)"
+     *  @param pluginAddress        Address of the contract containing the function.
+     */
+    struct Plugin {
+        bytes4 functionSelector;
+        string functionSignature;
+        address pluginAddress;
+    }
+
+    /// @dev Emitted when a function selector is mapped to a particular plug-in smart contract, during construction of Map.
+    event PluginSet(bytes4 indexed functionSelector, string indexed functionSignature, address indexed pluginAddress);
+
+    /// @dev Returns the plug-in contract for a given function.
+    function getPluginForFunction(bytes4 functionSelector) external view returns (address);
+
+    /// @dev Returns all functions that are mapped to the given plug-in contract.
+    function getAllFunctionsOfPlugin(address pluginAddress) external view returns (bytes4[] memory);
+
+    /// @dev Returns all plug-ins known by Map.
+    function getAllPlugins() external view returns (Plugin[] memory);
+}

contracts/extension/interface/plugin/IPluginMap.sol

@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity ^0.8.11;
+
+import "./IPluginMap.sol";
+
+interface IRouter is IPluginMap {
+    /// @dev Emitted when a functionality is added, or plugged-in.
+    event PluginAdded(bytes4 indexed functionSelector, address indexed pluginAddress);
+
+    /// @dev Emitted when a functionality is updated or overridden.
+    event PluginUpdated(
+        bytes4 indexed functionSelector,
+        address indexed oldPluginAddress,
+        address indexed newPluginAddress
+    );
+
+    /// @dev Emitted when a functionality is removed.
+    event PluginRemoved(bytes4 indexed functionSelector, address indexed pluginAddress);
+
+    /// @dev Add a new plugin to the contract.
+    function addPlugin(Plugin memory plugin) external;
+
+    /// @dev Update / override an existing plugin.
+    function updatePlugin(Plugin memory plugin) external;
+
+    /// @dev Remove an existing plugin from the contract.
+    function removePlugin(bytes4 functionSelector) external;
+}

contracts/extension/interface/plugin/IRouter.sol

@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity ^0.8.0;
+
+import "../interface/plugin/IPluginMap.sol";
+import "../../openzeppelin-presets/utils/EnumerableSet.sol";
+
+contract PluginMap is IPluginMap {
+    using EnumerableSet for EnumerableSet.Bytes32Set;
+
+    EnumerableSet.Bytes32Set private allSelectors;
+
+    mapping(address => EnumerableSet.Bytes32Set) private selectorsForPlugin;
+    mapping(bytes4 => Plugin) private pluginForSelector;
+
+    /*///////////////////////////////////////////////////////////////
+                    Constructor + initializer logic
+    //////////////////////////////////////////////////////////////*/
+
+    constructor(Plugin[] memory _pluginsToAdd) {
+        uint256 len = _pluginsToAdd.length;
+        for (uint256 i = 0; i < len; i += 1) {
+            _setPlugin(_pluginsToAdd[i]);
+        }
+    }
+
+    /*///////////////////////////////////////////////////////////////
+                            View functions
+    //////////////////////////////////////////////////////////////*/
+
+    /// @dev View address of the plugged-in functionality contract for a given function signature.
+    function getPluginForFunction(bytes4 _selector) public view returns (address) {
+        address _pluginAddress = pluginForSelector[_selector].pluginAddress;
+        require(_pluginAddress != address(0), "Map: No plugin available for selector");
+
+        return _pluginAddress;
+    }
+
+    /// @dev View all funtionality as list of function signatures.
+    function getAllFunctionsOfPlugin(address _pluginAddress) external view returns (bytes4[] memory registered) {
+        uint256 len = selectorsForPlugin[_pluginAddress].length();
+        registered = new bytes4[](len);
+
+        for (uint256 i = 0; i < len; i += 1) {
+            registered[i] = bytes4(selectorsForPlugin[_pluginAddress].at(i));
+        }
+    }
+
+    /// @dev View all funtionality existing on the contract.
+    function getAllPlugins() external view returns (Plugin[] memory _plugins) {
+        uint256 len = allSelectors.length();
+        _plugins = new Plugin[](len);
+
+        for (uint256 i = 0; i < len; i += 1) {
+            bytes4 selector = bytes4(allSelectors.at(i));
+            _plugins[i] = pluginForSelector[selector];
+        }
+    }
+
+    /*///////////////////////////////////////////////////////////////
+                        Internal functions
+    //////////////////////////////////////////////////////////////*/
+
+    /// @dev Add functionality to the contract.
+    function _setPlugin(Plugin memory _plugin) internal {
+        require(allSelectors.add(bytes32(_plugin.functionSelector)), "Map: Selector exists");
+        require(
+            _plugin.functionSelector == bytes4(keccak256(abi.encodePacked(_plugin.functionSignature))),
+            "Map: Incorrect selector"
+        );
+
+        pluginForSelector[_plugin.functionSelector] = _plugin;
+        selectorsForPlugin[_plugin.pluginAddress].add(bytes32(_plugin.functionSelector));
+
+        emit PluginSet(_plugin.functionSelector, _plugin.functionSignature, _plugin.pluginAddress);
+    }
+}