refactor(query,store-sync): rework query types, refactor query client (…

…#2436)

packages/query/src/api.ts

@@ -1,48 +1,86 @@
 import { Hex } from "viem";
 import { StaticPrimitiveType, DynamicPrimitiveType } from "@latticexyz/schema-type";
 import { satisfy } from "@latticexyz/common/type-utils";
+import { SchemaToPrimitives } from "@latticexyz/store";
+import { Table } from "@latticexyz/store/config/v2";
 
 /**
  * These types represent the "over the wire" protocol (i.e. JSON) for the query API.
+ *
+ * Currently always returns matching records for each subject. We may add separate endpoints and return types for just subjects later.
  */
 
-export type TableField = {
+// TODO: decide if we want to support stronger types here (e.g. table generic that constrains subjects, records, etc.)
+// TODO: decide if/how we want to add block number throughout (esp as it relates to instant sequencing)
+// TODO: separate set of types for querying just
+
+export type QueryTable = {
   readonly tableId: Hex;
   readonly field: string;
 };
 
-export type TableSubject = {
+export type QuerySubject = {
   readonly tableId: Hex;
   readonly subject: readonly string[];
 };
 
-export type ConditionLiteral = boolean | number | bigint | string;
+// TODO: should we exclude arrays? might be hard to support array comparisons in SQL
+export type ConditionLiteral = StaticPrimitiveType | DynamicPrimitiveType;
 
 export type ComparisonCondition = {
-  readonly left: TableField;
+  readonly left: QueryTable;
   readonly op: "<" | "<=" | "=" | ">" | ">=" | "!=";
-  // TODO: add support for TableField
+  // TODO: add support for QueryTable
   readonly right: ConditionLiteral;
 };
 
 export type InCondition = {
-  readonly left: TableField;
+  readonly left: QueryTable;
   readonly op: "in";
   readonly right: readonly ConditionLiteral[];
 };
 
 export type QueryCondition = satisfy<{ readonly op: string }, ComparisonCondition | InCondition>;
 
 export type Query = {
-  readonly from: readonly TableSubject[];
-  readonly except?: readonly TableSubject[];
+  readonly from: readonly QuerySubject[];
+  readonly except?: readonly QuerySubject[];
   readonly where?: readonly QueryCondition[];
 };
 
-export type QueryResultSubject = readonly (StaticPrimitiveType | DynamicPrimitiveType)[];
+export type PrimitiveType = StaticPrimitiveType | DynamicPrimitiveType;
+
+export type ResultRecord = {
+  readonly tableId: Hex;
+  readonly keyTuple: readonly Hex[];
+  readonly primaryKey: readonly StaticPrimitiveType[];
+  readonly fields: SchemaToPrimitives<Table["schema"]>;
+};
+
+export type Subject = readonly PrimitiveType[];
+
+export type SubjectRecords = {
+  readonly subject: Subject;
+  readonly records: readonly ResultRecord[];
+};
+
+// TODO: consider flattening this to be more like `ResultRecord & { subject: Subject }`
+export type SubjectRecord = {
+  readonly subject: Subject;
+  readonly record: ResultRecord;
+};
+
+// TODO: for change event, should this include previous record?
+// TODO: use merge helper instead of `&` intersection?
+export type SubjectEvent = SubjectRecord & {
+  /**
+   * `enter` = a new subject+record pair matched
+   * `exit` = a subject+record pair no longer matches
+   * `change` = the record oft he subject+record pair changed
+   */
+  readonly type: "enter" | "exit" | "change";
+};
 
-export type QueryResult = {
-  subjects: readonly QueryResultSubject[];
-  // TODO: matched records
-  // TODO: block number
+export type Result = {
+  readonly subjects: readonly SubjectRecords[];
 };

packages/query/src/common.ts

@@ -1,10 +1,13 @@
-import { Table, Schema } from "@latticexyz/store/config/v2";
+import { StaticPrimitiveType } from "@latticexyz/schema-type";
+import { SchemaToPrimitives } from "@latticexyz/store";
+import { Table } from "@latticexyz/store/config/v2";
+import { Hex } from "viem";
 
-export type schemaAbiTypes<schema extends Schema> = {
-  [key in keyof schema]: schema[key]["type"];
-};
-
-export type TableRecord<table extends Table = Table> = {
+export type TableRecord<table extends Table> = {
   readonly table: table;
-  readonly fields: schemaAbiTypes<table["schema"]>;
+  // TODO: refine to just static types
+  // TODO: add helper to extract primary key of primitive types from table primary key + field values
+  readonly primaryKey: readonly StaticPrimitiveType[];
+  readonly keyTuple: readonly Hex[];
+  readonly fields: SchemaToPrimitives<table["schema"]>;
 };

packages/query/src/findSubjects.ts

@@ -1,8 +1,8 @@
 import { encodeAbiParameters } from "viem";
 import { Table } from "@latticexyz/store/config/v2";
 import { groupBy, uniqueBy } from "@latticexyz/common/utils";
-import { Query, QueryResultSubject } from "./api";
-import { matchesCondition } from "./matchesCondition";
+import { Query, SubjectRecords } from "./api";
+import { matchRecords } from "./matchRecords";
 import { TableRecord } from "./common";
 
 // This assumes upstream has fully validated query
@@ -14,16 +14,12 @@ export type FindSubjectsParameters<table extends Table> = {
   readonly query: Query;
 };
 
-export type FindSubjectsResult = {
-  readonly subjects: readonly QueryResultSubject[];
-};
-
 // TODO: make condition types smarter? so condition literal matches the field primitive type
 
 export function findSubjects<table extends Table>({
   records: initialRecords,
   query,
-}: FindSubjectsParameters<table>): FindSubjectsResult {
+}: FindSubjectsParameters<table>): readonly SubjectRecords[] {
   const targetTables = Object.fromEntries(
     uniqueBy([...query.from, ...(query.except ?? [])], (subject) => subject.tableId).map((subject) => [
       subject.tableId,
@@ -64,9 +60,28 @@ export function findSubjects<table extends Table>({
       const tableIds = new Set(records.map((record) => record.table.tableId));
       return tableIds.size === fromTableIds.size;
     })
-    .filter((match) => (query.where ? query.where.every((condition) => matchesCondition(condition, match)) : true));
+    .map((match) => {
+      if (!query.where) return match;
+
+      let records: readonly TableRecord<table>[] = match.records;
+      for (const condition of query.where) {
+        if (!records.length) break;
+        records = matchRecords(condition, records);
+      }
+
+      return { ...match, records };
+    })
+    .filter((match) => match.records.length > 0);
 
-  const subjects = matchedSubjects.map((match) => match.subject);
+  const subjects = matchedSubjects.map((match) => ({
+    subject: match.subject,
+    records: match.records.map((record) => ({
+      tableId: record.table.tableId,
+      primaryKey: record.primaryKey,
+      keyTuple: record.keyTuple,
+      fields: record.fields,
+    })),
+  }));
 
-  return { subjects };
+  return subjects;
 }

packages/query/src/index.ts

@@ -1,3 +1,3 @@
 export * from "./api";
 export * from "./findSubjects";
-export * from "./matchesCondition";
+export * from "./matchRecords";

packages/query/src/matchesCondition.ts → packages/query/src/matchRecords.ts

@@ -2,11 +2,6 @@ import { Table } from "@latticexyz/store/config/v2";
 import { ComparisonCondition, ConditionLiteral, QueryCondition } from "./api";
 import { TableRecord } from "./common";
 
-export type MatchedSubject<table extends Table = Table> = {
-  readonly subject: readonly string[];
-  readonly records: readonly TableRecord<table>[];
-};
-
 const comparisons = {
   "<": (left, right) => left < right,
   "<=": (left, right) => left <= right,
@@ -16,26 +11,24 @@ const comparisons = {
   "!=": (left, right) => left !== right,
 } as const satisfies Record<ComparisonCondition["op"], (left: ConditionLiteral, right: ConditionLiteral) => boolean>;
 
-// TODO: adapt this to return matching records, not just a boolean
-
-export function matchesCondition<table extends Table = Table>(
+export function matchRecords<table extends Table = Table>(
   condition: QueryCondition,
-  subject: MatchedSubject<table>,
-): boolean {
+  records: readonly TableRecord<table>[],
+): readonly TableRecord<table>[] {
   switch (condition.op) {
     case "<":
     case "<=":
     case "=":
     case ">":
     case ">=":
     case "!=":
-      return subject.records.some(
+      return records.filter(
         (record) =>
           record.table.tableId === condition.left.tableId &&
           comparisons[condition.op](record.fields[condition.left.field], condition.right),
       );
     case "in":
-      return subject.records.some(
+      return records.filter(
         (record) =>
           record.table.tableId === condition.left.tableId &&
           condition.right.includes(record.fields[condition.left.field]),

packages/store-sync/src/query-cache/common.ts

@@ -51,12 +51,4 @@ export type Query<tables extends Tables = Tables> = {
   readonly where?: readonly queryConditions<tables>[];
 };
 
-export type queryToResultSubject<query extends Query<tables>, tables extends Tables> = {
-  [table in keyof query["from"]]: table extends keyof tables
-    ? subjectSchemaToPrimitive<mapTuple<query["from"][table], schemaAbiTypes<tables[table]["schema"]>>>
-    : never;
-}[keyof query["from"]];
-
-export type QueryResult<query extends Query<tables>, tables extends Tables = Tables> = {
-  readonly subjects: readonly queryToResultSubject<query, tables>[];
-};
+export type extractTables<T> = T extends Query<infer tables> ? tables : never;

packages/store-sync/src/query-cache/createStorageAdapter.ts

@@ -1,21 +1,27 @@
 import { StorageAdapter } from "../common";
 import { QueryCacheStore, RawTableRecord, TableRecord } from "./createStore";
 import { hexToResource, resourceToLabel, spliceHex } from "@latticexyz/common";
-import { size } from "viem";
+import { Hex, concatHex, size } from "viem";
 import { decodeKey, decodeValueArgs } from "@latticexyz/protocol-parser";
 import { flattenSchema } from "../flattenSchema";
-import { getId } from "./getId";
 import debug from "debug";
-import { KeySchema } from "@latticexyz/store";
 import { Tables } from "./common";
 
-export type CreateStorageAdapterOptions<tables extends Tables> = {
-  store: QueryCacheStore<tables>;
+function getRecordId({ tableId, keyTuple }: { tableId: Hex; keyTuple: readonly Hex[] }): string {
+  return `${tableId}:${concatHex(keyTuple)}`;
+}
+
+export type CreateStorageAdapterOptions<store extends QueryCacheStore> = {
+  store: store;
 };
 
-export function createStorageAdapter<tables extends Tables>({
+// TS isn't happy when we use the strongly typed store for the function definition so we
+// overload the strongly typed variant here and allow the more generic version in the function.
+export function createStorageAdapter<store extends QueryCacheStore>({
   store,
-}: CreateStorageAdapterOptions<tables>): StorageAdapter {
+}: CreateStorageAdapterOptions<store>): StorageAdapter;
+
+export function createStorageAdapter({ store }: CreateStorageAdapterOptions<QueryCacheStore<Tables>>): StorageAdapter {
   return async function queryCacheStorageAdapter({ logs }) {
     const touchedIds = new Set<string>();
 
@@ -35,7 +41,7 @@ export function createStorageAdapter<tables extends Tables>({
         continue;
       }
 
-      const id = getId(log.args);
+      const id = getRecordId(log.args);
 
       if (log.eventName === "Store_SetRecord") {
         // debug("setting record", { namespace: table.namespace, name: table.name, id, log });
@@ -103,19 +109,15 @@ export function createStorageAdapter<tables extends Tables>({
     const records: readonly TableRecord[] = [
       ...previousRecords.filter((record) => !touchedIds.has(record.id)),
       ...Object.values(updatedRawRecords).map((rawRecord): TableRecord => {
-        // TODO: figure out how to define this without casting
-        const key = decodeKey(flattenSchema(rawRecord.table.keySchema as KeySchema), rawRecord.keyTuple) as TableRecord<
-          tables[keyof tables]
-        >["key"];
-
-        // TODO: figure out how to define this without casting
-        const value = decodeValueArgs(flattenSchema(rawRecord.table.valueSchema), rawRecord) as TableRecord<
-          tables[keyof tables]
-        >["value"];
+        const key = decodeKey(flattenSchema(rawRecord.table.keySchema), rawRecord.keyTuple);
+        const value = decodeValueArgs(flattenSchema(rawRecord.table.valueSchema), rawRecord);
 
         return {
           table: rawRecord.table,
           id: rawRecord.id,
+          keyTuple: rawRecord.keyTuple,
+          // TODO: do something to make sure this stays ordered?
+          primaryKey: Object.values(key),
           key,
           value,
           fields: { ...key, ...value },
@@ -124,8 +126,8 @@ export function createStorageAdapter<tables extends Tables>({
     ];
 
     store.setState({
-      rawRecords: rawRecords as readonly RawTableRecord<tables[keyof tables]>[],
-      records: records as readonly TableRecord<tables[keyof tables]>[],
+      rawRecords,
+      records,
     });
   };
 }

packages/store-sync/src/query-cache/createStore.ts

@@ -1,7 +1,9 @@
 import { StoreApi, UseBoundStore, create } from "zustand";
 import { Table } from "@latticexyz/store/config/v2";
-import { Tables, schemaAbiTypes } from "./common";
+import { Tables } from "./common";
 import { Hex } from "viem";
+import { StaticPrimitiveType } from "@latticexyz/schema-type";
+import { SchemaToPrimitives } from "@latticexyz/store";
 
 export type RawTableRecord<table extends Table = Table> = {
   readonly table: table;
@@ -17,9 +19,11 @@ export type TableRecord<table extends Table = Table> = {
   readonly table: table;
   /** @internal Internal unique ID */
   readonly id: string;
-  readonly key: schemaAbiTypes<table["keySchema"]>;
-  readonly value: schemaAbiTypes<table["valueSchema"]>;
-  readonly fields: schemaAbiTypes<table["schema"]>;
+  readonly keyTuple: readonly Hex[];
+  readonly primaryKey: readonly StaticPrimitiveType[];
+  readonly key: SchemaToPrimitives<table["keySchema"]>;
+  readonly value: SchemaToPrimitives<table["valueSchema"]>;
+  readonly fields: SchemaToPrimitives<table["schema"]>;
 };
 
 export type QueryCacheState<tables extends Tables = Tables> = {
@@ -28,7 +32,10 @@ export type QueryCacheState<tables extends Tables = Tables> = {
   readonly records: readonly TableRecord<tables[keyof tables]>[];
 };
 
-export type QueryCacheStore<tables extends Tables = Tables> = UseBoundStore<StoreApi<QueryCacheState<tables>>>;
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export type QueryCacheStore<tables extends Tables = any> = UseBoundStore<StoreApi<QueryCacheState<tables>>>;
+
+export type extractTables<T> = T extends QueryCacheStore<infer tables> ? tables : never;
 
 export type CreateStoreOptions<tables extends Tables = Tables> = {
   tables: tables;