fix: add missing storage bases for erc7779

foundry.toml

@@ -1,5 +1,5 @@
 [profile.default]
-evm_version = "cancun"
+evm_version = "prague"
 src = "src"
 out = "out"
 libs = ["node_modules"]

src/MSAAdvanced.sol

@@ -419,6 +419,7 @@ contract MSAAdvanced is
         if (isERC7702) {
             _addStorageBase(MODULEMANAGER_STORAGE_LOCATION);
             _addStorageBase(HOOKMANAGER_STORAGE_LOCATION);
+            _addStorageBase(PREVALIDATION_HOOKMANAGER_STORAGE_LOCATION);
         }
 
         // bootstrap the account
@@ -438,10 +439,30 @@ contract MSAAdvanced is
         if (!success) revert();
     }
 
-    function _onRedelegation() internal override {
+    function _onRedelegation(bytes calldata context) internal override {
+        // Decode the context data which should contain arrays of uninstall data for validators and
+        // executors
+        // (bytes[] memory validatorUninstallData, bytes[] memory executorUninstallData) =
+        //     abi.decode(context, (bytes[], bytes[]));
+
+        // Call the uninstall functions with the decoded data
+        // _tryUninstallValidators(validatorUninstallData);
+        // _tryUninstallExecutors(executorUninstallData);
+
         _tryUninstallValidators();
         _tryUninstallExecutors();
+
+        // Continue with other uninstallations
         _tryUninstallHook(_getHook());
+        _tryUninstallPreValidationHook(
+            _getPreValidationHook(MODULE_TYPE_PREVALIDATION_HOOK_ERC1271),
+            MODULE_TYPE_PREVALIDATION_HOOK_ERC1271
+        );
+        _tryUninstallPreValidationHook(
+            _getPreValidationHook(MODULE_TYPE_PREVALIDATION_HOOK_ERC4337),
+            MODULE_TYPE_PREVALIDATION_HOOK_ERC4337
+        );
+        _tryUninstallFallbacks();
         _initModuleManager();
     }
 }

src/core/ERC7779Adapter.sol

@@ -31,7 +31,7 @@
 
     function accountStorageBases() external view returns (bytes32[] memory) {
         ERC7779Storage storage $;
         assembly {
             $.slot := ERC7779_STORAGE_BASE
         }
         return $.storageBases;
@@ -39,9 +39,15 @@
 
     function _addStorageBase(bytes32 storageBase) internal {
         ERC7779Storage storage $;
         assembly {
             $.slot := ERC7779_STORAGE_BASE
         }
+        // Check if the storageBase already exists to avoid duplication
+        for (uint256 i = 0; i < $.storageBases.length; i++) {
+            if ($.storageBases[i] == storageBase) {
+                return; // Exit if the storageBase is already present
+            }
+        }
         $.storageBases.push(storageBase);
     }
 
@@ -54,14 +60,17 @@
     *         It should uninitialize storages if needed and execute wallet-specific logic to prepare
     for redelegation.
     *         msg.sender should be the owner of the account.
+    * @param context - The context of the redelegation. Additional data required to uninitialize
+    storages.
     */
-    function onRedelegation() external returns (bool) {
+    function onRedelegation(bytes calldata context) external returns (bool) {
         require(msg.sender == address(this), NonAuthorizedOnRedelegationCaller());
-        _onRedelegation();
+        _onRedelegation(context);
         return true;
     }
 
     /// @dev This function is called before redelegation.
     /// @dev Account should override this function to implement the specific logic.
-    function _onRedelegation() internal virtual;
+    /// @param context The context is the additional data required to uninitialize storages.
+    function _onRedelegation(bytes calldata context) internal virtual;
 }

src/core/ModuleManager.sol

@@ -46,6 +46,8 @@ abstract contract ModuleManager is AccountBase, Receiver {
         // single fallback handler for all fallbacks
         // account vendors may implement this differently. This is just a reference implementation
         mapping(bytes4 selector => FallbackHandler fallbackHandler) $fallbacks;
+        ///< List of fallback selectors, initialized upon contract deployment.
+        bytes4[] fallbackSelectors;
     }
 
     function $moduleManager() internal pure virtual returns (ModuleManagerStorage storage $ims) {
@@ -95,36 +97,35 @@ abstract contract ModuleManager is AccountBase, Receiver {
         IValidator(validator).onUninstall(disableModuleData);
     }
 
-    /*
-    function _tryUninstallValidators(bytes[] calldata data) internal {
+    function _tryUninstallValidators() internal virtual {
         SentinelListLib.SentinelList storage $valdiators = $moduleManager().$valdiators;
-        uint256 length = data.length;
-        uint256 index;
         address validator = $valdiators.getNext(SENTINEL);
         while (validator != SENTINEL) {
-            bytes memory uninstallData;
-            if (index < length) {
-                uninstallData = data[index];
-            }
-            try IValidator(validator).onUninstall(uninstallData) {} catch {
-                emit ValidatorUninstallFailed(validator, uninstallData);
+            try IValidator(validator).onUninstall("") { }
+            catch {
+                emit ValidatorUninstallFailed(validator, "");
             }
             validator = $valdiators.getNext(validator);
-            index++;
         }
         $valdiators.popAll();
     }
-    */
 
-    function _tryUninstallValidators() internal {
+    function _tryUninstallValidators(bytes[] memory data) internal virtual {
         SentinelListLib.SentinelList storage $valdiators = $moduleManager().$valdiators;
+        uint256 length = data.length;
+        uint256 index;
         address validator = $valdiators.getNext(SENTINEL);
         while (validator != SENTINEL) {
-            try IValidator(validator).onUninstall("") { }
+            bytes memory uninstallData;
+            if (index < length) {
+                uninstallData = data[index];
+            }
+            try IValidator(validator).onUninstall(uninstallData) { }
             catch {
-                emit ValidatorUninstallFailed(validator, "");
+                emit ValidatorUninstallFailed(validator, uninstallData);
             }
             validator = $valdiators.getNext(validator);
+            index++;
         }
         $valdiators.popAll();
     }
@@ -168,36 +169,35 @@ abstract contract ModuleManager is AccountBase, Receiver {
         IExecutor(executor).onUninstall(disableModuleData);
     }
 
-    /*
-    function _tryUninstallExecutors(bytes[] calldata data) internal {
+    function _tryUninstallExecutors() internal virtual {
         SentinelListLib.SentinelList storage $executors = $moduleManager().$executors;
-        uint256 length = data.length;
-        uint256 index;
         address executor = $executors.getNext(SENTINEL);
         while (executor != SENTINEL) {
-            bytes memory uninstallData;
-            if (index < length) {
-                uninstallData = data[index];
-            }
-            try IExecutor(executor).onUninstall(uninstallData) {} catch {
-                emit ExecutorUninstallFailed(executor, uninstallData);
+            try IExecutor(executor).onUninstall("") { }
+            catch {
+                emit ExecutorUninstallFailed(executor, "");
             }
             executor = $executors.getNext(executor);
-            index++;
         }
         $executors.popAll();
     }
-    */
 
-    function _tryUninstallExecutors() internal {
+    function _tryUninstallExecutors(bytes[] memory data) internal virtual {
         SentinelListLib.SentinelList storage $executors = $moduleManager().$executors;
+        uint256 length = data.length;
+        uint256 index;
         address executor = $executors.getNext(SENTINEL);
         while (executor != SENTINEL) {
-            try IExecutor(executor).onUninstall("") { }
+            bytes memory uninstallData;
+            if (index < length) {
+                uninstallData = data[index];
+            }
+            try IExecutor(executor).onUninstall(uninstallData) { }
             catch {
-                emit ExecutorUninstallFailed(executor, "");
+                emit ExecutorUninstallFailed(executor, uninstallData);
             }
             executor = $executors.getNext(executor);
+            index++;
         }
         $executors.popAll();
     }
@@ -237,6 +237,8 @@ abstract contract ModuleManager is AccountBase, Receiver {
             revert("Function selector already used");
         }
         $moduleManager().$fallbacks[selector] = FallbackHandler(handler, calltype);
+        // Add the selector to the maintained list of fallback selectors
+        $moduleManager().fallbackSelectors.push(selector);
         IFallback(handler).onInstall(initData);
     }
 
@@ -262,9 +264,37 @@ abstract contract ModuleManager is AccountBase, Receiver {
 
         $moduleManager().$fallbacks[selector] = FallbackHandler(address(0), CallType.wrap(0x00));
 
+        // Remove selector from fallbackSelectors via swap-and-pop
+        uint256 len = $moduleManager().fallbackSelectors.length;
+        for (uint256 i = 0; i < len; i++) {
+            if ($moduleManager().fallbackSelectors[i] == selector) {
+                $moduleManager().fallbackSelectors[i] = $moduleManager().fallbackSelectors[len - 1];
+                $moduleManager().fallbackSelectors.pop();
+                break;
+            }
+        }
+
         IFallback(handler).onUninstall(_deInitData);
     }
 
+    // Review: if uninstalling selectors also need some data.
+    function _tryUninstallFallbacks() internal {
+        uint256 len = $moduleManager().fallbackSelectors.length;
+
+        for (uint256 i = 0; i < len; i++) {
+            bytes4 selector = $moduleManager().fallbackSelectors[i];
+            FallbackHandler memory fallbackHandler = $moduleManager().$fallbacks[selector];
+
+            if (address(fallbackHandler.handler) == address(0)) continue;
+
+            IFallback(fallbackHandler.handler).onUninstall(abi.encodePacked(selector));
+
+            $moduleManager().$fallbacks[selector] = FallbackHandler(address(0), CallType.wrap(0x00));
+        }
+
+        delete $moduleManager().fallbackSelectors;
+    }
+
     function _isFallbackHandlerInstalled(bytes4 functionSig) internal view virtual returns (bool) {
         FallbackHandler storage $fallback = $moduleManager().$fallbacks[functionSig];
         return $fallback.handler != address(0);

src/core/PreValidationHookManager.sol

@@ -92,6 +92,7 @@ abstract contract PreValidationHookManager {
     }
 
     function _tryUninstallPreValidationHook(address hook, uint256 hookType) internal virtual {
+        if (hook == address(0)) return;
         PreValidationHookManagerStorage storage $ = _getStorage();
         if (hookType == MODULE_TYPE_PREVALIDATION_HOOK_ERC1271) {
             try $.hook1271.onUninstall("") { }

src/interfaces/IERC7779.sol

@@ -27,6 +27,8 @@ interface IERC7779 {
     *         It should uninitialize storages if needed and execute wallet-specific logic to prepare
     for redelegation.
     *         msg.sender should be the owner of the account.
+    * @param context - The context of the redelegation. Additional data required to uninitialize
+    storages.
     */
-    function onRedelegation() external returns (bool);
+    function onRedelegation(bytes calldata context) external returns (bool);
 }

test/advanced/EIP7702.t.sol

@@ -308,7 +308,7 @@ contract EIP7702 is TestBaseUtilAdvanced {
         assertTrue(IMSA(account).isModuleInstalled(MODULE_TYPE_HOOK, address(hook), ""));
         // storage is cleared
         vm.prank(address(account));
-        IMSA(account).onRedelegation();
+        IMSA(account).onRedelegation(bytes(""));
         assertFalse(
             IMSA(account).isModuleInstalled(MODULE_TYPE_VALIDATOR, address(defaultValidator), "")
         );

test/mocks/MockERC7779.sol

@@ -8,7 +8,7 @@ contract MockERC7779 is ERC7779Adapter {
         _addStorageBase(storageBase);
     }
 
-    function _onRedelegation() internal override {
+    function _onRedelegation(bytes calldata context) internal override {
         // do nothing
     }
 }