feat(oquery): add defKeyQuery(), refactor/fix types

- add conditional types to fix return type inference (QueryFn/KeyQueryFn)
- add defKeyQuery()
- extract arrayQuery()/objQuery()
- optimize arrayQuery() by pre-selecting query impl

packages/oquery/src/api.ts

@@ -23,6 +23,8 @@ export type QueryObj = IObjectOf<any>;
  * - l => literal
  * - n => null / wildcard
  * - f => function / predicate
+ *
+ * @internal
  */
 export type QueryType =
     | "lll"
@@ -53,6 +55,9 @@ export type QueryType =
     | "nnf"
     | "nnn";
 
+/**
+ * @internal
+ */
 export type QueryImpl = Fn6<
     QueryObj,
     QueryObj,
@@ -63,42 +68,85 @@ export type QueryImpl = Fn6<
     void
 >;
 
+/**
+ * @internal
+ */
 export type QueryImpls = Record<QueryType, QueryImpl>;
 
 /**
- * Query function overloads.
+ * Takes an object of this structure `{ s1: { p1: o, p2: ... }, s2: { p1: o
+ * }...}`, matches all entries using provided `s`(ubject), `p`(redicate) and
+ * `o`(object) terms. Returns new object of matched results (format depends
+ * on query config given to {@link defQuery}).
+ *
+ * @remarks
+ * If `res` is provided, results will be injected in that object. Otherwise
+ * a new result object will be created.
  */
-export interface QueryFn {
-    /**
-     * Takes an object of this structure `{ s1: { p1: o, p2: ... }, s2: { p1: o
-     * }...}`, matches all entries using provided `s`(ubject), `p`(redicate) and
-     * `o`(object) terms. Returns new object of matched results (format depends
-     * on query config given to {@link defQuery}).
-     *
-     * @remarks
-     * If `res` is provided, results will be injected in that object. Otherwise
-     * a new result object will be created.
-     */
-    (
-        obj: QueryObj,
-        s: SPInputTerm,
-        p: SPInputTerm,
-        o: OTerm,
-        res?: QueryObj
-    ): QueryObj;
+export type ObjQueryFn<T extends QueryObj> = (
+    obj: T,
+    s: SPInputTerm,
+    p: SPInputTerm,
+    o: OTerm,
+    res?: QueryObj
+) => QueryObj;
 
-    /**
-     * Takes a source array of objects with this structure: [{p1: o, p2: ...},
-     * ...]`, and matches each using provided `p`(redicate) and `o`bject terms.
-     * Returns new array of matched results (result object format depends on
-     * query config given to {@link defQuery}).
-     * @remarks
-     * If `res` is provided, results will be appended to that array. Otherwise
-     * a new result array will be created.
-     */
-    (obj: QueryObj[], p: SPInputTerm, o: OTerm, res?: any[]): QueryObj[];
-}
+/**
+ * Takes a source array of objects with this structure: [{p1: o, p2: ...},
+ * ...]`, and matches each item using provided `p`(redicate) and `o`bject terms.
+ * Returns new array of matched results (result object format depends on query
+ * config given to {@link defQuery}).
+ * @remarks
+ * If `res` is provided, results will be appended to that array. Otherwise a new
+ * result array will be created.
+ */
+export type ArrayQueryFn<T extends QueryObj[]> = (
+    src: T,
+    p: SPInputTerm,
+    o: OTerm,
+    res?: QueryObj[]
+) => QueryObj[];
 
+/**
+ * Similar to {@link ObjQueryFn}, but only collects and returns a set of
+ * matching `s` keys.
+ */
+export type ObjKeyQueryFn<T extends QueryObj> = (
+    obj: T,
+    s: SPInputTerm,
+    p: SPInputTerm,
+    o: OTerm,
+    res?: Set<string>
+) => Set<string>;
+
+/**
+ * Similar to {@link ArrayQueryFn}, but only collects and returns a set of
+ * indices of matching objects.
+ */
+export type ArrayKeyQueryFn<T extends QueryObj[]> = (
+    src: T,
+    p: SPInputTerm,
+    o: OTerm,
+    res?: Set<number>
+) => Set<number>;
+
+/**
+ * Conditional return type for {@link defQuery}.
+ */
+export type QueryFn<T extends QueryObj | QueryObj[]> = T extends QueryObj[]
+    ? ArrayQueryFn<T>
+    : ObjQueryFn<T>;
+
+/**
+ * Conditional return type for {@link defKeyQuery}.
+ */
+export type KeyQueryFn<T extends QueryObj | QueryObj[]> = T extends QueryObj[]
+    ? ArrayKeyQueryFn<T>
+    : ObjKeyQueryFn<T>;
+
+/**
+ * Query behavior options.
+ */
 export interface QueryOpts {
     /**
      * If false, an entire object is included in the solution as soon as any of
@@ -144,3 +192,8 @@ export interface QueryOpts {
      */
     equiv: Predicate2<any>;
 }
+
+/**
+ * Subset of {@link QueryOpts} applicable to {@link defKeyQuery}.
+ */
+export interface KeyQueryOpts extends Pick<QueryOpts, "cwise" | "equiv"> {}

packages/oquery/src/query.ts

@@ -1,15 +1,19 @@
-import type { Predicate } from "@thi.ng/api";
+import type { Fn2, Predicate } from "@thi.ng/api";
 import { isArray, isFunction, isSet } from "@thi.ng/checks";
 import { defmulti } from "@thi.ng/defmulti";
 import { equiv } from "@thi.ng/equiv";
 import type {
     FTerm,
+    KeyQueryFn,
+    KeyQueryOpts,
     OTerm,
     QueryFn,
     QueryImpl,
     QueryImpls,
     QueryObj,
     QueryOpts,
+    QueryType,
+    SPInputTerm,
     SPTerm,
 } from "./api";
 
@@ -212,17 +216,7 @@ const queryO: QueryImpl = (res, db: any, s, p, _, opts) => {
         (opts.partial ? addTriple(res, s, p, val) : collectFull(res, s, sval));
 };
 
-const impl = defmulti<
-    QueryObj,
-    QueryObj,
-    SPTerm,
-    SPTerm,
-    OTerm,
-    QueryOpts,
-    void
->((_, __, s, p, o) => classify(s) + classify(p) + classify(o));
-
-impl.addAll(<QueryImpls>{
+const IMPLS = <QueryImpls>{
     lll: queryLL,
     llf: queryLL,
     lln: queryO,
@@ -307,34 +301,104 @@ impl.addAll(<QueryImpls>{
     nnl: queryNN,
     nnf: queryNN,
     nnn: (res, db) => Object.assign(res, db),
-});
+};
 
-export const defQuery = (opts?: Partial<QueryOpts>): QueryFn => {
-    opts = { partial: false, cwise: true, equiv, ...opts };
-    return (src: any, ...args: any[]) => {
+/**
+ * Query function implementation, dispatches to one of the 27 optimized
+ * functions based on given query pattern.
+ *
+ * @internal
+ */
+const impl = defmulti<
+    QueryObj,
+    QueryObj,
+    SPTerm,
+    SPTerm,
+    OTerm,
+    QueryOpts,
+    void
+>((_, __, s, p, o) => classify(s) + classify(p) + classify(o));
+impl.addAll(IMPLS);
+
+const objQuery = (src: QueryObj[], opts: QueryOpts, args: any[]) => {
+    let [s, p, o, out] = <[SPInputTerm, SPInputTerm, OTerm, QueryObj?]>args;
+    out = out || {};
+    impl(out, src, coerceStr(s), coerceStr(p), coerce(o), <QueryOpts>opts);
+    return out;
+};
+
+const arrayQuery = (
+    src: QueryObj[],
+    opts: QueryOpts,
+    p: SPInputTerm,
+    o: OTerm,
+    collect: Fn2<any, number, void>
+) => {
+    const $p = coerceStr(p);
+    const $o = coerce(o);
+    // pre-select implementation to avoid dynamic dispatch
+    const impl = IMPLS[<QueryType>("n" + classify($p) + classify($o))];
+    for (let i = 0, n = src.length; i < n; i++) {
+        const res: QueryObj = {};
+        impl(res, { _: src[i] }, null, $p, $o, opts);
+        res._ && collect(res._, i);
+    }
+};
+
+/**
+ * Generic Higher-order function to return an actual query function based on
+ * given behavior options.
+ *
+ * @remarks
+ * @see {@link QueryOpts}
+ * @see {@link ObjQueryFn}
+ * @see {@link ArrayQueryFn}
+ * @see {@link defKeyQuery}
+ *
+ * @param opts
+ */
+export const defQuery = <T extends QueryObj | QueryObj[] = QueryObj>(
+    opts?: Partial<QueryOpts>
+): QueryFn<T> => {
+    const $opts: QueryOpts = { partial: false, cwise: true, equiv, ...opts };
+    return <QueryFn<T>>((src: any, ...args: any[]): any => {
         if (isArray(src)) {
-            let [p, o, res] = args;
-            res = res || [];
-            p = coerceStr(p);
-            o = coerce(o);
-            for (let i = 0, n = src.length; i < n; i++) {
-                const curr: QueryObj = {};
-                impl(curr, { _: src[i] }, null, p, o, <QueryOpts>opts);
-                curr._ && res.push(curr._);
-            }
-            return res;
+            const out: QueryObj[] = args[2] || [];
+            arrayQuery(src, $opts, args[0], args[1], (x) => out.push(x));
+            return out;
+        } else {
+            return objQuery(src, $opts, args);
+        }
+    });
+};
+
+/**
+ * Generic Higher-order function to return an actual query function based on
+ * given behavior options. Unlike {@link defQuery}, key query functions only
+ * return sets of keys (or indices) of matching objects.
+ *
+ * @remarks
+ * @see {@link KeyQueryOpts}
+ * @see {@link ObjKeyQueryFn}
+ * @see {@link ArrayKeyQueryFn}
+ *
+ * @param opts
+ */
+export const defKeyQuery = <T extends QueryObj | QueryObj[] = QueryObj>(
+    opts?: Partial<KeyQueryOpts>
+) => {
+    const $opts: QueryOpts = { partial: false, cwise: true, equiv, ...opts };
+    return <KeyQueryFn<T>>((src: any, ...args: any[]): any => {
+        if (isArray(src)) {
+            const out = args[2] || new Set<number>();
+            arrayQuery(src, $opts, args[0], args[1], (_, i) => out.add(i));
+            return out;
         } else {
-            let [s, p, o, res] = args;
-            res = res || {};
-            impl(
-                res,
-                src,
-                coerceStr(s),
-                coerceStr(p),
-                coerce(o),
-                <QueryOpts>opts
-            );
-            return res;
+            const res = objQuery(src, $opts, args.slice(0, 3));
+            const out = args[3];
+            if (!out) return new Set<string>(Object.keys(res));
+            for (let k in res) out.add(k);
+            return out;
         }
-    };
+    });
 };

packages/oquery/test/index.ts

@@ -1,6 +1,6 @@
 import { isNumber } from "@thi.ng/checks";
 import * as assert from "assert";
-import { defQuery, QueryType, SPInputTerm, OTerm } from "../src";
+import { defKeyQuery, defQuery, OTerm, QueryType, SPInputTerm } from "../src";
 
 const DB = {
     alice: {
@@ -20,6 +20,8 @@ const DB = {
     },
 };
 
+const DB_A: any[] = [{ id: 1 }, { id: 11, name: "b" }, { name: "c" }];
+
 describe("oquery", () => {
     it("all patterns", () => {
         const { alice, bob, charlie } = DB;
@@ -333,7 +335,6 @@ describe("oquery", () => {
     });
 
     it("arrays", () => {
-        const db = [{ id: 1 }, { id: 11, name: "b" }, { name: "c" }];
         const tests: Record<
             "ff" | "fl" | "fn" | "lf" | "ll" | "ln" | "nf" | "nl" | "nn",
             [SPInputTerm, OTerm, any]
@@ -346,14 +347,14 @@ describe("oquery", () => {
             ln: ["id", null, [{ id: 1 }, { id: 11 }]],
             nf: [null, isNumber, [{ id: 1 }, { id: 11 }]],
             nl: [null, 11, [{ id: 11 }]],
-            nn: [null, null, [...db]],
+            nn: [null, null, [...DB_A]],
         };
 
-        const query = defQuery({ partial: true });
+        const query = defQuery<any[]>({ partial: true });
         for (let id in tests) {
             const t = tests[<keyof typeof tests>id];
             if (t) {
-                const res = query(db, t[0], t[1]);
+                const res = query(DB_A, t[0], t[1]);
                 assert.deepStrictEqual(
                     res,
                     t[2],
@@ -362,4 +363,25 @@ describe("oquery", () => {
             }
         }
     });
+
+    it("key query", () => {
+        assert.deepStrictEqual(
+            defKeyQuery()(DB, null, "type", "person"),
+            new Set(["alice", "bob"])
+        );
+        const res1 = new Set(["xxx"]);
+        assert.deepStrictEqual(
+            defKeyQuery()(DB, null, "type", "person", res1),
+            new Set(["alice", "bob", "xxx"])
+        );
+        assert.deepStrictEqual(
+            defKeyQuery<any[]>()(DB_A, "name", null),
+            new Set([1, 2])
+        );
+        const res2 = new Set([-1]);
+        assert.deepStrictEqual(
+            defKeyQuery<any[]>()(DB_A, "name", null, res2),
+            new Set([1, 2, -1])
+        );
+    });
 });