Update IERC7496 and DynamicTraits based on ethereum/ERCs PR 1400

- Add TraitDoesNotExist error to IERC7496 interface
- Add _validTraitKeys mapping to DynamicTraitsStorage
- Require trait keys to be registered before get/set operations
- Add _registerTraitKey and _isTraitKeyRegistered internal functions
- Add registerTraitKey public function to ERC721DynamicTraits
- Update OnchainTraits to register trait keys in _setTraitLabel
- Remove duplicate TraitDoesNotExist error from OnchainTraits
- Update tests to register trait keys before use
- Add tests for unregistered trait key behavior

Author: ryanio

src/dynamic-traits/DynamicTraits.sol

@@ -9,6 +9,8 @@ library DynamicTraitsStorage {
         mapping(uint256 tokenId => mapping(bytes32 traitKey => bytes32 traitValue)) _traits;
         /// @dev An offchain string URI that points to a JSON file containing trait metadata.
         string _traitMetadataURI;
+        /// @dev A mapping of valid trait keys.
+        mapping(bytes32 traitKey => bool isValid) _validTraitKeys;
     }
 
     bytes32 internal constant STORAGE_SLOT = keccak256("contracts.storage.erc7496-dynamictraits");
@@ -40,8 +42,14 @@ contract DynamicTraits is IERC7496 {
      * @param traitKey The trait key to get the value of
      */
     function getTraitValue(uint256 tokenId, bytes32 traitKey) public view virtual returns (bytes32 traitValue) {
+        // Revert if the trait key does not exist.
+        DynamicTraitsStorage.Layout storage layout = DynamicTraitsStorage.layout();
+        if (!layout._validTraitKeys[traitKey]) {
+            revert TraitDoesNotExist(traitKey);
+        }
+
         // Return the trait value.
-        return DynamicTraitsStorage.layout()._traits[tokenId][traitKey];
+        return layout._traits[tokenId][traitKey];
     }
 
     /**
@@ -86,8 +94,15 @@ contract DynamicTraits is IERC7496 {
      * @param newValue The new trait value to set
      */
     function setTrait(uint256 tokenId, bytes32 traitKey, bytes32 newValue) public virtual {
+        DynamicTraitsStorage.Layout storage layout = DynamicTraitsStorage.layout();
+
+        // Revert if the trait key does not exist.
+        if (!layout._validTraitKeys[traitKey]) {
+            revert TraitDoesNotExist(traitKey);
+        }
+
         // Revert if the new value is the same as the existing value.
-        bytes32 existingValue = DynamicTraitsStorage.layout()._traits[tokenId][traitKey];
+        bytes32 existingValue = layout._traits[tokenId][traitKey];
         if (existingValue == newValue) {
             revert TraitValueUnchanged();
         }
@@ -122,6 +137,22 @@ contract DynamicTraits is IERC7496 {
         emit TraitMetadataURIUpdated();
     }
 
+    /**
+     * @notice Register a trait key as valid.
+     * @param traitKey The trait key to register.
+     */
+    function _registerTraitKey(bytes32 traitKey) internal virtual {
+        DynamicTraitsStorage.layout()._validTraitKeys[traitKey] = true;
+    }
+
+    /**
+     * @notice Check if a trait key is registered.
+     * @param traitKey The trait key to check.
+     */
+    function _isTraitKeyRegistered(bytes32 traitKey) internal view virtual returns (bool) {
+        return DynamicTraitsStorage.layout()._validTraitKeys[traitKey];
+    }
+
     /**
      * @dev See {IERC165-supportsInterface}.
      */

src/dynamic-traits/ERC721DynamicTraits.sol

@@ -51,6 +51,11 @@ contract ERC721DynamicTraits is DynamicTraits, Ownable, ERC721 {
         _setTraitMetadataURI(uri);
     }
 
+    function registerTraitKey(bytes32 traitKey) external onlyOwner {
+        // Register the trait key as valid.
+        _registerTraitKey(traitKey);
+    }
+
     function supportsInterface(bytes4 interfaceId) public view virtual override(ERC721, DynamicTraits) returns (bool) {
         return ERC721.supportsInterface(interfaceId) || DynamicTraits.supportsInterface(interfaceId);
     }

src/dynamic-traits/OnchainTraits.sol

@@ -45,8 +45,6 @@ abstract contract OnchainTraits is Ownable, DynamicTraits {
 
     /// @notice Thrown when the caller does not have the privilege to set a trait
     error InsufficientPrivilege();
-    /// @notice Thrown when trying to set a trait that does not exist
-    error TraitDoesNotExist(bytes32 traitKey);
 
     constructor() {
         _initializeOwner(msg.sender);
@@ -170,6 +168,8 @@ abstract contract OnchainTraits is Ownable, DynamicTraits {
             valuesRequireValidation: _traitLabel.acceptableValues.length > 0,
             storedLabel: TraitLabelLib.store(_traitLabel)
         });
+        // Register the trait key in the base DynamicTraits contract.
+        _registerTraitKey(traitKey);
     }
 
     /**

src/dynamic-traits/interfaces/IERC7496.sol

@@ -2,6 +2,13 @@
 pragma solidity ^0.8.19;
 
 interface IERC7496 {
+    /* Errors */
+    /// @notice Thrown when trying to set a trait that does not exist.
+    error TraitDoesNotExist(bytes32 traitKey);
+    // Note: TokenDoesNotExist(uint256 tokenId) is specified in ERC-7496 but omitted here
+    // to avoid conflicts with token contracts that define their own error (e.g., Solady ERC721).
+    // Implementations MAY use their underlying token's error for non-existent tokens.
+
     /* Events */
     event TraitUpdated(bytes32 indexed traitKey, uint256 tokenId, bytes32 traitValue);
     event TraitUpdatedRange(bytes32 indexed traitKey, uint256 fromTokenId, uint256 toTokenId);

test/dynamic-traits/ERC721DynamicTraits.t.sol

@@ -36,6 +36,9 @@ contract ERC721DynamicTraitsTest is Test {
         bytes32 key = bytes32("testKey");
         bytes32 value = bytes32("foo");
         uint256 tokenId = 12345;
+
+        // Register the trait key before using it.
+        token.registerTraitKey(key);
         token.mint(address(this), tokenId);
 
         vm.expectEmit(true, true, true, true);
@@ -47,16 +50,23 @@ contract ERC721DynamicTraitsTest is Test {
     }
 
     function testOnlyOwnerCanSetValues() public {
+        bytes32 key = bytes32("test");
+        // Register the trait key before testing access control.
+        token.registerTraitKey(key);
+
         address alice = makeAddr("alice");
         vm.prank(alice);
         vm.expectRevert(abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, alice));
-        token.setTrait(0, bytes32("test"), bytes32("test"));
+        token.setTrait(0, key, bytes32("test"));
     }
 
     function testSetTrait_Unchanged() public {
         bytes32 key = bytes32("testKey");
         bytes32 value = bytes32("foo");
         uint256 tokenId = 1;
+
+        // Register the trait key before using it.
+        token.registerTraitKey(key);
         token.mint(address(this), tokenId);
 
         token.setTrait(tokenId, key, value);
@@ -70,6 +80,10 @@ contract ERC721DynamicTraitsTest is Test {
         bytes32 value1 = bytes32("foo");
         bytes32 value2 = bytes32("bar");
         uint256 tokenId = 1;
+
+        // Register the trait keys before using them.
+        token.registerTraitKey(key1);
+        token.registerTraitKey(key2);
         token.mint(address(this), tokenId);
 
         token.setTrait(tokenId, key1, value1);
@@ -99,6 +113,9 @@ contract ERC721DynamicTraitsTest is Test {
         bytes32 value = bytes32(uint256(1));
         uint256 tokenId = 1;
 
+        // Register the trait key so we can test token non-existence.
+        token.registerTraitKey(key);
+
         vm.expectRevert(abi.encodeWithSelector(IERC721Errors.ERC721NonexistentToken.selector, tokenId));
         token.setTrait(tokenId, key, value);
 
@@ -112,6 +129,9 @@ contract ERC721DynamicTraitsTest is Test {
     function testGetTraitValue_DefaultZeroValue() public {
         bytes32 key = bytes32("testKey");
         uint256 tokenId = 1;
+
+        // Register the trait key before using it.
+        token.registerTraitKey(key);
         token.mint(address(this), tokenId);
 
         bytes32 value = token.getTraitValue(tokenId, key);
@@ -120,4 +140,26 @@ contract ERC721DynamicTraitsTest is Test {
         bytes32[] memory values = token.getTraitValues(tokenId, Solarray.bytes32s(key));
         assertEq(values[0], bytes32(0), "should return bytes32(0)");
     }
+
+    function testGetTraitValue_UnregisteredTraitKey() public {
+        bytes32 key = bytes32("unregisteredKey");
+        uint256 tokenId = 1;
+        token.mint(address(this), tokenId);
+
+        vm.expectRevert(abi.encodeWithSelector(IERC7496.TraitDoesNotExist.selector, key));
+        token.getTraitValue(tokenId, key);
+
+        vm.expectRevert(abi.encodeWithSelector(IERC7496.TraitDoesNotExist.selector, key));
+        token.getTraitValues(tokenId, Solarray.bytes32s(key));
+    }
+
+    function testSetTrait_UnregisteredTraitKey() public {
+        bytes32 key = bytes32("unregisteredKey");
+        bytes32 value = bytes32("foo");
+        uint256 tokenId = 1;
+        token.mint(address(this), tokenId);
+
+        vm.expectRevert(abi.encodeWithSelector(IERC7496.TraitDoesNotExist.selector, key));
+        token.setTrait(tokenId, key, value);
+    }
 }

test/dynamic-traits/ERC721DynamicTraitsMultiUpdate.t.sol

@@ -26,6 +26,9 @@ contract ERC721DynamicTraitsMultiUpdateTest is Test {
         uint256 fromTokenId = 1;
         uint256 toTokenId = 100;
 
+        // Register the trait key before using it.
+        token.registerTraitKey(key);
+
         for (uint256 tokenId = fromTokenId; tokenId <= toTokenId; tokenId++) {
             token.mint(address(this), tokenId);
         }
@@ -49,6 +52,9 @@ contract ERC721DynamicTraitsMultiUpdateTest is Test {
             values[i] = bytes32(i);
         }
 
+        // Register the trait key before using it.
+        token.registerTraitKey(key);
+
         for (uint256 tokenId = fromTokenId; tokenId <= toTokenId; tokenId++) {
             token.mint(address(this), tokenId);
         }
@@ -68,6 +74,9 @@ contract ERC721DynamicTraitsMultiUpdateTest is Test {
         bytes32 value = bytes32("foo");
         uint256[] memory tokenIds = Solarray.uint256s(1, 10, 20, 50);
 
+        // Register the trait key before using it.
+        token.registerTraitKey(key);
+
         for (uint256 i = 0; i < tokenIds.length; i++) {
             token.mint(address(this), tokenIds[i]);
         }
@@ -90,6 +99,9 @@ contract ERC721DynamicTraitsMultiUpdateTest is Test {
             values[i] = bytes32(i * 1000);
         }
 
+        // Register the trait key before using it.
+        token.registerTraitKey(key);
+
         for (uint256 i = 0; i < tokenIds.length; i++) {
             token.mint(address(this), tokenIds[i]);
         }