feat(store): add Storage.loadField for optimized loading of 32 bytes …

…or less from storage (#1512) Co-authored-by: alvarius <[email protected]>
latticexyz · Sep 16, 2023 · 0f3e2e0 · 0f3e2e0
1 parent be31306
commit 0f3e2e0
Show file tree

Hide file tree

Showing 5 changed files with 83 additions and 1 deletion.
diff --git a/.changeset/hungry-rings-doubt.md b/.changeset/hungry-rings-doubt.md
@@ -0,0 +1,5 @@
+---
+"@latticexyz/store": patch
+---
+
+Added `Storage.loadField` to optimize loading 32 bytes or less from storage (which is always the case when loading data for static fields).
diff --git a/packages/store/gas-report.json b/packages/store/gas-report.json
@@ -281,17 +281,29 @@
     "name": "MUD storage load (warm, 1 word)",
     "gasUsed": 412
   },
+  {
+    "file": "test/GasStorageLoad.t.sol",
+    "test": "testCompareStorageLoadMUD",
+    "name": "MUD storage load field (warm, 1 word)",
+    "gasUsed": 245
+  },
   {
     "file": "test/GasStorageLoad.t.sol",
     "test": "testCompareStorageLoadMUD",
     "name": "MUD storage load (warm, 1 word, partial)",
     "gasUsed": 460
   },
+  {
+    "file": "test/GasStorageLoad.t.sol",
+    "test": "testCompareStorageLoadMUD",
+    "name": "MUD storage load field (warm, 1 word, partial)",
+    "gasUsed": 378
+  },
   {
     "file": "test/GasStorageLoad.t.sol",
     "test": "testCompareStorageLoadMUD",
     "name": "MUD storage load (warm, 10 words)",
-    "gasUsed": 1914
+    "gasUsed": 1916
   },
   {
     "file": "test/GasStorageLoad.t.sol",

diff --git a/packages/store/src/Storage.sol b/packages/store/src/Storage.sol
@@ -209,4 +209,39 @@ library Storage {
       }
     }
   }
+
+  /**
+   * Load up to 32 bytes from storage at the given storagePointer and offset.
+   * The return value is left-aligned, the bytes beyond the length are not zeroed out,
+   * and the caller is expected to truncate as needed.
+   * Since fields are tightly packed, they can span more than one slot.
+   * Since the they're max 32 bytes, they can span at most 2 slots.
+   */
+  function loadField(uint256 storagePointer, uint256 length, uint256 offset) internal view returns (bytes32 result) {
+    if (offset >= 32) {
+      unchecked {
+        storagePointer += offset / 32;
+        offset %= 32;
+      }
+    }
+
+    // Extra data past length is not truncated
+    // This assumes that the caller will handle the overflow bits appropriately
+    assembly {
+      result := shl(mul(offset, 8), sload(storagePointer))
+    }
+
+    uint256 wordRemainder;
+    // (safe because of `offset %= 32` at the start)
+    unchecked {
+      wordRemainder = 32 - offset;
+    }
+
+    // Read from the next slot if field spans 2 slots
+    if (length > wordRemainder) {
+      assembly {
+        result := or(result, shr(mul(wordRemainder, 8), sload(add(storagePointer, 1))))
+      }
+    }
+  }
 }
diff --git a/packages/store/test/GasStorageLoad.t.sol b/packages/store/test/GasStorageLoad.t.sol
@@ -67,6 +67,9 @@ contract GasStorageLoadTest is Test, GasReporter {
     bytes memory encodedPartial = abi.encodePacked(valuePartial);
     bytes memory encoded9Words = abi.encodePacked(value9Words.length, value9Words);
 
+    bytes32 encodedFieldSimple = valueSimple;
+    bytes32 encodedFieldPartial = valuePartial;
+
     startGasReport("MUD storage load (cold, 1 word)");
     encodedSimple = Storage.load(SolidityStorage.STORAGE_SLOT_SIMPLE, encodedSimple.length, 0);
     endGasReport();
@@ -85,10 +88,19 @@ contract GasStorageLoadTest is Test, GasReporter {
     encodedSimple = Storage.load(SolidityStorage.STORAGE_SLOT_SIMPLE, encodedSimple.length, 0);
     endGasReport();
 
+    startGasReport("MUD storage load field (warm, 1 word)");
+    encodedFieldSimple = Storage.loadField(SolidityStorage.STORAGE_SLOT_SIMPLE, encodedSimple.length, 0);
+    endGasReport();
+
     startGasReport("MUD storage load (warm, 1 word, partial)");
     encodedPartial = Storage.load(SolidityStorage.STORAGE_SLOT_PARTIAL, encodedPartial.length, 16);
     endGasReport();
 
+    encodedFieldPartial = Storage.loadField(SolidityStorage.STORAGE_SLOT_PARTIAL, encodedSimple.length, 16);
+    startGasReport("MUD storage load field (warm, 1 word, partial)");
+    encodedFieldPartial = Storage.loadField(SolidityStorage.STORAGE_SLOT_PARTIAL, encodedSimple.length, 16);
+    endGasReport();
+
     startGasReport("MUD storage load (warm, 10 words)");
     encoded9Words = Storage.load(SolidityStorage.STORAGE_SLOT_BYTES, encoded9Words.length, 0);
     endGasReport();

diff --git a/packages/store/test/Storage.t.sol b/packages/store/test/Storage.t.sol
@@ -75,4 +75,22 @@ contract StorageTest is Test, GasReporter {
     Storage.store({ storagePointer: uint256(storagePointer), offset: offset, data: data });
     assertEq(Storage.load({ storagePointer: uint256(storagePointer), length: data.length, offset: offset }), data);
   }
+
+  function testStoreLoadFieldBytes32Fuzzy(bytes32 data, uint256 storagePointer, uint256 offset) public {
+    vm.assume(offset < type(uint256).max);
+    vm.assume(storagePointer > 0);
+    vm.assume(storagePointer < type(uint256).max - offset);
+
+    Storage.store({ storagePointer: storagePointer, offset: offset, data: abi.encodePacked((data)) });
+    assertEq(Storage.loadField({ storagePointer: storagePointer, length: 32, offset: offset }), data);
+  }
+
+  function testStoreLoadFieldBytes16Fuzzy(bytes16 data, uint256 storagePointer, uint256 offset) public {
+    vm.assume(offset < type(uint256).max);
+    vm.assume(storagePointer > 0);
+    vm.assume(storagePointer < type(uint256).max - offset);
+
+    Storage.store({ storagePointer: storagePointer, offset: offset, data: abi.encodePacked((data)) });
+    assertEq(bytes16(Storage.loadField({ storagePointer: storagePointer, length: 16, offset: offset })), data);
+  }
 }