Add transient storage slot support in StorageSlot.sol (#4980)

Co-authored-by: ernestognw <ernestognw@gmail.com>

scripts/generate/run.js

@@ -39,6 +39,7 @@ for (const [file, template] of Object.entries({
   'utils/SlotDerivation.sol': './templates/SlotDerivation.js',
   'utils/StorageSlot.sol': './templates/StorageSlot.js',
   'utils/Arrays.sol': './templates/Arrays.js',
+  'mocks/StorageSlotMock.sol': './templates/StorageSlotMock.js',
 })) {
   generateFromTemplate(file, template, './contracts/');
 }

scripts/generate/templates/StorageSlot.js

@@ -2,7 +2,7 @@ const format = require('../format-lines');
 const { TYPES } = require('./Slot.opts');
 
 const header = `\
-pragma solidity ^0.8.20;
+pragma solidity ^0.8.24;
 
 /**
  * @dev Library for reading and writing primitive types to specific storage slots.
@@ -28,7 +28,25 @@ pragma solidity ^0.8.20;
  *     }
  * }
  * \`\`\`
- * 
+ *
+ * Since version 5.1, this library also support writing and reading value types to and from transient storage.
+ *
+ *  * Example using transient storage:
+ * \`\`\`solidity
+ * contract Lock {
+ *     // Define the slot. Alternatively, use the SlotDerivation library to derive the slot.
+ *     bytes32 internal constant _LOCK_SLOT = 0xf4678858b2b588224636b8522b729e7722d32fc491da849ed75b3fdf3c84f542;
+ *
+ *     modifier locked() {
+ *         require(!_LOCK_SLOT.asBoolean().tload());
+ *
+ *         _LOCK_SLOT.asBoolean().tstore(true);
+ *         _;
+ *         _LOCK_SLOT.asBoolean().tstore(false);
+ *     }
+ * }
+ * \`\`\`
+ *
  * TIP: Consider using this library along with {SlotDerivation}.
  */
 `;
@@ -63,11 +81,47 @@ function get${name}Slot(${type} storage store) internal pure returns (${name}Slo
 }
 `;
 
+const udvt = ({ type, name }) => `\
+/**
+ * @dev UDVT that represent a slot holding a ${type}.
+ */
+type ${name}SlotType is bytes32;
+/**
+ * @dev Cast an arbitrary slot to a ${name}SlotType.
+ */
+function as${name}(bytes32 slot) internal pure returns (${name}SlotType) {
+  return ${name}SlotType.wrap(slot);
+}
+`;
+
+const transient = ({ type, name }) => `\
+/**
+ * @dev Load the value held at location \`slot\` in transient storage.
+ */
+function tload(${name}SlotType slot) internal view returns (${type} value) {
+  /// @solidity memory-safe-assembly
+  assembly {
+    value := tload(slot)
+  }
+}
+/**
+ * @dev Store \`value\` at location \`slot\` in transient storage.
+ */
+function tstore(${name}SlotType slot, ${type} value) internal {
+  /// @solidity memory-safe-assembly
+  assembly {
+    tstore(slot, value)
+  }
+}
+`;
+
 // GENERATE
 module.exports = format(
   header.trimEnd(),
   'library StorageSlot {',
   TYPES.map(type => struct(type)),
   TYPES.flatMap(type => [get(type), type.isValueType ? '' : getStorage(type)]),
+  TYPES.filter(type => type.isValueType).map(type => udvt(type)),
+  TYPES.filter(type => type.isValueType).map(type => transient(type)),
   '}',
 );

scripts/generate/templates/StorageSlotMock.js

@@ -0,0 +1,65 @@
+const format = require('../format-lines');
+const { TYPES } = require('./Slot.opts');
+
+const header = `\
+pragma solidity ^0.8.24;
+
+import {Multicall} from "../utils/Multicall.sol";
+import {StorageSlot} from "../utils/StorageSlot.sol";
+`;
+
+const storageSetValueType = ({ type, name }) => `\
+function set${name}Slot(bytes32 slot, ${type} value) public {
+    slot.get${name}Slot().value = value;
+}
+`;
+
+const storageGetValueType = ({ type, name }) => `\
+function get${name}Slot(bytes32 slot) public view returns (${type}) {
+    return slot.get${name}Slot().value;
+}
+`;
+
+const storageSetNonValueType = ({ type, name }) => `\
+mapping(uint256 key => ${type}) public ${type}Map;
+
+function set${name}Slot(bytes32 slot, ${type} calldata value) public {
+    slot.get${name}Slot().value = value;
+}
+
+function set${name}Storage(uint256 key, ${type} calldata value) public {
+    ${type}Map[key].get${name}Slot().value = value;
+}
+
+function get${name}Slot(bytes32 slot) public view returns (${type} memory) {
+    return slot.get${name}Slot().value;
+}
+
+function get${name}Storage(uint256 key) public view returns (${type} memory) {
+    return ${type}Map[key].get${name}Slot().value;
+}
+`;
+
+const transient = ({ type, name }) => `\
+event ${name}Value(bytes32 slot, ${type} value);
+
+function tload${name}(bytes32 slot) public {
+  emit ${name}Value(slot, slot.as${name}().tload());
+}
+
+function tstore(bytes32 slot, ${type} value) public {
+  slot.as${name}().tstore(value);
+}
+`;
+
+// GENERATE
+module.exports = format(
+  header.trimEnd(),
+  'contract StorageSlotMock is Multicall {',
+  'using StorageSlot for *;',
+  TYPES.filter(type => type.isValueType).map(type => storageSetValueType(type)),
+  TYPES.filter(type => type.isValueType).map(type => storageGetValueType(type)),
+  TYPES.filter(type => !type.isValueType).map(type => storageSetNonValueType(type)),
+  TYPES.filter(type => type.isValueType).map(type => transient(type)),
+  '}',
+);

scripts/helpers.js

@@ -1,10 +1,7 @@
 const iterate = require('../test/helpers/iterate');
+const strings = require('../test/helpers/strings');
 
 module.exports = {
-  // Capitalize the first char of a string
-  // Example: capitalize('uint256') → 'Uint256'
-  capitalize: str => str.charAt(0).toUpperCase() + str.slice(1),
-
-  // Iterate tools for the test helpers
   ...iterate,
+  ...strings,
 };