From 1168e6d62dd44be2103ed199d718401acf62c4a8 Mon Sep 17 00:00:00 2001
From: Michael de Hoog <michael.dehoog@coinbase.com>
Date: Sat, 19 Oct 2024 22:01:15 -1000
Subject: [PATCH] Pass IGnosisSafe to Simulation library methods

---
 script/universal/MultisigBase.sol |  2 +-
 script/universal/Simulation.sol   | 22 +++++++++-------------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/script/universal/MultisigBase.sol b/script/universal/MultisigBase.sol
index 5b230e4..cc46694 100644
--- a/script/universal/MultisigBase.sol
+++ b/script/universal/MultisigBase.sol
@@ -195,7 +195,7 @@ abstract contract MultisigBase is CommonBase {
     // not be reflected in the prod execution.
     function _safeOverrides(IGnosisSafe _safe, address _owner) internal virtual view returns (Simulation.StateOverride memory) {
         uint256 _nonce = _getNonce(_safe);
-        return Simulation.overrideSafeThresholdOwnerAndNonce(address(_safe), _owner, _nonce);
+        return Simulation.overrideSafeThresholdOwnerAndNonce(_safe, _owner, _nonce);
     }
 
     // Tenderly simulations can accept generic state overrides. This hook enables this functionality.
diff --git a/script/universal/Simulation.sol b/script/universal/Simulation.sol
index 10ff69b..77445b3 100644
--- a/script/universal/Simulation.sol
+++ b/script/universal/Simulation.sol
@@ -3,6 +3,7 @@ pragma solidity ^0.8.15;
 
 import { console } from "forge-std/console.sol";
 import { Vm } from "forge-std/Vm.sol";
+import { IGnosisSafe } from "./IGnosisSafe.sol";
 
 library Simulation {
     address internal constant VM_ADDRESS = address(uint160(uint256(keccak256("hevm cheat code"))));
@@ -50,9 +51,9 @@ library Simulation {
         return accesses;
     }
 
-    function overrideSafeThresholdOwnerAndNonce(address _safe, address _owner, uint256 _nonce) public view returns (StateOverride memory) {
+    function overrideSafeThresholdOwnerAndNonce(IGnosisSafe _safe, address _owner, uint256 _nonce) public view returns (StateOverride memory) {
         StateOverride memory state = StateOverride({
-            contractAddress: _safe,
+            contractAddress: address(_safe),
             overrides: new StorageOverride[](0)
         });
         state = addThresholdOverride(_safe, state);
@@ -61,11 +62,9 @@ library Simulation {
         return state;
     }
 
-    function addThresholdOverride(address _safe, StateOverride memory _state) internal view returns (StateOverride memory) {
+    function addThresholdOverride(IGnosisSafe _safe, StateOverride memory _state) internal view returns (StateOverride memory) {
         // get the threshold and check if we need to override it
-        (, bytes memory thresholdBytes) = _safe.staticcall(abi.encodeWithSignature("getThreshold()"));
-        uint256 threshold = abi.decode(thresholdBytes, (uint256));
-        if (threshold == 1) return _state;
+        if (_safe.getThreshold() == 1) return _state;
 
         // set the threshold (slot 4) to 1
         return addOverride(_state, StorageOverride({
@@ -74,10 +73,9 @@ library Simulation {
         }));
     }
 
-    function addOwnerOverride(address _safe, StateOverride memory _state, address _owner) internal view returns (StateOverride memory) {
+    function addOwnerOverride(IGnosisSafe _safe, StateOverride memory _state, address _owner) internal view returns (StateOverride memory) {
         // get the owners and check if _owner is an owner
-        (, bytes memory ownersBytes) = _safe.staticcall(abi.encodeWithSignature("getOwners()"));
-        address[] memory owners = abi.decode(ownersBytes, (address[]));
+        address[] memory owners = _safe.getOwners();
         for (uint256 i; i < owners.length; i++) {
             if (owners[i] == _owner) return _state;
         }
@@ -98,11 +96,9 @@ library Simulation {
         }));
     }
 
-    function addNonceOverride(address _safe, StateOverride memory _state, uint256 _nonce) internal view returns (StateOverride memory) {
+    function addNonceOverride(IGnosisSafe _safe, StateOverride memory _state, uint256 _nonce) internal view returns (StateOverride memory) {
         // get the nonce and check if we need to override it
-        (, bytes memory nonceBytes) = _safe.staticcall(abi.encodeWithSignature("nonce()"));
-        uint256 nonce = abi.decode(nonceBytes, (uint256));
-        if (nonce == _nonce) return _state;
+        if (_safe.nonce() == _nonce) return _state;
 
         // set the nonce (slot 5) to the desired value
         return addOverride(_state, StorageOverride({