From 1869bd42fd4c4da8c209cb399e22af932c8955ea Mon Sep 17 00:00:00 2001
From: Jay Paik <jay@alchemy.com>
Date: Sun, 21 Jan 2024 23:10:54 -0500
Subject: [PATCH] feat: run duplicate pre-hooks just once

---
 src/account/AccountLoupe.sol              | 45 +++++++++++------------
 src/account/UpgradeableModularAccount.sol | 15 ++------
 2 files changed, 24 insertions(+), 36 deletions(-)

diff --git a/src/account/AccountLoupe.sol b/src/account/AccountLoupe.sol
index 55b1566c..fa24b521 100644
--- a/src/account/AccountLoupe.sol
+++ b/src/account/AccountLoupe.sol
@@ -44,30 +44,7 @@ abstract contract AccountLoupe is IAccountLoupe {
 
     /// @inheritdoc IAccountLoupe
     function getExecutionHooks(bytes4 selector) external view returns (ExecutionHooks[] memory execHooks) {
-        execHooks = _getHooks(getAccountStorage().selectorData[selector].executionHooks);
-    }
-
-    /// @inheritdoc IAccountLoupe
-    function getPreValidationHooks(bytes4 selector)
-        external
-        view
-        returns (
-            FunctionReference[] memory preUserOpValidationHooks,
-            FunctionReference[] memory preRuntimeValidationHooks
-        )
-    {
-        preUserOpValidationHooks =
-            toFunctionReferenceArray(getAccountStorage().selectorData[selector].preUserOpValidationHooks);
-        preRuntimeValidationHooks =
-            toFunctionReferenceArray(getAccountStorage().selectorData[selector].preRuntimeValidationHooks);
-    }
-
-    /// @inheritdoc IAccountLoupe
-    function getInstalledPlugins() external view returns (address[] memory pluginAddresses) {
-        pluginAddresses = getAccountStorage().plugins.values();
-    }
-
-    function _getHooks(HookGroup storage hooks) internal view returns (ExecutionHooks[] memory execHooks) {
+        HookGroup storage hooks = getAccountStorage().selectorData[selector].executionHooks;
         uint256 preExecHooksLength = hooks.preHooks.length();
         uint256 postOnlyExecHooksLength = hooks.postOnlyHooks.length();
         uint256 maxExecHooksLength = postOnlyExecHooksLength;
@@ -129,4 +106,24 @@ abstract contract AccountLoupe is IAccountLoupe {
             mstore(execHooks, actualExecHooksLength)
         }
     }
+
+    /// @inheritdoc IAccountLoupe
+    function getPreValidationHooks(bytes4 selector)
+        external
+        view
+        returns (
+            FunctionReference[] memory preUserOpValidationHooks,
+            FunctionReference[] memory preRuntimeValidationHooks
+        )
+    {
+        preUserOpValidationHooks =
+            toFunctionReferenceArray(getAccountStorage().selectorData[selector].preUserOpValidationHooks);
+        preRuntimeValidationHooks =
+            toFunctionReferenceArray(getAccountStorage().selectorData[selector].preRuntimeValidationHooks);
+    }
+
+    /// @inheritdoc IAccountLoupe
+    function getInstalledPlugins() external view returns (address[] memory pluginAddresses) {
+        pluginAddresses = getAccountStorage().plugins.values();
+    }
 }
diff --git a/src/account/UpgradeableModularAccount.sol b/src/account/UpgradeableModularAccount.sol
index 950caeb5..875d3d29 100644
--- a/src/account/UpgradeableModularAccount.sol
+++ b/src/account/UpgradeableModularAccount.sol
@@ -470,14 +470,6 @@ contract UpgradeableModularAccount is
         returns (PostExecToRun[] memory postHooksToRun)
     {
         HookGroup storage hooks = getAccountStorage().selectorData[selector].executionHooks;
-
-        return _doPreHooks(hooks, data);
-    }
-
-    function _doPreHooks(HookGroup storage hooks, bytes calldata data)
-        internal
-        returns (PostExecToRun[] memory postHooksToRun)
-    {
         uint256 preExecHooksLength = hooks.preHooks.length();
         uint256 postOnlyHooksLength = hooks.postOnlyHooks.length();
         uint256 maxPostExecHooksLength = postOnlyHooksLength;
@@ -517,21 +509,20 @@ contract UpgradeableModularAccount is
                 revert AlwaysDenyRule();
             }
 
+            bytes memory preExecHookReturnData = _runPreExecHook(preExecHook, data);
+
             uint256 associatedPostExecHooksLength = hooks.associatedPostHooks[preExecHook].length();
             if (associatedPostExecHooksLength > 0) {
                 for (uint256 j = 0; j < associatedPostExecHooksLength;) {
                     (key,) = hooks.associatedPostHooks[preExecHook].at(j);
                     postHooksToRun[actualPostHooksToRunLength].postExecHook = _toFunctionReference(key);
-                    postHooksToRun[actualPostHooksToRunLength].preExecHookReturnData =
-                        _runPreExecHook(preExecHook, data);
+                    postHooksToRun[actualPostHooksToRunLength].preExecHookReturnData = preExecHookReturnData;
 
                     unchecked {
                         ++actualPostHooksToRunLength;
                         ++j;
                     }
                 }
-            } else {
-                _runPreExecHook(preExecHook, data);
             }
 
             unchecked {