1// SPDX-License-Identifier: MIT
  2pragma solidity >=0.7.0 <0.9.0;
  3
  4import "../src/StdStorage.sol";
  5import "../src/Test.sol";
  6
  7contract StdStorageTest is Test {
  8    using stdStorage for StdStorage;
  9
 10    StorageTest internal test;
 11
 12    function setUp() public {
 13        test = new StorageTest();
 14    }
 15
 16    function testStorageHidden() public {
 17        assertEq(uint256(keccak256("my.random.var")), stdstore.target(address(test)).sig("hidden()").find());
 18    }
 19
 20    function testStorageObvious() public {
 21        assertEq(uint256(0), stdstore.target(address(test)).sig("exists()").find());
 22    }
 23
 24    function testStorageCheckedWriteHidden() public {
 25        stdstore.target(address(test)).sig(test.hidden.selector).checked_write(100);
 26        assertEq(uint256(test.hidden()), 100);
 27    }
 28
 29    function testStorageCheckedWriteObvious() public {
 30        stdstore.target(address(test)).sig(test.exists.selector).checked_write(100);
 31        assertEq(test.exists(), 100);
 32    }
 33
 34    function testStorageMapStructA() public {
 35        uint256 slot =
 36            stdstore.target(address(test)).sig(test.map_struct.selector).with_key(address(this)).depth(0).find();
 37        assertEq(uint256(keccak256(abi.encode(address(this), 4))), slot);
 38    }
 39
 40    function testStorageMapStructB() public {
 41        uint256 slot =
 42            stdstore.target(address(test)).sig(test.map_struct.selector).with_key(address(this)).depth(1).find();
 43        assertEq(uint256(keccak256(abi.encode(address(this), 4))) + 1, slot);
 44    }
 45
 46    function testStorageDeepMap() public {
 47        uint256 slot = stdstore.target(address(test)).sig(test.deep_map.selector).with_key(address(this)).with_key(
 48            address(this)
 49        ).find();
 50        assertEq(uint256(keccak256(abi.encode(address(this), keccak256(abi.encode(address(this), uint256(5)))))), slot);
 51    }
 52
 53    function testStorageCheckedWriteDeepMap() public {
 54        stdstore.target(address(test)).sig(test.deep_map.selector).with_key(address(this)).with_key(address(this))
 55            .checked_write(100);
 56        assertEq(100, test.deep_map(address(this), address(this)));
 57    }
 58
 59    function testStorageDeepMapStructA() public {
 60        uint256 slot = stdstore.target(address(test)).sig(test.deep_map_struct.selector).with_key(address(this))
 61            .with_key(address(this)).depth(0).find();
 62        assertEq(
 63            bytes32(uint256(keccak256(abi.encode(address(this), keccak256(abi.encode(address(this), uint256(6)))))) + 0),
 64            bytes32(slot)
 65        );
 66    }
 67
 68    function testStorageDeepMapStructB() public {
 69        uint256 slot = stdstore.target(address(test)).sig(test.deep_map_struct.selector).with_key(address(this))
 70            .with_key(address(this)).depth(1).find();
 71        assertEq(
 72            bytes32(uint256(keccak256(abi.encode(address(this), keccak256(abi.encode(address(this), uint256(6)))))) + 1),
 73            bytes32(slot)
 74        );
 75    }
 76
 77    function testStorageCheckedWriteDeepMapStructA() public {
 78        stdstore.target(address(test)).sig(test.deep_map_struct.selector).with_key(address(this)).with_key(
 79            address(this)
 80        ).depth(0).checked_write(100);
 81        (uint256 a, uint256 b) = test.deep_map_struct(address(this), address(this));
 82        assertEq(100, a);
 83        assertEq(0, b);
 84    }
 85
 86    function testStorageCheckedWriteDeepMapStructB() public {
 87        stdstore.target(address(test)).sig(test.deep_map_struct.selector).with_key(address(this)).with_key(
 88            address(this)
 89        ).depth(1).checked_write(100);
 90        (uint256 a, uint256 b) = test.deep_map_struct(address(this), address(this));
 91        assertEq(0, a);
 92        assertEq(100, b);
 93    }
 94
 95    function testStorageCheckedWriteMapStructA() public {
 96        stdstore.target(address(test)).sig(test.map_struct.selector).with_key(address(this)).depth(0).checked_write(100);
 97        (uint256 a, uint256 b) = test.map_struct(address(this));
 98        assertEq(a, 100);
 99        assertEq(b, 0);
100    }
101
102    function testStorageCheckedWriteMapStructB() public {
103        stdstore.target(address(test)).sig(test.map_struct.selector).with_key(address(this)).depth(1).checked_write(100);
104        (uint256 a, uint256 b) = test.map_struct(address(this));
105        assertEq(a, 0);
106        assertEq(b, 100);
107    }
108
109    function testStorageStructA() public {
110        uint256 slot = stdstore.target(address(test)).sig(test.basic.selector).depth(0).find();
111        assertEq(uint256(7), slot);
112    }
113
114    function testStorageStructB() public {
115        uint256 slot = stdstore.target(address(test)).sig(test.basic.selector).depth(1).find();
116        assertEq(uint256(7) + 1, slot);
117    }
118
119    function testStorageCheckedWriteStructA() public {
120        stdstore.target(address(test)).sig(test.basic.selector).depth(0).checked_write(100);
121        (uint256 a, uint256 b) = test.basic();
122        assertEq(a, 100);
123        assertEq(b, 1337);
124    }
125
126    function testStorageCheckedWriteStructB() public {
127        stdstore.target(address(test)).sig(test.basic.selector).depth(1).checked_write(100);
128        (uint256 a, uint256 b) = test.basic();
129        assertEq(a, 1337);
130        assertEq(b, 100);
131    }
132
133    function testStorageMapAddrFound() public {
134        uint256 slot = stdstore.target(address(test)).sig(test.map_addr.selector).with_key(address(this)).find();
135        assertEq(uint256(keccak256(abi.encode(address(this), uint256(1)))), slot);
136    }
137
138    function testStorageMapUintFound() public {
139        uint256 slot = stdstore.target(address(test)).sig(test.map_uint.selector).with_key(100).find();
140        assertEq(uint256(keccak256(abi.encode(100, uint256(2)))), slot);
141    }
142
143    function testStorageCheckedWriteMapUint() public {
144        stdstore.target(address(test)).sig(test.map_uint.selector).with_key(100).checked_write(100);
145        assertEq(100, test.map_uint(100));
146    }
147
148    function testStorageCheckedWriteMapAddr() public {
149        stdstore.target(address(test)).sig(test.map_addr.selector).with_key(address(this)).checked_write(100);
150        assertEq(100, test.map_addr(address(this)));
151    }
152
153    function testStorageCheckedWriteMapBool() public {
154        stdstore.target(address(test)).sig(test.map_bool.selector).with_key(address(this)).checked_write(true);
155        assertTrue(test.map_bool(address(this)));
156    }
157
158    function testFailStorageCheckedWriteMapPacked() public {
159        // expect PackedSlot error but not external call so cant expectRevert
160        stdstore.target(address(test)).sig(test.read_struct_lower.selector).with_key(address(uint160(1337)))
161            .checked_write(100);
162    }
163
164    function testStorageCheckedWriteMapPackedSuccess() public {
165        uint256 full = test.map_packed(address(1337));
166        // keep upper 128, set lower 128 to 1337
167        full = (full & (uint256((1 << 128) - 1) << 128)) | 1337;
168        stdstore.target(address(test)).sig(test.map_packed.selector).with_key(address(uint160(1337))).checked_write(
169            full
170        );
171        assertEq(1337, test.read_struct_lower(address(1337)));
172    }
173
174    function testFailStorageConst() public {
175        // vm.expectRevert(abi.encodeWithSignature("NotStorage(bytes4)", bytes4(keccak256("const()"))));
176        stdstore.target(address(test)).sig("const()").find();
177    }
178
179    function testFailStorageNativePack() public {
180        stdstore.target(address(test)).sig(test.tA.selector).find();
181        stdstore.target(address(test)).sig(test.tB.selector).find();
182
183        // these both would fail
184        stdstore.target(address(test)).sig(test.tC.selector).find();
185        stdstore.target(address(test)).sig(test.tD.selector).find();
186    }
187
188    function testStorageReadBytes32() public {
189        bytes32 val = stdstore.target(address(test)).sig(test.tE.selector).read_bytes32();
190        assertEq(val, hex"1337");
191    }
192
193    function testStorageReadBool_False() public {
194        bool val = stdstore.target(address(test)).sig(test.tB.selector).read_bool();
195        assertEq(val, false);
196    }
197
198    function testStorageReadBool_True() public {
199        bool val = stdstore.target(address(test)).sig(test.tH.selector).read_bool();
200        assertEq(val, true);
201    }
202
203    function testStorageReadBool_Revert() public {
204        vm.expectRevert("stdStorage read_bool(StdStorage): Cannot decode. Make sure you are reading a bool.");
205        this.readNonBoolValue();
206    }
207
208    function readNonBoolValue() public {
209        stdstore.target(address(test)).sig(test.tE.selector).read_bool();
210    }
211
212    function testStorageReadAddress() public {
213        address val = stdstore.target(address(test)).sig(test.tF.selector).read_address();
214        assertEq(val, address(1337));
215    }
216
217    function testStorageReadUint() public {
218        uint256 val = stdstore.target(address(test)).sig(test.exists.selector).read_uint();
219        assertEq(val, 1);
220    }
221
222    function testStorageReadInt() public {
223        int256 val = stdstore.target(address(test)).sig(test.tG.selector).read_int();
224        assertEq(val, type(int256).min);
225    }
226}
227
228contract StorageTest {
229    uint256 public exists = 1;
230    mapping(address => uint256) public map_addr;
231    mapping(uint256 => uint256) public map_uint;
232    mapping(address => uint256) public map_packed;
233    mapping(address => UnpackedStruct) public map_struct;
234    mapping(address => mapping(address => uint256)) public deep_map;
235    mapping(address => mapping(address => UnpackedStruct)) public deep_map_struct;
236    UnpackedStruct public basic;
237
238    uint248 public tA;
239    bool public tB;
240
241    bool public tC = false;
242    uint248 public tD = 1;
243
244    struct UnpackedStruct {
245        uint256 a;
246        uint256 b;
247    }
248
249    mapping(address => bool) public map_bool;
250
251    bytes32 public tE = hex"1337";
252    address public tF = address(1337);
253    int256 public tG = type(int256).min;
254    bool public tH = true;
255
256    constructor() {
257        basic = UnpackedStruct({a: 1337, b: 1337});
258
259        uint256 two = (1 << 128) | 1;
260        map_packed[msg.sender] = two;
261        map_packed[address(uint160(1337))] = 1 << 128;
262    }
263
264    function read_struct_upper(address who) public view returns (uint256) {
265        return map_packed[who] >> 128;
266    }
267
268    function read_struct_lower(address who) public view returns (uint256) {
269        return map_packed[who] & ((1 << 128) - 1);
270    }
271
272    function hidden() public view returns (bytes32 t) {
273        bytes32 slot = keccak256("my.random.var");
274        /// @solidity memory-safe-assembly
275        assembly {
276            t := sload(slot)
277        }
278    }
279
280    function const() public pure returns (bytes32 t) {
281        t = bytes32(hex"1337");
282    }
283}