Lua: early work on LuaQueryCollections

common/space_lua/query_collection.test.ts

@@ -0,0 +1,149 @@
+import { parseExpressionString } from "$common/space_lua/parse.ts";
+import { ArrayQueryCollection } from "./query_collection.ts";
+import {
+  LuaEnv,
+  LuaNativeJSFunction,
+  LuaStackFrame,
+} from "$common/space_lua/runtime.ts";
+import { assert, assertEquals } from "@std/assert";
+
+Deno.test("ArrayQueryCollection", async () => {
+  const rootEnv = new LuaEnv();
+  rootEnv.setLocal(
+    "build_name",
+    new LuaNativeJSFunction((a, b) => {
+      return Promise.resolve(a + " " + b);
+    }),
+  );
+
+  const collection = new ArrayQueryCollection([{ x: 1, y: 1 }, { x: 2, y: 2 }, {
+    x: 3,
+    y: 3,
+  }]);
+  const result = await collection.query(
+    {
+      where: parseExpressionString("x >= 2"),
+    },
+    rootEnv,
+    LuaStackFrame.lostFrame,
+  );
+  // console.log(result);
+  assert(result.length === 2);
+
+  // Test limit
+  const result2 = await collection.query(
+    {
+      limit: 1,
+    },
+    rootEnv,
+    LuaStackFrame.lostFrame,
+  );
+  assert(result2.length === 1);
+  assert(result2[0].x === 1);
+
+  // Test offset
+  const result3 = await collection.query(
+    {
+      offset: 1,
+    },
+    rootEnv,
+    LuaStackFrame.lostFrame,
+  );
+  assert(result3.length === 2);
+  assert(result3[0].x === 2);
+
+  // Test order by
+  const result4 = await collection.query(
+    {
+      orderBy: [{ expr: parseExpressionString("x"), desc: false }],
+    },
+    rootEnv,
+    LuaStackFrame.lostFrame,
+  );
+  assert(result4.length === 3);
+  assert(result4[0].x === 1);
+  assert(result4[1].x === 2);
+  assert(result4[2].x === 3);
+
+  // Test order by desc
+  const result5 = await collection.query(
+    {
+      orderBy: [{ expr: parseExpressionString("x"), desc: true }],
+    },
+    rootEnv,
+    LuaStackFrame.lostFrame,
+  );
+  assert(result5.length === 3);
+  assert(result5[0].x === 3);
+  assert(result5[1].x === 2);
+  assert(result5[2].x === 1);
+
+  // Test order by multiple fields
+  const collection2 = new ArrayQueryCollection([
+    { firstName: "John", lastName: "Doe" },
+    { firstName: "Alice", lastName: "Johnson" },
+    { firstName: "Jane", lastName: "Doe" },
+    { firstName: "Bob", lastName: "Johnson" },
+  ]);
+  const result6 = await collection2.query(
+    {
+      orderBy: [
+        { expr: parseExpressionString("lastName"), desc: false },
+        { expr: parseExpressionString("firstName"), desc: true },
+      ],
+    },
+    rootEnv,
+    LuaStackFrame.lostFrame,
+  );
+  assertEquals(result6[0].firstName, "John");
+  assertEquals(result6[0].lastName, "Doe");
+  assertEquals(result6[1].firstName, "Jane");
+  assertEquals(result6[1].lastName, "Doe");
+  assertEquals(result6[2].firstName, "Bob");
+  assertEquals(result6[2].lastName, "Johnson");
+  assertEquals(result6[3].firstName, "Alice");
+  assertEquals(result6[3].lastName, "Johnson");
+
+  // Test select
+  const result7 = await collection2.query(
+    {
+      select: [{ name: "firstName" }],
+    },
+    rootEnv,
+    LuaStackFrame.lostFrame,
+  );
+  assertEquals(result7[0].firstName, "John");
+  assertEquals(result7[0].lastName, undefined);
+
+  // Test select with expression
+  const result8 = await collection2.query(
+    {
+      select: [{
+        name: "fullName",
+        expr: parseExpressionString("firstName .. ' ' .. lastName"),
+      }],
+    },
+    rootEnv,
+    LuaStackFrame.lostFrame,
+  );
+  assertEquals(result8[0].fullName, "John Doe");
+  assertEquals(result8[1].fullName, "Alice Johnson");
+  assertEquals(result8[2].fullName, "Jane Doe");
+  assertEquals(result8[3].fullName, "Bob Johnson");
+
+  // Test select with native function
+  const result9 = await collection2.query(
+    {
+      select: [{
+        name: "fullName",
+        expr: parseExpressionString("build_name(firstName, lastName)"),
+      }],
+    },
+    rootEnv,
+    LuaStackFrame.lostFrame,
+  );
+  assertEquals(result9[0].fullName, "John Doe");
+  assertEquals(result9[1].fullName, "Alice Johnson");
+  assertEquals(result9[2].fullName, "Jane Doe");
+  assertEquals(result9[3].fullName, "Bob Johnson");
+});

common/space_lua/query_collection.ts

@@ -0,0 +1,131 @@
+import type { LuaExpression } from "$common/space_lua/ast.ts";
+import {
+  LuaEnv,
+  luaGet,
+  luaKeys,
+  type LuaStackFrame,
+} from "$common/space_lua/runtime.ts";
+import { evalExpression } from "$common/space_lua/eval.ts";
+import { asyncQuickSort } from "$common/space_lua/util.ts";
+
+function buildItemEnv(item: any, env: LuaEnv, sf: LuaStackFrame): LuaEnv {
+  const itemEnv = new LuaEnv(env);
+  for (const key of luaKeys(item)) {
+    itemEnv.setLocal(key, luaGet(item, key, sf));
+  }
+  return itemEnv;
+}
+
+export type LuaOrderBy = {
+  expr: LuaExpression;
+  desc: boolean;
+};
+
+export type LuaSelect = {
+  name: string;
+  expr?: LuaExpression;
+};
+
+/**
+ * Represents a query for a collection
+ */
+export type LuaCollectionQuery = {
+  // The filter expression evaluated with Lua
+  where?: LuaExpression;
+  // The order by expression evaluated with Lua
+  orderBy?: LuaOrderBy[];
+  // The select expression evaluated with Lua
+  select?: LuaSelect[];
+  // The limit of the query
+  limit?: number;
+  // The offset of the query
+  offset?: number;
+};
+
+export interface LuaQueryCollection {
+  query(
+    query: LuaCollectionQuery,
+    env: LuaEnv,
+    sf: LuaStackFrame,
+  ): Promise<any[]>;
+}
+
+/**
+ * Implements a query collection for a regular JavaScript array
+ */
+export class ArrayQueryCollection<T> implements LuaQueryCollection {
+  constructor(private readonly array: T[]) {}
+
+  async query(
+    query: LuaCollectionQuery,
+    env: LuaEnv,
+    sf: LuaStackFrame,
+  ): Promise<any[]> {
+    let result: any[] = [];
+
+    // Filter the array
+    for (const item of this.array) {
+      const itemEnv = buildItemEnv(item, env, sf);
+      if (query.where && !await evalExpression(query.where, itemEnv, sf)) {
+        continue;
+      }
+      result.push(item);
+    }
+
+    // Apply the select
+    if (query.select) {
+      const newResult = [];
+      for (const item of result) {
+        const itemEnv = buildItemEnv(item, env, sf);
+        const newItem: Record<string, any> = {};
+        for (const select of query.select) {
+          if (select.expr) {
+            newItem[select.name] = await evalExpression(
+              select.expr,
+              itemEnv,
+              sf,
+            );
+          } else {
+            newItem[select.name] = item[select.name];
+          }
+        }
+        newResult.push(newItem);
+      }
+      result = newResult;
+    }
+
+    // Apply the order by
+    if (query.orderBy) {
+      result = await asyncQuickSort(result, async (a, b) => {
+        // Compare each orderBy clause until we find a difference
+        for (const { expr, desc } of query.orderBy!) {
+          const aEnv = buildItemEnv(a, env, sf);
+          const bEnv = buildItemEnv(b, env, sf);
+
+          const aVal = await evalExpression(expr, aEnv, sf);
+          const bVal = await evalExpression(expr, bEnv, sf);
+
+          if (aVal < bVal) {
+            return desc ? 1 : -1;
+          }
+          if (aVal > bVal) {
+            return desc ? -1 : 1;
+          }
+          // If equal, continue to next orderBy clause
+        }
+        return 0; // All orderBy clauses were equal
+      });
+    }
+
+    // Apply the limit and offset
+    if (query.limit !== undefined && query.offset !== undefined) {
+      result = result.slice(query.offset, query.offset + query.limit);
+    } else if (query.limit !== undefined) {
+      result = result.slice(0, query.limit);
+    } else if (query.offset !== undefined) {
+      result = result.slice(query.offset);
+    }
+
+    return Promise.resolve(result);
+  }
+}