feat: add support for key schemas

packages/cli/src/render-solidity/renderTable.ts

@@ -4,7 +4,16 @@ import { renderRecordMethods } from "./record.js";
 import { RenderTableOptions } from "./types.js";
 
 export function renderTable(options: RenderTableOptions) {
-  const { imports, libraryName, structName, staticResourceData, storeImportPath, fields, withRecordMethods } = options;
+  const {
+    imports,
+    libraryName,
+    structName,
+    staticResourceData,
+    storeImportPath,
+    fields,
+    withRecordMethods,
+    primaryKeys,
+  } = options;
 
   const { _typedTableId, _typedKeyArgs, _primaryKeysDefinition } = renderCommonData(options);
 
@@ -60,6 +69,13 @@ library ${libraryName} {
     return SchemaLib.encode(_schema);
   }
 
+  function getKeySchema() internal pure returns (Schema) {
+    SchemaType[] memory _schema = new SchemaType[](${primaryKeys.length});
+    ${renderList(primaryKeys, ({ enumName }, index) => `_schema[${index}] = SchemaType.${enumName};`)}
+
+    return SchemaLib.encode(_schema);
+  }
+
   /** Get the table's metadata */
   function getMetadata() internal pure returns (string memory, string[] memory) {
     string[] memory _fieldNames = new string[](${fields.length});
@@ -69,7 +85,7 @@ library ${libraryName} {
 
   /** Register the table's schema */
   function registerSchema(${_typedTableId}) internal {
-    StoreSwitch.registerSchema(_tableId, getSchema());
+    StoreSwitch.registerSchema(_tableId, getSchema(), getKeySchema());
   }
 
   /** Set the table's metadata */
@@ -84,7 +100,7 @@ ${
     : `
   /** Register the table's schema for the specified store */
   function registerSchema(${renderArguments([_typedTableId, "IStore _store"])}) internal {
-    _store.registerSchema(_tableId, getSchema());
+    _store.registerSchema(_tableId, getSchema(), getKeySchema());
   }
 
   /** Set the table's metadata for the specified store */

packages/cli/src/utils/deploy-v2.ts

@@ -71,18 +71,23 @@ export async function deploy(mudConfig: MUDConfig, deployConfig: DeployConfig):
 
   // Register tables
   promises.push(
-    ...Object.entries(mudConfig.tables).map(async ([tableName, { fileSelector, schema }]) => {
+    ...Object.entries(mudConfig.tables).map(async ([tableName, { fileSelector, schema, primaryKeys }]) => {
       console.log(chalk.blue(`Registering table ${tableName} at ${namespace}/${fileSelector}`));
 
       // Register table
       const schemaTypes = Object.values(schema).map((schemaOrUserType) => {
         return resolveSchemaOrUserTypeSimple(schemaOrUserType, mudConfig.userTypes);
       });
 
+      const keyTypes = Object.values(primaryKeys).map((schemaOrUserType) => {
+        return resolveSchemaOrUserTypeSimple(schemaOrUserType, mudConfig.userTypes);
+      });
+
       await fastTxExecute(WorldContract, "registerTable", [
         toBytes16(namespace),
         toBytes16(fileSelector),
         encodeSchema(schemaTypes),
+        encodeSchema(keyTypes),
       ]);
 
       // Register table metadata

packages/store/gas-report.txt

@@ -14,9 +14,9 @@
 (test/Gas.t.sol) | pass abi encoded bytes to external contract [someContract.doSomethingWithBytes(abiEncoded)]: 6554
 (test/Gas.t.sol) | pass custom encoded bytes to external contract [someContract.doSomethingWithBytes(customEncoded)]: 1381
 (test/Mixed.t.sol) | store Mixed struct in storage (native solidity) [testMixed = mixed]: 92050
-(test/Mixed.t.sol) | register Mixed schema [Mixed.registerSchema()]: 35823
-(test/Mixed.t.sol) | set record in Mixed [Mixed.set({ key: key, u32: 1, u128: 2, a32: a32, s: s })]: 111978
-(test/Mixed.t.sol) | get record from Mixed [MixedData memory mixed = Mixed.get(key)]: 13395
+(test/Mixed.t.sol) | register Mixed schema [Mixed.registerSchema()]: 61103
+(test/Mixed.t.sol) | set record in Mixed [Mixed.set({ key: key, u32: 1, u128: 2, a32: a32, s: s })]: 111931
+(test/Mixed.t.sol) | get record from Mixed [MixedData memory mixed = Mixed.get(key)]: 13374
 (test/PackedCounter.t.sol) | get value at index of PackedCounter [packedCounter.atIndex(3)]: 261
 (test/PackedCounter.t.sol) | set value at index of PackedCounter [packedCounter = packedCounter.setAtIndex(2, 5)]: 799
 (test/PackedCounter.t.sol) | pack uint16 array into PackedCounter [PackedCounter packedCounter = PackedCounterLib.pack(counters)]: 2152
@@ -44,48 +44,49 @@
 (test/Storage.t.sol) | store 1 storage slot [Storage.store({ storagePointer: storagePointer, data: originalDataFirstSlot })]: 22509
 (test/Storage.t.sol) | store 34 bytes over 3 storage slots (with offset and safeTrail)) [Storage.store({ storagePointer: storagePointer, offset: 31, data: data1 })]: 23164
 (test/Storage.t.sol) | load 34 bytes over 3 storage slots (with offset and safeTrail)) [bytes memory data = Storage.load({ storagePointer: storagePointer, length: 34, offset: 31 })]: 1104
-(test/StoreCore.t.sol) | access non-existing record [bytes memory data1 = StoreCore.getRecord(table, key)]: 7326
-(test/StoreCore.t.sol) | access static field of non-existing record [bytes memory data2 = StoreCore.getField(table, key, 0)]: 2990
-(test/StoreCore.t.sol) | access dynamic field of non-existing record [bytes memory data3 = StoreCore.getField(table, key, 1)]: 3605
-(test/StoreCore.t.sol) | delete record (complex data, 3 slots) [StoreCore.deleteRecord(table, key)]: 10992
-(test/StoreCore.t.sol) | Check for existence of table (existent) [StoreCore.hasTable(table)]: 1120
-(test/StoreCore.t.sol) | check for existence of table (non-existent) [StoreCore.hasTable(table2)]: 3144
-(test/StoreCore.t.sol) | register subscriber [StoreCore.registerStoreHook(table, subscriber)]: 67345
-(test/StoreCore.t.sol) | set record on table with subscriber [StoreCore.setRecord(table, key, data)]: 73830
-(test/StoreCore.t.sol) | set static field on table with subscriber [StoreCore.setField(table, key, 0, data)]: 29585
-(test/StoreCore.t.sol) | delete record on table with subscriber [StoreCore.deleteRecord(table, key)]: 24432
-(test/StoreCore.t.sol) | register subscriber [StoreCore.registerStoreHook(table, subscriber)]: 67345
-(test/StoreCore.t.sol) | set (dynamic) record on table with subscriber [StoreCore.setRecord(table, key, data)]: 167237
-(test/StoreCore.t.sol) | set (dynamic) field on table with subscriber [StoreCore.setField(table, key, 1, arrayDataBytes)]: 32628
-(test/StoreCore.t.sol) | delete (dynamic) record on table with subscriber [StoreCore.deleteRecord(table, key)]: 25910
-(test/StoreCore.t.sol) | push to field (1 slot, 1 uint32 item) [StoreCore.pushToField(table, key, 1, secondDataToPush)]: 16952
-(test/StoreCore.t.sol) | push to field (2 slots, 10 uint32 items) [StoreCore.pushToField(table, key, 2, thirdDataToPush)]: 39673
-(test/StoreCore.t.sol) | StoreCore: register schema [StoreCore.registerSchema(table, schema)]: 30539
-(test/StoreCore.t.sol) | StoreCore: get schema (warm) [Schema loadedSchema = StoreCore.getSchema(table)]: 1148
-(test/StoreCore.t.sol) | set complex record with dynamic data (4 slots) [StoreCore.setRecord(table, key, data)]: 107606
-(test/StoreCore.t.sol) | get complex record with dynamic data (4 slots) [bytes memory loadedData = StoreCore.getRecord(table, key)]: 6458
+(test/StoreCore.t.sol) | access non-existing record [bytes memory data1 = StoreCore.getRecord(table, key)]: 7322
+(test/StoreCore.t.sol) | access static field of non-existing record [bytes memory data2 = StoreCore.getField(table, key, 0)]: 2991
+(test/StoreCore.t.sol) | access dynamic field of non-existing record [bytes memory data3 = StoreCore.getField(table, key, 1)]: 3604
+(test/StoreCore.t.sol) | delete record (complex data, 3 slots) [StoreCore.deleteRecord(table, key)]: 10970
+(test/StoreCore.t.sol) | Check for existence of table (existent) [StoreCore.hasTable(table)]: 1117
+(test/StoreCore.t.sol) | check for existence of table (non-existent) [StoreCore.hasTable(table2)]: 3143
+(test/StoreCore.t.sol) | register subscriber [StoreCore.registerStoreHook(table, subscriber)]: 67279
+(test/StoreCore.t.sol) | set record on table with subscriber [StoreCore.setRecord(table, key, data)]: 73742
+(test/StoreCore.t.sol) | set static field on table with subscriber [StoreCore.setField(table, key, 0, data)]: 29511
+(test/StoreCore.t.sol) | delete record on table with subscriber [StoreCore.deleteRecord(table, key)]: 24365
+(test/StoreCore.t.sol) | register subscriber [StoreCore.registerStoreHook(table, subscriber)]: 67279
+(test/StoreCore.t.sol) | set (dynamic) record on table with subscriber [StoreCore.setRecord(table, key, data)]: 167149
+(test/StoreCore.t.sol) | set (dynamic) field on table with subscriber [StoreCore.setField(table, key, 1, arrayDataBytes)]: 32562
+(test/StoreCore.t.sol) | delete (dynamic) record on table with subscriber [StoreCore.deleteRecord(table, key)]: 25845
+(test/StoreCore.t.sol) | push to field (1 slot, 1 uint32 item) [StoreCore.pushToField(table, key, 1, secondDataToPush)]: 16930
+(test/StoreCore.t.sol) | push to field (2 slots, 10 uint32 items) [StoreCore.pushToField(table, key, 2, thirdDataToPush)]: 39652
+(test/StoreCore.t.sol) | StoreCore: register schema [StoreCore.registerSchema(table, schema, keySchema)]: 54796
+(test/StoreCore.t.sol) | StoreCore: get schema (warm) [Schema loadedSchema = StoreCore.getSchema(table)]: 1145
+(test/StoreCore.t.sol) | StoreCore: get key schema (warm) [Schema loadedKeySchema = StoreCore.getKeySchema(table)]: 1213
+(test/StoreCore.t.sol) | set complex record with dynamic data (4 slots) [StoreCore.setRecord(table, key, data)]: 107585
+(test/StoreCore.t.sol) | get complex record with dynamic data (4 slots) [bytes memory loadedData = StoreCore.getRecord(table, key)]: 6459
 (test/StoreCore.t.sol) | compare: Set complex record with dynamic data using native solidity [testStruct = _testStruct]: 116842
 (test/StoreCore.t.sol) | compare: Set complex record with dynamic data using abi.encode [testMapping[1234] = abi.encode(_testStruct)]: 267376
-(test/StoreCore.t.sol) | set dynamic length of dynamic index 0 [StoreCoreInternal._setDynamicDataLengthAtIndex(table, key, 0, 10)]: 23602
+(test/StoreCore.t.sol) | set dynamic length of dynamic index 0 [StoreCoreInternal._setDynamicDataLengthAtIndex(table, key, 0, 10)]: 23600
 (test/StoreCore.t.sol) | set dynamic length of dynamic index 1 [StoreCoreInternal._setDynamicDataLengthAtIndex(table, key, 1, 99)]: 1719
-(test/StoreCore.t.sol) | reduce dynamic length of dynamic index 0 [StoreCoreInternal._setDynamicDataLengthAtIndex(table, key, 0, 5)]: 1707
-(test/StoreCore.t.sol) | set static field (1 slot) [StoreCore.setField(table, key, 0, abi.encodePacked(firstDataBytes))]: 38017
+(test/StoreCore.t.sol) | reduce dynamic length of dynamic index 0 [StoreCoreInternal._setDynamicDataLengthAtIndex(table, key, 0, 5)]: 1708
+(test/StoreCore.t.sol) | set static field (1 slot) [StoreCore.setField(table, key, 0, abi.encodePacked(firstDataBytes))]: 37994
 (test/StoreCore.t.sol) | get static field (1 slot) [bytes memory loadedData = StoreCore.getField(table, key, 0)]: 2996
-(test/StoreCore.t.sol) | set static field (overlap 2 slot) [StoreCore.setField(table, key, 1, abi.encodePacked(secondDataBytes))]: 33030
+(test/StoreCore.t.sol) | set static field (overlap 2 slot) [StoreCore.setField(table, key, 1, abi.encodePacked(secondDataBytes))]: 33008
 (test/StoreCore.t.sol) | get static field (overlap 2 slot) [loadedData = StoreCore.getField(table, key, 1)]: 3884
-(test/StoreCore.t.sol) | set dynamic field (1 slot, first dynamic field) [StoreCore.setField(table, key, 2, thirdDataBytes)]: 55332
+(test/StoreCore.t.sol) | set dynamic field (1 slot, first dynamic field) [StoreCore.setField(table, key, 2, thirdDataBytes)]: 55311
 (test/StoreCore.t.sol) | get dynamic field (1 slot, first dynamic field) [loadedData = StoreCore.getField(table, key, 2)]: 3836
-(test/StoreCore.t.sol) | set dynamic field (1 slot, second dynamic field) [StoreCore.setField(table, key, 3, fourthDataBytes)]: 35472
-(test/StoreCore.t.sol) | get dynamic field (1 slot, second dynamic field) [loadedData = StoreCore.getField(table, key, 3)]: 3855
-(test/StoreCore.t.sol) | set static record (1 slot) [StoreCore.setRecord(table, key, data)]: 37297
+(test/StoreCore.t.sol) | set dynamic field (1 slot, second dynamic field) [StoreCore.setField(table, key, 3, fourthDataBytes)]: 35450
+(test/StoreCore.t.sol) | get dynamic field (1 slot, second dynamic field) [loadedData = StoreCore.getField(table, key, 3)]: 3856
+(test/StoreCore.t.sol) | set static record (1 slot) [StoreCore.setRecord(table, key, data)]: 37275
 (test/StoreCore.t.sol) | get static record (1 slot) [bytes memory loadedData = StoreCore.getRecord(table, key, schema)]: 1335
-(test/StoreCore.t.sol) | set static record (2 slots) [StoreCore.setRecord(table, key, data)]: 59861
+(test/StoreCore.t.sol) | set static record (2 slots) [StoreCore.setRecord(table, key, data)]: 59840
 (test/StoreCore.t.sol) | get static record (2 slots) [bytes memory loadedData = StoreCore.getRecord(table, key, schema)]: 1580
-(test/StoreCore.t.sol) | StoreCore: set table metadata [StoreCore.setMetadata(table, tableName, fieldNames)]: 251689
-(test/StoreMetadata.t.sol) | set record in StoreMetadataTable [StoreMetadata.set({ tableId: tableId, tableName: tableName, abiEncodedFieldNames: abi.encode(fieldNames) })]: 250156
-(test/StoreMetadata.t.sol) | get record from StoreMetadataTable [StoreMetadataData memory metadata = StoreMetadata.get(tableId)]: 12132
-(test/StoreSwitch.t.sol) | check if delegatecall [isDelegate = StoreSwitch.isDelegateCall()]: 693
+(test/StoreCore.t.sol) | StoreCore: set table metadata [StoreCore.setMetadata(table, tableName, fieldNames)]: 251645
+(test/StoreMetadata.t.sol) | set record in StoreMetadataTable [StoreMetadata.set({ tableId: tableId, tableName: tableName, abiEncodedFieldNames: abi.encode(fieldNames) })]: 250112
+(test/StoreMetadata.t.sol) | get record from StoreMetadataTable [StoreMetadataData memory metadata = StoreMetadata.get(tableId)]: 12110
+(test/StoreSwitch.t.sol) | check if delegatecall [isDelegate = StoreSwitch.isDelegateCall()]: 671
 (test/StoreSwitch.t.sol) | check if delegatecall [isDelegate = StoreSwitch.isDelegateCall()]: 627
-(test/Vector2.t.sol) | register Vector2 schema [Vector2.registerSchema()]: 32600
-(test/Vector2.t.sol) | set Vector2 record [Vector2.set({ key: key, x: 1, y: 2 })]: 38541
-(test/Vector2.t.sol) | get Vector2 record [Vector2Data memory vector = Vector2.get(key)]: 5063
+(test/Vector2.t.sol) | register Vector2 schema [Vector2.registerSchema()]: 57901
+(test/Vector2.t.sol) | set Vector2 record [Vector2.set({ key: key, x: 1, y: 2 })]: 38539
+(test/Vector2.t.sol) | get Vector2 record [Vector2Data memory vector = Vector2.get(key)]: 5064

packages/store/src/IStore.sol

@@ -8,10 +8,12 @@ interface IStore {
   event StoreSetField(uint256 table, bytes32[] key, uint8 schemaIndex, bytes data);
   event StoreDeleteRecord(uint256 table, bytes32[] key);
 
-  function registerSchema(uint256 table, Schema schema) external;
+  function registerSchema(uint256 table, Schema schema, Schema keySchema) external;
 
   function getSchema(uint256 table) external view returns (Schema schema);
 
+  function getKeySchema(uint256 table) external view returns (Schema schema);
+
   function setMetadata(uint256 table, string calldata tableName, string[] calldata fieldNames) external;
 
   // Set full record (including full dynamic data)

packages/store/src/Store.sol

@@ -14,6 +14,10 @@ abstract contract Store is IStore {
     schema = StoreCore.getSchema(table);
   }
 
+  function getKeySchema(uint256 table) public view virtual returns (Schema schema) {
+    schema = StoreCore.getKeySchema(table);
+  }
+
   // Get full record (static and dynamic data, load schema from storage)
   function getRecord(uint256 table, bytes32[] calldata key) public view virtual returns (bytes memory data) {
     data = StoreCore.getRecord(table, key);

packages/store/src/StoreCore.sol

@@ -40,7 +40,11 @@ library StoreCore {
    */
   function initialize() internal {
     // Register internal schema table
-    registerSchema(StoreCoreInternal.SCHEMA_TABLE, SchemaLib.encode(SchemaType.BYTES32));
+    registerSchema(
+      StoreCoreInternal.SCHEMA_TABLE,
+      SchemaLib.encode(SchemaType.BYTES32, SchemaType.BYTES32), // The Schema table's valueSchema is { valueSchema: BYTES32, keySchema: BYTES32 }
+      SchemaLib.encode(SchemaType.UINT256) // The Schema table's keySchema is { tableId: UINT256 }
+    );
 
     // Register other internal tables
     //
@@ -70,6 +74,16 @@ library StoreCore {
     }
   }
 
+  /**
+   * Get the key schema for the given tableId
+   */
+  function getKeySchema(uint256 tableId) internal view returns (Schema keySchema) {
+    keySchema = StoreCoreInternal._getKeySchema(tableId);
+    if (keySchema.isEmpty()) {
+      revert StoreCore_TableNotFound(tableId, tableId.toString());
+    }
+  }
+
   /**
    * Check if a table with the given tableId exists
    */
@@ -80,17 +94,18 @@ library StoreCore {
   /**
    * Register a new tableId schema
    */
-  function registerSchema(uint256 tableId, Schema schema) internal {
+  function registerSchema(uint256 tableId, Schema valueSchema, Schema keySchema) internal {
     // Verify the schema is valid
-    schema.validate();
+    valueSchema.validate();
+    keySchema.validate();
 
     // Verify the schema doesn't exist yet
     if (hasTable(tableId)) {
       revert StoreCore_TableAlreadyExists(tableId, tableId.toString());
     }
 
     // Register the schema
-    StoreCoreInternal._registerSchemaUnchecked(tableId, schema);
+    StoreCoreInternal._registerSchemaUnchecked(tableId, valueSchema, keySchema);
   }
 
   /**
@@ -372,17 +387,25 @@ library StoreCoreInternal {
     return Schema.wrap(Storage.load({ storagePointer: location }));
   }
 
+  function _getKeySchema(uint256 tableId) internal view returns (Schema) {
+    bytes32[] memory key = new bytes32[](1);
+    key[0] = bytes32(tableId);
+    uint256 location = StoreCoreInternal._getStaticDataLocation(SCHEMA_TABLE, key);
+    return Schema.wrap(Storage.load({ storagePointer: location + 0x20 }));
+  }
+
   /**
    * Register a new tableId schema without validity checks
    */
-  function _registerSchemaUnchecked(uint256 tableId, Schema schema) internal {
+  function _registerSchemaUnchecked(uint256 tableId, Schema valueSchema, Schema keySchema) internal {
     bytes32[] memory key = new bytes32[](1);
     key[0] = bytes32(tableId);
     uint256 location = _getStaticDataLocation(SCHEMA_TABLE, key);
-    Storage.store({ storagePointer: location, data: schema.unwrap() });
+    Storage.store({ storagePointer: location, data: valueSchema.unwrap() });
+    Storage.store({ storagePointer: location + 0x20, data: keySchema.unwrap() });
 
     // Emit an event to notify indexers
-    emit StoreCore.StoreSetRecord(SCHEMA_TABLE, key, abi.encodePacked(schema.unwrap()));
+    emit StoreCore.StoreSetRecord(SCHEMA_TABLE, key, abi.encodePacked(valueSchema.unwrap(), keySchema.unwrap()));
   }
 
   /************************************************************************

packages/store/src/StoreSwitch.sol

@@ -34,11 +34,11 @@ library StoreSwitch {
     }
   }
 
-  function registerSchema(uint256 table, Schema schema) internal {
+  function registerSchema(uint256 table, Schema schema, Schema keySchema) internal {
     if (isDelegateCall()) {
-      StoreCore.registerSchema(table, schema);
+      StoreCore.registerSchema(table, schema, keySchema);
     } else {
-      IStore(msg.sender).registerSchema(table, schema);
+      IStore(msg.sender).registerSchema(table, schema, keySchema);
     }
   }
 

packages/store/src/StoreView.sol

@@ -13,7 +13,7 @@ contract StoreView is Store {
   /**
    * Not implemented in StoreView
    */
-  function registerSchema(uint256, Schema) public virtual {
+  function registerSchema(uint256, Schema, Schema) public virtual {
     revert StoreView_NotImplemented();
   }
 

packages/store/src/tables/Callbacks.sol

@@ -28,6 +28,13 @@ library Callbacks {
     return SchemaLib.encode(_schema);
   }
 
+  function getKeySchema() internal pure returns (Schema) {
+    SchemaType[] memory _schema = new SchemaType[](1);
+    _schema[0] = SchemaType.BYTES32;
+
+    return SchemaLib.encode(_schema);
+  }
+
   /** Get the table's metadata */
   function getMetadata() internal pure returns (string memory, string[] memory) {
     string[] memory _fieldNames = new string[](1);
@@ -37,7 +44,7 @@ library Callbacks {
 
   /** Register the table's schema */
   function registerSchema() internal {
-    StoreSwitch.registerSchema(_tableId, getSchema());
+    StoreSwitch.registerSchema(_tableId, getSchema(), getKeySchema());
   }
 
   /** Set the table's metadata */