feat: Remove hook group (#39)

Author: adam-alchemy

foundry.toml

@@ -7,7 +7,6 @@ libs = ['lib']
 out = 'out'
 optimizer = true
 optimizer_runs = 200
-ignored_error_codes = []
 fs_permissions = [
   { access = "read", path = "./out-optimized" }
 ]

src/account/AccountLoupe.sol

@@ -8,7 +8,7 @@ import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet
 import {IAccountLoupe} from "../interfaces/IAccountLoupe.sol";
 import {FunctionReference, IPluginManager} from "../interfaces/IPluginManager.sol";
 import {IStandardExecutor} from "../interfaces/IStandardExecutor.sol";
-import {AccountStorage, getAccountStorage, HookGroup, toFunctionReferenceArray} from "./AccountStorage.sol";
+import {AccountStorage, getAccountStorage, SelectorData, toFunctionReferenceArray} from "./AccountStorage.sol";
 
 abstract contract AccountLoupe is IAccountLoupe {
     using EnumerableMap for EnumerableMap.Bytes32ToUintMap;
@@ -43,14 +43,14 @@ abstract contract AccountLoupe is IAccountLoupe {
 
     /// @inheritdoc IAccountLoupe
     function getExecutionHooks(bytes4 selector) external view returns (ExecutionHooks[] memory execHooks) {
-        HookGroup storage hooks = getAccountStorage().selectorData[selector].executionHooks;
-        uint256 preExecHooksLength = hooks.preHooks.length();
-        uint256 postOnlyExecHooksLength = hooks.postOnlyHooks.length();
+        SelectorData storage selectorData = getAccountStorage().selectorData[selector];
+        uint256 preExecHooksLength = selectorData.preHooks.length();
+        uint256 postOnlyExecHooksLength = selectorData.postOnlyHooks.length();
         uint256 maxExecHooksLength = postOnlyExecHooksLength;
 
         // There can only be as many associated post hooks to run as there are pre hooks.
         for (uint256 i = 0; i < preExecHooksLength;) {
-            (, uint256 count) = hooks.preHooks.at(i);
+            (, uint256 count) = selectorData.preHooks.at(i);
             unchecked {
                 maxExecHooksLength += (count + 1);
                 ++i;
@@ -62,14 +62,14 @@ abstract contract AccountLoupe is IAccountLoupe {
         uint256 actualExecHooksLength;
 
         for (uint256 i = 0; i < preExecHooksLength;) {
-            (bytes32 key,) = hooks.preHooks.at(i);
+            (bytes32 key,) = selectorData.preHooks.at(i);
             FunctionReference preExecHook = FunctionReference.wrap(bytes21(key));
 
-            uint256 associatedPostExecHooksLength = hooks.associatedPostHooks[preExecHook].length();
+            uint256 associatedPostExecHooksLength = selectorData.associatedPostHooks[preExecHook].length();
             if (associatedPostExecHooksLength > 0) {
                 for (uint256 j = 0; j < associatedPostExecHooksLength;) {
                     execHooks[actualExecHooksLength].preExecHook = preExecHook;
-                    (key,) = hooks.associatedPostHooks[preExecHook].at(j);
+                    (key,) = selectorData.associatedPostHooks[preExecHook].at(j);
                     execHooks[actualExecHooksLength].postExecHook = FunctionReference.wrap(bytes21(key));
 
                     unchecked {
@@ -91,7 +91,7 @@ abstract contract AccountLoupe is IAccountLoupe {
         }
 
         for (uint256 i = 0; i < postOnlyExecHooksLength;) {
-            (bytes32 key,) = hooks.postOnlyHooks.at(i);
+            (bytes32 key,) = selectorData.postOnlyHooks.at(i);
             execHooks[actualExecHooksLength].postExecHook = FunctionReference.wrap(bytes21(key));
 
             unchecked {

src/account/AccountStorage.sol

@@ -12,8 +12,7 @@ bytes32 constant _ACCOUNT_STORAGE_SLOT = 0x9f09680beaa4e5c9f38841db2460c40149916
 
 struct PluginData {
     bool anyExternalExecPermitted;
-    // boolean to indicate if the plugin can spend native tokens, if any of the execution function can spend
-    // native tokens, a plugin is considered to be able to spend native tokens of the accounts
+    // boolean to indicate if the plugin can spend native tokens from the account.
     bool canSpendNativeToken;
     bytes32 manifestHash;
     FunctionReference[] dependencies;
@@ -31,14 +30,6 @@ struct PermittedExternalCallData {
     mapping(bytes4 => bool) permittedSelectors;
 }
 
-// Represets a set of pre- and post- hooks.
-struct HookGroup {
-    EnumerableMap.Bytes32ToUintMap preHooks;
-    // bytes21 key = pre hook function reference
-    mapping(FunctionReference => EnumerableMap.Bytes32ToUintMap) associatedPostHooks;
-    EnumerableMap.Bytes32ToUintMap postOnlyHooks;
-}
-
 // Represents data associated with a specifc function selector.
 struct SelectorData {
     // The plugin that implements this execution function.
@@ -50,7 +41,10 @@ struct SelectorData {
     EnumerableMap.Bytes32ToUintMap preUserOpValidationHooks;
     EnumerableMap.Bytes32ToUintMap preRuntimeValidationHooks;
     // The execution hooks for this function selector.
-    HookGroup executionHooks;
+    EnumerableMap.Bytes32ToUintMap preHooks;
+    // bytes21 key = pre hook function reference
+    mapping(FunctionReference => EnumerableMap.Bytes32ToUintMap) associatedPostHooks;
+    EnumerableMap.Bytes32ToUintMap postOnlyHooks;
 }
 
 struct AccountStorage {

src/account/PluginManagerInternals.sol

@@ -21,7 +21,6 @@ import {
     getAccountStorage,
     SelectorData,
     getPermittedCallKey,
-    HookGroup,
     PermittedExternalCallData
 } from "./AccountStorage.sol";
 
@@ -126,50 +125,38 @@ abstract contract PluginManagerInternals is IPluginManager {
     {
         SelectorData storage _selectorData = getAccountStorage().selectorData[selector];
 
-        _addHooks(_selectorData.executionHooks, preExecHook, postExecHook);
-    }
-
-    function _removeExecHooks(bytes4 selector, FunctionReference preExecHook, FunctionReference postExecHook)
-        internal
-    {
-        SelectorData storage _selectorData = getAccountStorage().selectorData[selector];
-
-        _removeHooks(_selectorData.executionHooks, preExecHook, postExecHook);
-    }
-
-    function _addHooks(HookGroup storage hooks, FunctionReference preExecHook, FunctionReference postExecHook)
-        internal
-    {
         if (!preExecHook.isEmpty()) {
-            _addOrIncrement(hooks.preHooks, _toSetValue(preExecHook));
+            _addOrIncrement(_selectorData.preHooks, _toSetValue(preExecHook));
 
             if (!postExecHook.isEmpty()) {
-                _addOrIncrement(hooks.associatedPostHooks[preExecHook], _toSetValue(postExecHook));
+                _addOrIncrement(_selectorData.associatedPostHooks[preExecHook], _toSetValue(postExecHook));
             }
         } else {
             if (postExecHook.isEmpty()) {
                 // both pre and post hooks cannot be null
                 revert NullFunctionReference();
             }
 
-            _addOrIncrement(hooks.postOnlyHooks, _toSetValue(postExecHook));
+            _addOrIncrement(_selectorData.postOnlyHooks, _toSetValue(postExecHook));
         }
     }
 
-    function _removeHooks(HookGroup storage hooks, FunctionReference preExecHook, FunctionReference postExecHook)
+    function _removeExecHooks(bytes4 selector, FunctionReference preExecHook, FunctionReference postExecHook)
         internal
     {
+        SelectorData storage _selectorData = getAccountStorage().selectorData[selector];
+
         if (!preExecHook.isEmpty()) {
-            _removeOrDecrement(hooks.preHooks, _toSetValue(preExecHook));
+            _removeOrDecrement(_selectorData.preHooks, _toSetValue(preExecHook));
 
             if (!postExecHook.isEmpty()) {
-                _removeOrDecrement(hooks.associatedPostHooks[preExecHook], _toSetValue(postExecHook));
+                _removeOrDecrement(_selectorData.associatedPostHooks[preExecHook], _toSetValue(postExecHook));
             }
         } else {
             // The case where both pre and post hooks are null was checked during installation.
 
             // May ignore return value, as the manifest hash is validated to ensure that the hook exists.
-            _removeOrDecrement(hooks.postOnlyHooks, _toSetValue(postExecHook));
+            _removeOrDecrement(_selectorData.postOnlyHooks, _toSetValue(postExecHook));
         }
     }
 

src/account/UpgradeableModularAccount.sol

@@ -17,7 +17,7 @@ import {FunctionReference, IPluginManager} from "../interfaces/IPluginManager.so
 import {IStandardExecutor, Call} from "../interfaces/IStandardExecutor.sol";
 import {AccountExecutor} from "./AccountExecutor.sol";
 import {AccountLoupe} from "./AccountLoupe.sol";
-import {AccountStorage, HookGroup, getAccountStorage, getPermittedCallKey} from "./AccountStorage.sol";
+import {AccountStorage, getAccountStorage, getPermittedCallKey, SelectorData} from "./AccountStorage.sol";
 import {AccountStorageInitializable} from "./AccountStorageInitializable.sol";
 import {PluginManagerInternals} from "./PluginManagerInternals.sol";
 
@@ -459,14 +459,14 @@ contract UpgradeableModularAccount is
         internal
         returns (PostExecToRun[] memory postHooksToRun)
     {
-        HookGroup storage hooks = getAccountStorage().selectorData[selector].executionHooks;
-        uint256 preExecHooksLength = hooks.preHooks.length();
-        uint256 postOnlyHooksLength = hooks.postOnlyHooks.length();
+        SelectorData storage selectorData = getAccountStorage().selectorData[selector];
+        uint256 preExecHooksLength = selectorData.preHooks.length();
+        uint256 postOnlyHooksLength = selectorData.postOnlyHooks.length();
         uint256 maxPostExecHooksLength = postOnlyHooksLength;
 
         // There can only be as many associated post hooks to run as there are pre hooks.
         for (uint256 i = 0; i < preExecHooksLength;) {
-            (, uint256 count) = hooks.preHooks.at(i);
+            (, uint256 count) = selectorData.preHooks.at(i);
             unchecked {
                 maxPostExecHooksLength += (count + 1);
                 ++i;
@@ -479,7 +479,7 @@ contract UpgradeableModularAccount is
 
         // Copy post-only hooks to the array.
         for (uint256 i = 0; i < postOnlyHooksLength;) {
-            (bytes32 key,) = hooks.postOnlyHooks.at(i);
+            (bytes32 key,) = selectorData.postOnlyHooks.at(i);
             postHooksToRun[actualPostHooksToRunLength].postExecHook = _toFunctionReference(key);
             unchecked {
                 ++actualPostHooksToRunLength;
@@ -490,7 +490,7 @@ contract UpgradeableModularAccount is
         // Then run the pre hooks and copy the associated post hooks (along with their pre hook's return data) to
         // the array.
         for (uint256 i = 0; i < preExecHooksLength;) {
-            (bytes32 key,) = hooks.preHooks.at(i);
+            (bytes32 key,) = selectorData.preHooks.at(i);
             FunctionReference preExecHook = _toFunctionReference(key);
 
             if (preExecHook.isEmptyOrMagicValue()) {
@@ -501,10 +501,10 @@ contract UpgradeableModularAccount is
 
             bytes memory preExecHookReturnData = _runPreExecHook(preExecHook, data);
 
-            uint256 associatedPostExecHooksLength = hooks.associatedPostHooks[preExecHook].length();
+            uint256 associatedPostExecHooksLength = selectorData.associatedPostHooks[preExecHook].length();
             if (associatedPostExecHooksLength > 0) {
                 for (uint256 j = 0; j < associatedPostExecHooksLength;) {
-                    (key,) = hooks.associatedPostHooks[preExecHook].at(j);
+                    (key,) = selectorData.associatedPostHooks[preExecHook].at(j);
                     postHooksToRun[actualPostHooksToRunLength].postExecHook = _toFunctionReference(key);
                     postHooksToRun[actualPostHooksToRunLength].preExecHookReturnData = preExecHookReturnData;