From 9b57f5ced364bfb170995f1aae8d31940f22980c Mon Sep 17 00:00:00 2001
From: Thomas Clement <thomasaclement98@gmail.com>
Date: Tue, 11 Jun 2024 18:16:51 -0400
Subject: [PATCH] Packed Slot Helpers

---
 src/FraxTest.sol        | 52 +++++++++++++++++++++++++++++++++++++++++
 test/TestSlotDump.t.sol | 26 +++++++++++++++++++--
 2 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/src/FraxTest.sol b/src/FraxTest.sol
index afc6395..a28014c 100644
--- a/src/FraxTest.sol
+++ b/src/FraxTest.sol
@@ -9,6 +9,10 @@ abstract contract FraxTest is VmHelper, Test {
     uint256[] internal snapShotIds;
     function()[] internal setupFunctions;
 
+    // ========================================================================
+    // ~~~~~~~~~~~~~~~~~~ Different State Testing Helpers ~~~~~~~~~~~~~~~~~~~~~
+    // ========================================================================
+
     modifier useMultipleSetupFunctions() {
         if (snapShotIds.length == 0) _;
         for (uint256 i = 0; i < snapShotIds.length; i++) {
@@ -30,6 +34,13 @@ abstract contract FraxTest is VmHelper, Test {
         }
     }
 
+    // ========================================================================
+    // ~~~~~~~~~~~~~~~~~~~~~~~ Storage Slot Helpers ~~~~~~~~~~~~~~~~~~~~~~~~~~~
+    // ========================================================================
+
+    /// @notice Helper function to dump the storage slots of a contract to the console
+    /// @param target      The target contract whose state we wish to view
+    /// @param slotsToDump The # of lower storage slots we want to log
     function dumpStorageLayout(address target, uint256 slotsToDump) internal view {
         console.log("===================================");
         console.log("Storage dump for: ", target);
@@ -41,5 +52,46 @@ abstract contract FraxTest is VmHelper, Test {
         }
     }
 
+    /// @notice Helper function for unpacking low level storage slots
+    /// @param dataToUnpack The bytes32|uint256 of the slot to unpack
+    /// @param offset       The bits to remove st. the target bits are LSB
+    /// @param lenOfTarget  The length target result in bits
+    /// @return result      The target bits expressed as a uint256
+    function unpackBits(
+        uint256 dataToUnpack,
+        uint256 offset,
+        uint256 lenOfTarget
+    ) internal pure returns (uint256 result) {
+        uint256 mask = (1 << lenOfTarget) - 1;
+        result = (dataToUnpack >> offset) & mask;
+    }
+
+    function unpackBits(
+        bytes32 dataToUnpack,
+        uint256 offset,
+        uint256 lenOfTarget
+    ) internal pure returns (uint256 result) {
+        uint256 mask = (1 << lenOfTarget) - 1;
+        result = (uint256(dataToUnpack) >> offset) & mask;
+    }
+
+    function unpackBitsAndLogUint(
+        uint256 dataToUnpack,
+        uint256 offset,
+        uint256 lenOfTarget
+    ) internal pure returns (uint256 result) {
+        result = unpackBits(dataToUnpack, offset, lenOfTarget);
+        console.log(result);
+    }
+
+    function unpackBitsAndLogUint(
+        bytes32 dataToUnpack,
+        uint256 offset,
+        uint256 lenOfTarget
+    ) internal pure returns (uint256 result) {
+        result = unpackBits(dataToUnpack, offset, lenOfTarget);
+        console.log(result);
+    }
+
     error VmDidNotRevert(uint256 _snapshotId);
 }
diff --git a/test/TestSlotDump.t.sol b/test/TestSlotDump.t.sol
index e514eaf..7f853f2 100644
--- a/test/TestSlotDump.t.sol
+++ b/test/TestSlotDump.t.sol
@@ -7,12 +7,29 @@ contract TestSlotDump is FraxTest {
     address instance;
 
     function testDumpSlots() public {
-        instance = address(new Bravo());
+        instance = address(new Charlie());
         dumpStorageLayout(instance, 15);
     }
+
+    function testUnpackSlot() public {
+        instance = address(new Charlie());
+        dumpStorageLayout(instance, 15);
+
+        bytes32 packedSlot = vm.load(address(instance), bytes32(uint(9)));
+        uint256 unpacked1 = unpackBits(packedSlot, 0, 96);
+        uint256 unpacked2 = unpackBits(packedSlot, 96, 160);
+
+        /// @notice `unpacked1` is `uint96` expressed as `uint256`
+        assertEq(22222222222222222222, unpacked1);
+        assertEq(22222222222222222222, uint96(unpacked1));
+
+        /// @notice `unpacked2` is `address` expressed as `uint256`
+        assertEq(uint256(uint160(0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE)), unpacked2);
+        assertEq(0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE, address(uint160(unpacked2)));
+    }
 }
 
-// ================== Helpers =============
+// ================== Helpers ==================
 
 contract Alpha {
     address owner = address(0xC0ffee);
@@ -24,3 +41,8 @@ contract Bravo is Alpha {
     uint256[5] gap;
     bytes32 someOtherValue = 0xAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA;
 }
+
+contract Charlie is Bravo {
+    uint96 packed1 = 22222222222222222222;
+    address packed2 = 0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE;
+}