feat(store-sync): initial query client (#2355)

packages/store-sync/package.json

@@ -71,6 +71,7 @@
     "change-case": "^5.2.0",
     "debug": "^4.3.4",
     "drizzle-orm": "^0.28.5",
+    "fast-deep-equal": "^3.1.3",
     "kysely": "^0.26.3",
     "postgres": "^3.3.5",
     "rxjs": "7.5.5",

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

@@ -0,0 +1,52 @@
+import { StoreConfig, Tables, ResolvedStoreConfig } from "@latticexyz/store";
+import { Hex } from "viem";
+import { storeTables, worldTables } from "../common";
+import { StaticPrimitiveType, DynamicPrimitiveType } from "@latticexyz/schema-type";
+
+// TODO: move to some common utils file/module/package
+export type satisfy<base, t extends base> = t;
+
+// TODO: make this better with new config resolver
+export type AllTables<
+  config extends StoreConfig,
+  extraTables extends Tables | undefined = undefined,
+> = ResolvedStoreConfig<config>["tables"] &
+  (extraTables extends Tables ? extraTables : Record<never, never>) &
+  typeof storeTables &
+  typeof worldTables;
+
+export type TableField = {
+  readonly tableId: Hex;
+  readonly field: string;
+};
+
+export type TableSubject = {
+  readonly tableId: Hex;
+  readonly subject: readonly string[];
+};
+
+export type ConditionLiteral = boolean | number | bigint | string;
+
+export type ComparisonCondition = {
+  readonly left: TableField;
+  readonly op: "<" | "<=" | "=" | ">" | ">=" | "!=";
+  // TODO: add support for TableField
+  readonly right: ConditionLiteral;
+};
+
+export type InCondition = {
+  readonly left: TableField;
+  readonly op: "in";
+  readonly right: readonly ConditionLiteral[];
+};
+
+export type QueryCondition = satisfy<{ readonly op: string }, ComparisonCondition | InCondition>;
+
+// TODO: move this into some "wire" type and then make this more client specific (uses config to validate)
+export type Query = {
+  readonly from: readonly TableSubject[];
+  readonly except?: readonly TableSubject[];
+  readonly where?: readonly QueryCondition[];
+};
+
+export type QueryResultSubject = readonly (StaticPrimitiveType | DynamicPrimitiveType)[];

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

@@ -0,0 +1,70 @@
+import { TableRecord } from "../zustand/common";
+import { Query, QueryResultSubject } from "./common";
+import { Table } from "@latticexyz/store";
+import { groupBy, uniqueBy } from "@latticexyz/common/utils";
+import { encodeAbiParameters } from "viem";
+import { matchesCondition } from "./matchesCondition";
+
+// This assumes upstream has fully validated query
+// This also assumes we have full records, which may not always be the case and we may need some way to request records for a given table subject
+// We don't carry around config types here for ease, they get handled by the wrapping `query` function
+
+type QueryParameters<table extends Table> = {
+  readonly records: readonly TableRecord<table>[];
+  readonly query: Query;
+};
+
+// TODO: make condition types smarter, so condition literal matches the field primitive type
+
+export function findSubjects<table extends Table>({
+  records: initialRecords,
+  query,
+}: QueryParameters<table>): readonly QueryResultSubject[] {
+  const targetTables = Object.fromEntries(
+    uniqueBy([...query.from, ...(query.except ?? [])], (subject) => subject.tableId).map((subject) => [
+      subject.tableId,
+      subject.subject,
+    ]),
+  );
+  const fromTableIds = new Set(query.from.map((subject) => subject.tableId));
+  const exceptTableIds = new Set((query.except ?? []).map((subject) => subject.tableId));
+
+  // TODO: store/lookup subjects separately rather than mapping each time so we can "memoize" better?
+  const records = initialRecords
+    .filter((record) => targetTables[record.table.tableId])
+    .map((record) => {
+      const subjectFields = targetTables[record.table.tableId];
+      const schema = { ...record.table.keySchema, ...record.table.valueSchema };
+      const fields = { ...record.key, ...record.value };
+      const subject = subjectFields.map((field) => fields[field]);
+      const subjectSchema = subjectFields.map((field) => schema[field]);
+      const id = encodeAbiParameters(subjectSchema, subject);
+      return {
+        ...record,
+        schema,
+        fields,
+        subjectSchema,
+        subject,
+        id,
+      };
+    });
+
+  const matchedSubjects = Array.from(groupBy(records, (record) => record.id).values())
+    .map((records) => ({
+      id: records[0].id,
+      subject: records[0].subject,
+      records,
+    }))
+    .filter(({ records }) => {
+      // make sure our matched subject has no records in `query.except` tables
+      return exceptTableIds.size ? !records.some((record) => exceptTableIds.has(record.table.tableId)) : true;
+    })
+    .filter(({ records }) => {
+      // make sure our matched subject has records in all `query.from` tables
+      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));
+
+  return matchedSubjects.map((match) => match.subject);
+}

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

@@ -0,0 +1,46 @@
+import { Table } from "@latticexyz/store";
+import { TableRecord } from "../zustand/common";
+import { ComparisonCondition, ConditionLiteral, QueryCondition, TableSubject } from "./common";
+
+type MatchedSubjectRecord<table extends Table> = TableRecord<table> & {
+  fields: TableRecord<table>["key"] & TableRecord<table>["value"];
+};
+
+type MatchedSubject<table extends Table> = {
+  readonly subject: TableSubject;
+  readonly records: readonly MatchedSubjectRecord<table>[];
+};
+
+const comparisons = {
+  "<": (left, right) => left < right,
+  "<=": (left, right) => left <= right,
+  "=": (left, right) => left === right,
+  ">": (left, right) => left > right,
+  ">=": (left, right) => left >= right,
+  "!=": (left, right) => left !== right,
+} as const satisfies Record<ComparisonCondition["op"], (left: ConditionLiteral, right: ConditionLiteral) => boolean>;
+
+export function matchesCondition<table extends Table>(
+  condition: QueryCondition,
+  subject: MatchedSubject<table>,
+): boolean {
+  switch (condition.op) {
+    case "<":
+    case "<=":
+    case "=":
+    case ">":
+    case ">=":
+    case "!=":
+      return subject.records.some(
+        (record) =>
+          record.table.tableId === condition.left.tableId &&
+          comparisons[condition.op](record.fields[condition.left.field], condition.right),
+      );
+    case "in":
+      return subject.records.some(
+        (record) =>
+          record.table.tableId === condition.left.tableId &&
+          condition.right.includes(record.fields[condition.left.field]),
+      );
+  }
+}

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

@@ -0,0 +1,140 @@
+import { beforeAll, describe, expect, it } from "vitest";
+import { createHydratedStore, tables } from "./test/createHydratedStore";
+import { query } from "./query";
+import { deployMockGame } from "../../test/mockGame";
+import { Address } from "viem";
+
+describe("query", async () => {
+  let worldAddress: Address;
+  beforeAll(async () => {
+    worldAddress = await deployMockGame();
+  });
+
+  it("can get players with a position", async () => {
+    const { store } = await createHydratedStore(worldAddress);
+    const result = await query(store, {
+      from: [{ tableId: tables.Position.tableId, subject: ["player"] }],
+    });
+
+    expect(result).toMatchInlineSnapshot(`
+      [
+        [
+          "0x1D96F2f6BeF1202E4Ce1Ff6Dad0c2CB002861d3e",
+        ],
+        [
+          "0x328809Bc894f92807417D2dAD6b7C998c1aFdac6",
+        ],
+        [
+          "0x078cf0753dd50f7C56F20B3Ae02719EA199BE2eb",
+        ],
+        [
+          "0xdBa86119a787422C593ceF119E40887f396024E2",
+        ],
+      ]
+    `);
+  });
+
+  it("can get players at position (3, 5)", async () => {
+    const { store } = await createHydratedStore(worldAddress);
+    const result = await query(store, {
+      from: [{ tableId: tables.Position.tableId, subject: ["player"] }],
+      where: [
+        { left: { tableId: tables.Position.tableId, field: "x" }, op: "=", right: 3 },
+        { left: { tableId: tables.Position.tableId, field: "y" }, op: "=", right: 5 },
+      ],
+    });
+
+    expect(result).toMatchInlineSnapshot(`
+      [
+        [
+          "0x328809Bc894f92807417D2dAD6b7C998c1aFdac6",
+        ],
+        [
+          "0x078cf0753dd50f7C56F20B3Ae02719EA199BE2eb",
+        ],
+      ]
+    `);
+  });
+
+  it("can get players within the bounds of (-5, -5) and (5, 5)", async () => {
+    const { store } = await createHydratedStore(worldAddress);
+    const result = await query(store, {
+      from: [{ tableId: tables.Position.tableId, subject: ["player"] }],
+      where: [
+        { left: { tableId: tables.Position.tableId, field: "x" }, op: ">=", right: -5 },
+        { left: { tableId: tables.Position.tableId, field: "x" }, op: "<=", right: 5 },
+        { left: { tableId: tables.Position.tableId, field: "y" }, op: ">=", right: -5 },
+        { left: { tableId: tables.Position.tableId, field: "y" }, op: "<=", right: 5 },
+      ],
+    });
+
+    expect(result).toMatchInlineSnapshot(`
+      [
+        [
+          "0x1D96F2f6BeF1202E4Ce1Ff6Dad0c2CB002861d3e",
+        ],
+        [
+          "0x328809Bc894f92807417D2dAD6b7C998c1aFdac6",
+        ],
+        [
+          "0x078cf0753dd50f7C56F20B3Ae02719EA199BE2eb",
+        ],
+      ]
+    `);
+  });
+
+  it("can get players that are still alive", async () => {
+    const { store } = await createHydratedStore(worldAddress);
+    const result = await query(store, {
+      from: [
+        { tableId: tables.Position.tableId, subject: ["player"] },
+        { tableId: tables.Health.tableId, subject: ["player"] },
+      ],
+      where: [{ left: { tableId: tables.Health.tableId, field: "health" }, op: "!=", right: 0n }],
+    });
+
+    expect(result).toMatchInlineSnapshot(`
+      [
+        [
+          "0x1D96F2f6BeF1202E4Ce1Ff6Dad0c2CB002861d3e",
+        ],
+        [
+          "0x328809Bc894f92807417D2dAD6b7C998c1aFdac6",
+        ],
+      ]
+    `);
+  });
+
+  it("can get all players in grassland", async () => {
+    const { store } = await createHydratedStore(worldAddress);
+    const result = await query(store, {
+      from: [{ tableId: tables.Terrain.tableId, subject: ["x", "y"] }],
+      where: [{ left: { tableId: tables.Terrain.tableId, field: "terrainType" }, op: "=", right: 2 }],
+    });
+
+    expect(result).toMatchInlineSnapshot(`
+      [
+        [
+          3,
+          5,
+        ],
+      ]
+    `);
+  });
+
+  it("can get all players without health (e.g. spectator)", async () => {
+    const { store } = await createHydratedStore(worldAddress);
+    const result = await query(store, {
+      from: [{ tableId: tables.Position.tableId, subject: ["player"] }],
+      except: [{ tableId: tables.Health.tableId, subject: ["player"] }],
+    });
+
+    expect(result).toMatchInlineSnapshot(`
+      [
+        [
+          "0xdBa86119a787422C593ceF119E40887f396024E2",
+        ],
+      ]
+    `);
+  });
+});

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

@@ -0,0 +1,28 @@
+import { ZustandStore } from "../zustand";
+import { AllTables, Query, QueryResultSubject } from "./common";
+import { StoreConfig, Tables } from "@latticexyz/store";
+import { findSubjects } from "./findSubjects";
+
+// TODO: validate query
+//       - one subject per table
+//       - underlying subject field types match
+//       - only keys as subjects for now?
+//       - subjects and conditions all have valid fields
+//       - can only compare like types?
+//       - `where` tables are in `from`
+
+// TODO: make query smarter/config aware for shorthand
+// TODO: make condition types smarter, so condition literal matches the field primitive type
+// TODO: return matching records alongside subjects? because the record subset may be smaller than what querying for records with matching subjects
+
+type QueryResult<query extends Query> = readonly QueryResultSubject[];
+
+export async function query<config extends StoreConfig, extraTables extends Tables | undefined = undefined>(
+  store: ZustandStore<AllTables<config, extraTables>>,
+  query: Query,
+): Promise<QueryResult<typeof query>> {
+  const records = Object.values(store.getState().records);
+  const matches = findSubjects({ records, query });
+
+  return matches;
+}