1// SPDX-License-Identifier: MIT
  2pragma solidity >=0.6.2 <0.9.0;
  3
  4import {Vm} from "./Vm.sol";
  5
  6struct StdStorage {
  7    mapping(address => mapping(bytes4 => mapping(bytes32 => uint256))) slots;
  8    mapping(address => mapping(bytes4 => mapping(bytes32 => bool))) finds;
  9    bytes32[] _keys;
 10    bytes4 _sig;
 11    uint256 _depth;
 12    address _target;
 13    bytes32 _set;
 14}
 15
 16library stdStorageSafe {
 17    event SlotFound(address who, bytes4 fsig, bytes32 keysHash, uint256 slot);
 18    event WARNING_UninitedSlot(address who, uint256 slot);
 19
 20    Vm private constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));
 21
 22    function sigs(string memory sigStr) internal pure returns (bytes4) {
 23        return bytes4(keccak256(bytes(sigStr)));
 24    }
 25
 26    /// @notice find an arbitrary storage slot given a function sig, input data, address of the contract and a value to check against
 27    // slot complexity:
 28    //  if flat, will be bytes32(uint256(uint));
 29    //  if map, will be keccak256(abi.encode(key, uint(slot)));
 30    //  if deep map, will be keccak256(abi.encode(key1, keccak256(abi.encode(key0, uint(slot)))));
 31    //  if map struct, will be bytes32(uint256(keccak256(abi.encode(key1, keccak256(abi.encode(key0, uint(slot)))))) + structFieldDepth);
 32    function find(StdStorage storage self) internal returns (uint256) {
 33        address who = self._target;
 34        bytes4 fsig = self._sig;
 35        uint256 field_depth = self._depth;
 36        bytes32[] memory ins = self._keys;
 37
 38        // calldata to test against
 39        if (self.finds[who][fsig][keccak256(abi.encodePacked(ins, field_depth))]) {
 40            return self.slots[who][fsig][keccak256(abi.encodePacked(ins, field_depth))];
 41        }
 42        bytes memory cald = abi.encodePacked(fsig, flatten(ins));
 43        vm.record();
 44        bytes32 fdat;
 45        {
 46            (, bytes memory rdat) = who.staticcall(cald);
 47            fdat = bytesToBytes32(rdat, 32 * field_depth);
 48        }
 49
 50        (bytes32[] memory reads,) = vm.accesses(address(who));
 51        if (reads.length == 1) {
 52            bytes32 curr = vm.load(who, reads[0]);
 53            if (curr == bytes32(0)) {
 54                emit WARNING_UninitedSlot(who, uint256(reads[0]));
 55            }
 56            if (fdat != curr) {
 57                require(
 58                    false,
 59                    "stdStorage find(StdStorage): Packed slot. This would cause dangerous overwriting and currently isn't supported."
 60                );
 61            }
 62            emit SlotFound(who, fsig, keccak256(abi.encodePacked(ins, field_depth)), uint256(reads[0]));
 63            self.slots[who][fsig][keccak256(abi.encodePacked(ins, field_depth))] = uint256(reads[0]);
 64            self.finds[who][fsig][keccak256(abi.encodePacked(ins, field_depth))] = true;
 65        } else if (reads.length > 1) {
 66            for (uint256 i = 0; i < reads.length; i++) {
 67                bytes32 prev = vm.load(who, reads[i]);
 68                if (prev == bytes32(0)) {
 69                    emit WARNING_UninitedSlot(who, uint256(reads[i]));
 70                }
 71                // store
 72                vm.store(who, reads[i], bytes32(hex"1337"));
 73                bool success;
 74                bytes memory rdat;
 75                {
 76                    (success, rdat) = who.staticcall(cald);
 77                    fdat = bytesToBytes32(rdat, 32 * field_depth);
 78                }
 79
 80                if (success && fdat == bytes32(hex"1337")) {
 81                    // we found which of the slots is the actual one
 82                    emit SlotFound(who, fsig, keccak256(abi.encodePacked(ins, field_depth)), uint256(reads[i]));
 83                    self.slots[who][fsig][keccak256(abi.encodePacked(ins, field_depth))] = uint256(reads[i]);
 84                    self.finds[who][fsig][keccak256(abi.encodePacked(ins, field_depth))] = true;
 85                    vm.store(who, reads[i], prev);
 86                    break;
 87                }
 88                vm.store(who, reads[i], prev);
 89            }
 90        } else {
 91            revert("stdStorage find(StdStorage): No storage use detected for target.");
 92        }
 93
 94        require(
 95            self.finds[who][fsig][keccak256(abi.encodePacked(ins, field_depth))],
 96            "stdStorage find(StdStorage): Slot(s) not found."
 97        );
 98
 99        delete self._target;
100        delete self._sig;
101        delete self._keys;
102        delete self._depth;
103
104        return self.slots[who][fsig][keccak256(abi.encodePacked(ins, field_depth))];
105    }
106
107    function target(StdStorage storage self, address _target) internal returns (StdStorage storage) {
108        self._target = _target;
109        return self;
110    }
111
112    function sig(StdStorage storage self, bytes4 _sig) internal returns (StdStorage storage) {
113        self._sig = _sig;
114        return self;
115    }
116
117    function sig(StdStorage storage self, string memory _sig) internal returns (StdStorage storage) {
118        self._sig = sigs(_sig);
119        return self;
120    }
121
122    function with_key(StdStorage storage self, address who) internal returns (StdStorage storage) {
123        self._keys.push(bytes32(uint256(uint160(who))));
124        return self;
125    }
126
127    function with_key(StdStorage storage self, uint256 amt) internal returns (StdStorage storage) {
128        self._keys.push(bytes32(amt));
129        return self;
130    }
131
132    function with_key(StdStorage storage self, bytes32 key) internal returns (StdStorage storage) {
133        self._keys.push(key);
134        return self;
135    }
136
137    function depth(StdStorage storage self, uint256 _depth) internal returns (StdStorage storage) {
138        self._depth = _depth;
139        return self;
140    }
141
142    function read(StdStorage storage self) private returns (bytes memory) {
143        address t = self._target;
144        uint256 s = find(self);
145        return abi.encode(vm.load(t, bytes32(s)));
146    }
147
148    function read_bytes32(StdStorage storage self) internal returns (bytes32) {
149        return abi.decode(read(self), (bytes32));
150    }
151
152    function read_bool(StdStorage storage self) internal returns (bool) {
153        int256 v = read_int(self);
154        if (v == 0) return false;
155        if (v == 1) return true;
156        revert("stdStorage read_bool(StdStorage): Cannot decode. Make sure you are reading a bool.");
157    }
158
159    function read_address(StdStorage storage self) internal returns (address) {
160        return abi.decode(read(self), (address));
161    }
162
163    function read_uint(StdStorage storage self) internal returns (uint256) {
164        return abi.decode(read(self), (uint256));
165    }
166
167    function read_int(StdStorage storage self) internal returns (int256) {
168        return abi.decode(read(self), (int256));
169    }
170
171    function bytesToBytes32(bytes memory b, uint256 offset) private pure returns (bytes32) {
172        bytes32 out;
173
174        uint256 max = b.length > 32 ? 32 : b.length;
175        for (uint256 i = 0; i < max; i++) {
176            out |= bytes32(b[offset + i] & 0xFF) >> (i * 8);
177        }
178        return out;
179    }
180
181    function flatten(bytes32[] memory b) private pure returns (bytes memory) {
182        bytes memory result = new bytes(b.length * 32);
183        for (uint256 i = 0; i < b.length; i++) {
184            bytes32 k = b[i];
185            /// @solidity memory-safe-assembly
186            assembly {
187                mstore(add(result, add(32, mul(32, i))), k)
188            }
189        }
190
191        return result;
192    }
193}
194
195library stdStorage {
196    Vm private constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));
197
198    function sigs(string memory sigStr) internal pure returns (bytes4) {
199        return stdStorageSafe.sigs(sigStr);
200    }
201
202    function find(StdStorage storage self) internal returns (uint256) {
203        return stdStorageSafe.find(self);
204    }
205
206    function target(StdStorage storage self, address _target) internal returns (StdStorage storage) {
207        return stdStorageSafe.target(self, _target);
208    }
209
210    function sig(StdStorage storage self, bytes4 _sig) internal returns (StdStorage storage) {
211        return stdStorageSafe.sig(self, _sig);
212    }
213
214    function sig(StdStorage storage self, string memory _sig) internal returns (StdStorage storage) {
215        return stdStorageSafe.sig(self, _sig);
216    }
217
218    function with_key(StdStorage storage self, address who) internal returns (StdStorage storage) {
219        return stdStorageSafe.with_key(self, who);
220    }
221
222    function with_key(StdStorage storage self, uint256 amt) internal returns (StdStorage storage) {
223        return stdStorageSafe.with_key(self, amt);
224    }
225
226    function with_key(StdStorage storage self, bytes32 key) internal returns (StdStorage storage) {
227        return stdStorageSafe.with_key(self, key);
228    }
229
230    function depth(StdStorage storage self, uint256 _depth) internal returns (StdStorage storage) {
231        return stdStorageSafe.depth(self, _depth);
232    }
233
234    function checked_write(StdStorage storage self, address who) internal {
235        checked_write(self, bytes32(uint256(uint160(who))));
236    }
237
238    function checked_write(StdStorage storage self, uint256 amt) internal {
239        checked_write(self, bytes32(amt));
240    }
241
242    function checked_write(StdStorage storage self, bool write) internal {
243        bytes32 t;
244        /// @solidity memory-safe-assembly
245        assembly {
246            t := write
247        }
248        checked_write(self, t);
249    }
250
251    function checked_write(StdStorage storage self, bytes32 set) internal {
252        address who = self._target;
253        bytes4 fsig = self._sig;
254        uint256 field_depth = self._depth;
255        bytes32[] memory ins = self._keys;
256
257        bytes memory cald = abi.encodePacked(fsig, flatten(ins));
258        if (!self.finds[who][fsig][keccak256(abi.encodePacked(ins, field_depth))]) {
259            find(self);
260        }
261        bytes32 slot = bytes32(self.slots[who][fsig][keccak256(abi.encodePacked(ins, field_depth))]);
262
263        bytes32 fdat;
264        {
265            (, bytes memory rdat) = who.staticcall(cald);
266            fdat = bytesToBytes32(rdat, 32 * field_depth);
267        }
268        bytes32 curr = vm.load(who, slot);
269
270        if (fdat != curr) {
271            require(
272                false,
273                "stdStorage find(StdStorage): Packed slot. This would cause dangerous overwriting and currently isn't supported."
274            );
275        }
276        vm.store(who, slot, set);
277        delete self._target;
278        delete self._sig;
279        delete self._keys;
280        delete self._depth;
281    }
282
283    function read_bytes32(StdStorage storage self) internal returns (bytes32) {
284        return stdStorageSafe.read_bytes32(self);
285    }
286
287    function read_bool(StdStorage storage self) internal returns (bool) {
288        return stdStorageSafe.read_bool(self);
289    }
290
291    function read_address(StdStorage storage self) internal returns (address) {
292        return stdStorageSafe.read_address(self);
293    }
294
295    function read_uint(StdStorage storage self) internal returns (uint256) {
296        return stdStorageSafe.read_uint(self);
297    }
298
299    function read_int(StdStorage storage self) internal returns (int256) {
300        return stdStorageSafe.read_int(self);
301    }
302
303    // Private function so needs to be copied over
304    function bytesToBytes32(bytes memory b, uint256 offset) private pure returns (bytes32) {
305        bytes32 out;
306
307        uint256 max = b.length > 32 ? 32 : b.length;
308        for (uint256 i = 0; i < max; i++) {
309            out |= bytes32(b[offset + i] & 0xFF) >> (i * 8);
310        }
311        return out;
312    }
313
314    // Private function so needs to be copied over
315    function flatten(bytes32[] memory b) private pure returns (bytes memory) {
316        bytes memory result = new bytes(b.length * 32);
317        for (uint256 i = 0; i < b.length; i++) {
318            bytes32 k = b[i];
319            /// @solidity memory-safe-assembly
320            assembly {
321                mstore(add(result, add(32, mul(32, i))), k)
322            }
323        }
324
325        return result;
326    }
327}