feat(protocol-parser): add keySchema/valueSchema helpers (#1443)

Co-authored-by: alvarius <alvarius@lattice.xyz>

.changeset/tricky-comics-remain.md

@@ -0,0 +1,5 @@
+---
+"@latticexyz/store": minor
+---
+
+Moved `KeySchema`, `ValueSchema`, `SchemaToPrimitives` and `TableRecord` types into `@latticexyz/protocol-parser`

.changeset/wicked-tigers-return.md

@@ -0,0 +1,5 @@
+---
+"@latticexyz/protocol-parser": minor
+---
+
+Adds `decodeKey`, `decodeValue`, `encodeKey`, and `encodeValue` helpers to decode/encode from key/value schemas. Deprecates previous methods that use a schema object with static/dynamic field arrays, originally attempting to model our on-chain behavior but ended up not very ergonomic when working with table configs.

packages/protocol-parser/src/common.ts

@@ -1,11 +1,26 @@
-import { DynamicAbiType, StaticAbiType } from "@latticexyz/schema-type";
+import { DynamicAbiType, SchemaAbiType, SchemaAbiTypeToPrimitiveType, StaticAbiType } from "@latticexyz/schema-type";
 
+/** @deprecated use `KeySchema` or `ValueSchema` instead */
 export type Schema = {
   readonly staticFields: readonly StaticAbiType[];
   readonly dynamicFields: readonly DynamicAbiType[];
 };
 
+/** @deprecated use `KeySchema` and `ValueSchema` instead */
 export type TableSchema = {
   readonly keySchema: Schema; // TODO: refine to set dynamicFields to []
   readonly valueSchema: Schema;
 };
+
+export type KeySchema = Record<string, StaticAbiType>;
+export type ValueSchema = Record<string, SchemaAbiType>;
+
+/** Map a table schema like `{ value: "uint256" }` to its primitive types like `{ value: bigint }` */
+export type SchemaToPrimitives<TSchema extends ValueSchema> = {
+  [key in keyof TSchema]: SchemaAbiTypeToPrimitiveType<TSchema[key]>;
+};
+
+export type TableRecord<TKeySchema extends KeySchema = KeySchema, TValueSchema extends ValueSchema = ValueSchema> = {
+  key: SchemaToPrimitives<TKeySchema>;
+  value: SchemaToPrimitives<TValueSchema>;
+};

packages/protocol-parser/src/decodeKey.ts

@@ -0,0 +1,15 @@
+import { Hex } from "viem";
+import { SchemaToPrimitives, KeySchema } from "./common";
+import { decodeKeyTuple } from "./decodeKeyTuple";
+
+export function decodeKey<TSchema extends KeySchema>(
+  keySchema: TSchema,
+  data: readonly Hex[]
+): SchemaToPrimitives<TSchema> {
+  // TODO: refactor and move all decodeKeyTuple logic into this method so we can delete decodeKeyTuple
+  const keyValues = decodeKeyTuple({ staticFields: Object.values(keySchema), dynamicFields: [] }, data);
+
+  return Object.fromEntries(
+    Object.keys(keySchema).map((name, i) => [name, keyValues[i]])
+  ) as SchemaToPrimitives<TSchema>;
+}

packages/protocol-parser/src/decodeKeyTuple.ts

@@ -4,6 +4,7 @@ import { Schema } from "./common";
 
 // key tuples are encoded in the same way as abi.encode, so we can decode them with viem
 
+/** @deprecated use `decodeKey` instead */
 export function decodeKeyTuple(keySchema: Schema, keyTuple: readonly Hex[]): StaticPrimitiveType[] {
   if (keySchema.staticFields.length !== keyTuple.length) {
     throw new Error(

packages/protocol-parser/src/decodeRecord.test.ts

@@ -21,4 +21,16 @@ describe("decodeRecord", () => {
       ]
     `);
   });
+
+  it("can decode an out of bounds array", () => {
+    const schema = { staticFields: [], dynamicFields: ["uint32[]"] } as const;
+    const values = decodeRecord(schema, "0x0000000000000000000000000000000000000000000000000400000000000004");
+    expect(values).toMatchInlineSnapshot(`
+      [
+        [
+          0,
+        ],
+      ]
+    `);
+  });
 });

packages/protocol-parser/src/decodeRecord.ts

@@ -4,20 +4,22 @@ import {
   staticAbiTypeToByteLength,
   dynamicAbiTypeToDefaultValue,
 } from "@latticexyz/schema-type";
-import { Hex, sliceHex } from "viem";
+import { Hex } from "viem";
 import { Schema } from "./common";
 import { decodeDynamicField } from "./decodeDynamicField";
 import { decodeStaticField } from "./decodeStaticField";
 import { hexToPackedCounter } from "./hexToPackedCounter";
 import { staticDataLength } from "./staticDataLength";
+import { readHex } from "./readHex";
 
+/** @deprecated use `decodeValue` instead */
 export function decodeRecord(schema: Schema, data: Hex): readonly (StaticPrimitiveType | DynamicPrimitiveType)[] {
   const values: (StaticPrimitiveType | DynamicPrimitiveType)[] = [];
 
   let bytesOffset = 0;
   schema.staticFields.forEach((fieldType) => {
     const fieldByteLength = staticAbiTypeToByteLength[fieldType];
-    const value = decodeStaticField(fieldType, sliceHex(data, bytesOffset, bytesOffset + fieldByteLength));
+    const value = decodeStaticField(fieldType, readHex(data, bytesOffset, bytesOffset + fieldByteLength));
     bytesOffset += fieldByteLength;
     values.push(value);
   });
@@ -37,13 +39,13 @@ export function decodeRecord(schema: Schema, data: Hex): readonly (StaticPrimiti
   }
 
   if (schema.dynamicFields.length > 0) {
-    const dataLayout = hexToPackedCounter(sliceHex(data, bytesOffset, bytesOffset + 32));
+    const dataLayout = hexToPackedCounter(readHex(data, bytesOffset, bytesOffset + 32));
     bytesOffset += 32;
 
     schema.dynamicFields.forEach((fieldType, i) => {
       const dataLength = dataLayout.fieldByteLengths[i];
       if (dataLength > 0) {
-        const value = decodeDynamicField(fieldType, sliceHex(data, bytesOffset, bytesOffset + dataLength));
+        const value = decodeDynamicField(fieldType, readHex(data, bytesOffset, bytesOffset + dataLength));
         bytesOffset += dataLength;
         values.push(value);
       } else {

packages/protocol-parser/src/decodeValue.ts

@@ -0,0 +1,16 @@
+import { isStaticAbiType, isDynamicAbiType } from "@latticexyz/schema-type";
+import { Hex } from "viem";
+import { SchemaToPrimitives, ValueSchema } from "./common";
+import { decodeRecord } from "./decodeRecord";
+
+export function decodeValue<TSchema extends ValueSchema>(valueSchema: TSchema, data: Hex): SchemaToPrimitives<TSchema> {
+  const staticFields = Object.values(valueSchema).filter(isStaticAbiType);
+  const dynamicFields = Object.values(valueSchema).filter(isDynamicAbiType);
+
+  // TODO: refactor and move all decodeRecord logic into this method so we can delete decodeRecord
+  const valueTuple = decodeRecord({ staticFields, dynamicFields }, data);
+
+  return Object.fromEntries(
+    Object.keys(valueSchema).map((name, i) => [name, valueTuple[i]])
+  ) as SchemaToPrimitives<TSchema>;
+}

packages/protocol-parser/src/encodeField.ts

@@ -8,10 +8,13 @@ export function encodeField<TSchemaAbiType extends SchemaAbiType>(
 ): Hex {
   if (isArrayAbiType(fieldType) && Array.isArray(value)) {
     const staticFieldType = arrayAbiTypeToStaticAbiType(fieldType);
-    return encodePacked(
-      value.map(() => staticFieldType),
-      value
-    );
+    // TODO: we can remove conditional once this is fixed: https://github.com/wagmi-dev/viem/pull/1147
+    return value.length === 0
+      ? "0x"
+      : encodePacked(
+          value.map(() => staticFieldType),
+          value
+        );
   }
   return encodePacked([fieldType], [value]);
 }

packages/protocol-parser/src/encodeKey.ts

@@ -0,0 +1,10 @@
+import { isStaticAbiType } from "@latticexyz/schema-type";
+import { Hex } from "viem";
+import { SchemaToPrimitives, KeySchema } from "./common";
+import { encodeKeyTuple } from "./encodeKeyTuple";
+
+export function encodeKey<TSchema extends KeySchema>(keySchema: TSchema, key: SchemaToPrimitives<TSchema>): Hex[] {
+  const staticFields = Object.values(keySchema).filter(isStaticAbiType);
+  // TODO: refactor and move all encodeKeyTuple logic into this method so we can delete encodeKeyTuple
+  return encodeKeyTuple({ staticFields, dynamicFields: [] }, Object.values(key));
+}

packages/protocol-parser/src/encodeKeyTuple.ts

@@ -2,6 +2,7 @@ import { StaticPrimitiveType } from "@latticexyz/schema-type";
 import { Hex, encodeAbiParameters } from "viem";
 import { Schema } from "./common";
 
+/** @deprecated use `encodeKey` instead */
 export function encodeKeyTuple(keySchema: Schema, keyTuple: StaticPrimitiveType[]): Hex[] {
   return keyTuple.map((key, index) => encodeAbiParameters([{ type: keySchema.staticFields[index] }], [key]));
 }

packages/protocol-parser/src/encodeRecord.test.ts

@@ -6,7 +6,7 @@ describe("encodeRecord", () => {
     const schema = { staticFields: ["uint32", "uint128"], dynamicFields: ["uint32[]", "string"] } as const;
     const hex = encodeRecord(schema, [1, 2n, [3, 4], "some string"]);
     expect(hex).toBe(
-      "0x0000000100000000000000000000000000000002000000000000130000000008000000000b0000000000000000000000000000000000000300000004736f6d6520737472696e67"
+      "0x0000000100000000000000000000000000000002000000000000000000000000000000000000000b0000000008000000000000130000000300000004736f6d6520737472696e67"
     );
   });
 
@@ -15,4 +15,10 @@ describe("encodeRecord", () => {
     const hex = encodeRecord(schema, [1, 2n]);
     expect(hex).toBe("0x0000000100000000000000000000000000000002");
   });
+
+  it("can encode an array to hex", () => {
+    const schema = { staticFields: [], dynamicFields: ["uint32[]"] } as const;
+    const hex = encodeRecord(schema, [[42]]);
+    expect(hex).toBe("0x00000000000000000000000000000000000000000000000004000000000000040000002a");
+  });
 });

packages/protocol-parser/src/encodeRecord.ts

@@ -3,6 +3,7 @@ import { Hex } from "viem";
 import { encodeField } from "./encodeField";
 import { Schema } from "./common";
 
+/** @deprecated use `encodeValue` instead */
 export function encodeRecord(schema: Schema, values: readonly (StaticPrimitiveType | DynamicPrimitiveType)[]): Hex {
   const staticValues = values.slice(0, schema.staticFields.length) as readonly StaticPrimitiveType[];
   const dynamicValues = values.slice(schema.staticFields.length) as readonly DynamicPrimitiveType[];
@@ -17,14 +18,14 @@ export function encodeRecord(schema: Schema, values: readonly (StaticPrimitiveTy
     encodeField(schema.dynamicFields[i], value).replace(/^0x/, "")
   );
 
-  const dynamicFieldByteLengths = dynamicDataItems.map((value) => value.length / 2);
+  const dynamicFieldByteLengths = dynamicDataItems.map((value) => value.length / 2).reverse();
   const dynamicTotalByteLength = dynamicFieldByteLengths.reduce((total, length) => total + BigInt(length), 0n);
 
   const dynamicData = dynamicDataItems.join("");
 
-  const packedCounter = `${encodeField("uint56", dynamicTotalByteLength).replace(/^0x/, "")}${dynamicFieldByteLengths
+  const packedCounter = `${dynamicFieldByteLengths
     .map((length) => encodeField("uint40", length).replace(/^0x/, ""))
-    .join("")}`.padEnd(64, "0");
+    .join("")}${encodeField("uint56", dynamicTotalByteLength).replace(/^0x/, "")}`.padStart(64, "0");
 
   return `0x${staticData}${packedCounter}${dynamicData}`;
 }

packages/protocol-parser/src/encodeValue.ts

@@ -0,0 +1,18 @@
+import { isStaticAbiType, isDynamicAbiType } from "@latticexyz/schema-type";
+import { Hex } from "viem";
+import { SchemaToPrimitives, ValueSchema } from "./common";
+import { encodeRecord } from "./encodeRecord";
+
+export function encodeValue<TSchema extends ValueSchema>(
+  valueSchema: TSchema,
+  value: SchemaToPrimitives<TSchema>
+): Hex {
+  const staticFields = Object.values(valueSchema).filter(isStaticAbiType);
+  const dynamicFields = Object.values(valueSchema).filter(isDynamicAbiType);
+
+  // TODO: refactor and move all encodeRecord logic into this method so we can delete encodeRecord
+
+  // This currently assumes fields/values are ordered by static, dynamic
+  // TODO: make sure we preserve ordering based on schema definition
+  return encodeRecord({ staticFields, dynamicFields }, Object.values(value));
+}

packages/protocol-parser/src/hexToPackedCounter.ts

@@ -1,7 +1,8 @@
-import { Hex, sliceHex } from "viem";
+import { Hex } from "viem";
 import { decodeStaticField } from "./decodeStaticField";
 import { decodeDynamicField } from "./decodeDynamicField";
 import { InvalidHexLengthForPackedCounterError, PackedCounterLengthMismatchError } from "./errors";
+import { readHex } from "./readHex";
 
 // Keep this logic in sync with PackedCounter.sol
 
@@ -18,9 +19,9 @@ export function hexToPackedCounter(data: Hex): {
     throw new InvalidHexLengthForPackedCounterError(data);
   }
 
-  const totalByteLength = decodeStaticField("uint56", sliceHex(data, 32 - 7, 32));
+  const totalByteLength = decodeStaticField("uint56", readHex(data, 32 - 7, 32));
   // TODO: use schema to make sure we only parse as many as we need (rather than zeroes at the end)?
-  const reversedFieldByteLengths = decodeDynamicField("uint40[]", sliceHex(data, 0, 32 - 7));
+  const reversedFieldByteLengths = decodeDynamicField("uint40[]", readHex(data, 0, 32 - 7));
   // Reverse the lengths
   const fieldByteLengths = Object.freeze([...reversedFieldByteLengths].reverse());
 

packages/protocol-parser/src/index.ts

@@ -2,16 +2,23 @@ export * from "./abiTypesToSchema";
 export * from "./common";
 export * from "./decodeDynamicField";
 export * from "./decodeField";
+export * from "./decodeKey";
 export * from "./decodeKeyTuple";
 export * from "./decodeRecord";
 export * from "./decodeStaticField";
+export * from "./decodeValue";
 export * from "./encodeField";
+export * from "./encodeKey";
 export * from "./encodeKeyTuple";
 export * from "./encodeRecord";
+export * from "./encodeValue";
 export * from "./errors";
 export * from "./hexToPackedCounter";
 export * from "./hexToSchema";
 export * from "./hexToTableSchema";
+export * from "./keySchemaToHex";
+export * from "./readHex";
 export * from "./schemaIndexToAbiType";
 export * from "./schemaToHex";
 export * from "./staticDataLength";
+export * from "./valueSchemaToHex";

packages/protocol-parser/src/keySchemaToHex.ts

@@ -0,0 +1,8 @@
+import { isStaticAbiType } from "@latticexyz/schema-type";
+import { Hex } from "viem";
+import { KeySchema } from "./common";
+import { schemaToHex } from "./schemaToHex";
+
+export function keySchemaToHex(schema: KeySchema): Hex {
+  return schemaToHex({ staticFields: Object.values(schema).filter(isStaticAbiType), dynamicFields: [] });
+}

packages/protocol-parser/src/readHex.test.ts

@@ -0,0 +1,15 @@
+import { describe, expect, it } from "vitest";
+import { readHex } from "./readHex";
+
+describe("readHex", () => {
+  it("can slice empty hex", () => {
+    expect(readHex("0x", 6)).toBe("0x");
+    expect(readHex("0x", 6, 10)).toBe("0x00000000");
+  });
+  it("can slice hex out of bounds", () => {
+    expect(readHex("0x000100", 1)).toBe("0x0100");
+    expect(readHex("0x000100", 1, 4)).toBe("0x010000");
+    expect(readHex("0x000100", 3)).toBe("0x");
+    expect(readHex("0x000100", 3, 4)).toBe("0x00");
+  });
+});

packages/protocol-parser/src/readHex.ts

@@ -0,0 +1,15 @@
+import { Hex } from "viem";
+
+/**
+ * Get the hex value at start/end positions. This will always return a valid hex string.
+ *
+ * If `start` is out of range, this returns `"0x"`.
+ *
+ * If `end` is specified and out of range, the result is right zero-padded to the desired length (`end - start`).
+ */
+export function readHex(data: Hex, start: number, end?: number): Hex {
+  return `0x${data
+    .replace(/^0x/, "")
+    .slice(start * 2, end != null ? end * 2 : undefined)
+    .padEnd(((end ?? start) - start) * 2, "0")}`;
+}

packages/protocol-parser/src/schemaToHex.ts

@@ -3,6 +3,7 @@ import { Hex } from "viem";
 import { Schema } from "./common";
 import { staticDataLength } from "./staticDataLength";
 
+/** @deprecated use `keySchemaToHex` or `valueSchemaToHex` instead */
 export function schemaToHex(schema: Schema): Hex {
   const staticSchemaTypes = schema.staticFields.map((abiType) => schemaAbiTypes.indexOf(abiType));
   const dynamicSchemaTypes = schema.dynamicFields.map((abiType) => schemaAbiTypes.indexOf(abiType));

packages/protocol-parser/src/valueSchemaToHex.ts

@@ -0,0 +1,11 @@
+import { isDynamicAbiType, isStaticAbiType } from "@latticexyz/schema-type";
+import { Hex } from "viem";
+import { ValueSchema } from "./common";
+import { schemaToHex } from "./schemaToHex";
+
+export function valueSchemaToHex(schema: ValueSchema): Hex {
+  return schemaToHex({
+    staticFields: Object.values(schema).filter(isStaticAbiType),
+    dynamicFields: Object.values(schema).filter(isDynamicAbiType),
+  });
+}

packages/store-sync/src/common.ts

@@ -1,17 +1,15 @@
-import { Address, Block, Hex, Log, PublicClient, TransactionReceipt } from "viem";
+import { Address, Block, Hex, Log, PublicClient } from "viem";
 import { GroupLogsByBlockNumberResult } from "@latticexyz/block-logs-stream";
 import {
   StoreConfig,
-  KeySchema,
-  ValueSchema,
   ConfigToKeyPrimitives as Key,
   ConfigToValuePrimitives as Value,
-  TableRecord,
   StoreEventsAbiItem,
   StoreEventsAbi,
 } from "@latticexyz/store";
 import { Observable } from "rxjs";
 import { BlockStorageOperations } from "./blockLogsToStorage";
+import { KeySchema, ValueSchema, TableRecord } from "@latticexyz/protocol-parser";
 
 export type ChainId = number;
 export type WorldId = `${ChainId}:${Address}`;

packages/store-sync/src/postgres/buildInternalTables.ts

@@ -2,6 +2,7 @@ import { integer, pgSchema, text } from "drizzle-orm/pg-core";
 import { DynamicAbiType, StaticAbiType } from "@latticexyz/schema-type";
 import { transformSchemaName } from "./transformSchemaName";
 import { asAddress, asBigInt, asJson, asNumber } from "./columnTypes";
+import { KeySchema, ValueSchema } from "@latticexyz/protocol-parser";
 
 // eslint-disable-next-line @typescript-eslint/explicit-function-return-type
 export function buildInternalTables() {
@@ -22,8 +23,8 @@ export function buildInternalTables() {
       tableId: text("table_id").notNull(),
       namespace: text("namespace").notNull(),
       name: text("name").notNull(),
-      keySchema: asJson<Record<string, StaticAbiType>>("key_schema").notNull(),
-      valueSchema: asJson<Record<string, StaticAbiType | DynamicAbiType>>("value_schema").notNull(),
+      keySchema: asJson<KeySchema>("key_schema").notNull(),
+      valueSchema: asJson<ValueSchema>("value_schema").notNull(),
       lastUpdatedBlockNumber: asBigInt("last_updated_block_number", "numeric"),
       // TODO: last block hash?
       lastError: text("last_error"),