@mason-data/cli 0.8.0 → 0.9.0

package/dist/src/duck/libc-bound.js

@@ -0,0 +1,341 @@
+import { functionsCalled } from "./pgtree.js";
+import { BARRIER } from "./row-filter.js";
+/** the run's length in characters: the longest shared prefix a column's values can have and still be told apart by
+ *  their bounds (a URL's `https://example.com/item/` is 25), measured (docs/DESIGN.md "The libc key for the candidates
+ *  alone"); a longer one costs each row's bounds more and saves little */
+export const BOUND_CHARS = 48;
+/** the characters a run may hold: printable ASCII (no control character, so `E`, `L` and `H` are none of them) */
+const FIRST = 0x20, LAST = 0x7e;
+/** `L` below `E` below every character a run writes, below `H` */
+const L = 1, E = 2, H = 0x7f;
+const S = (sval) => ({ String: { sval } });
+const str = (s) => ({ A_Const: { sval: { sval: s } } });
+const int = (n) => ({ A_Const: { ival: n === 0 ? {} : { ival: n } } });
+const call = (name, args) => ({ FuncCall: { funcname: [S(name)], args, funcformat: "COERCE_EXPLICIT_CALL" } });
+const op = (name, l, r) => ({ A_Expr: { kind: "AEXPR_OP", name: [S(name)], lexpr: l, rexpr: r } });
+/** an AND of `args`, one AND's own arguments in place, as Postgres's grammar reads `(a AND b) AND c` back */
+const and = (args) => {
+    const flat = args.flatMap((a) => isNodeOf(a, "BoolExpr") && bodyOf(a).boolop === "AND_EXPR" ? bodyOf(a).args : [a]);
+    return flat.length === 1 ? flat[0] : { BoolExpr: { boolop: "AND_EXPR", args: flat } };
+};
+const nullTest = (arg, type) => ({ NullTest: { arg, nulltesttype: type } });
+const chr = (code) => call("chr", [int(code)]);
+/** a character class of code points, as RE2 reads it: each as `\x{…}`, runs as ranges */
+function charClass(points) {
+    const sorted = [...points].sort((a, b) => a - b);
+    let out = "";
+    for (let i = 0; i < sorted.length;) {
+        let j = i;
+        while (j + 1 < sorted.length && sorted[j + 1] === sorted[j] + 1)
+            j++;
+        const hex = (g) => `\\x{${g.toString(16)}}`;
+        out += i === j ? hex(sorted[i]) : `${hex(sorted[i])}-${hex(sorted[j])}`;
+        i = j + 1;
+    }
+    return `[${out}]`;
+}
+/** the engine's `lower()` of a printable ASCII character (tests/duck-libc-order.test.ts runs the engine's own) */
+const folded = (g) => g >= 0x41 && g <= 0x5a ? g + 32 : g;
+const BOUNDS = new WeakMap();
+/**
+ * The order's bounds, as the module's header says, or null where it has none to give: an order not followed (every
+ * value refused, `libc-locale-unread`), byte order, no character mapped, or mapped characters whose weights their
+ * folded bytes do not order.
+ */
+export function boundOf(o) {
+    const known = BOUNDS.get(o);
+    if (known !== undefined)
+        return known;
+    let bound = null;
+    if (!o.unread && !o.bytes) {
+        // a character whose first level is its own: covered, listed, no sequence of weights
+        const mapped = new Set();
+        for (let g = FIRST; g <= LAST; g++)
+            if (o.uncovered[g] === 0 && o.unlisted[g] === 0 && !o.sequences.has(g))
+                mapped.add(g);
+        // a contraction a mapped character could complete makes its first character's weight depend on what follows: not mapped
+        for (const g of [...mapped])
+            if ((o.contractions.get(g) ?? []).some((c) => c.rest.length > 0 && mapped.has(c.rest[0])))
+                mapped.delete(g);
+        // (the rest begin a contraction only before an unmapped character: known where a mapped one follows, or the text ends)
+        const plain = [...mapped].filter((g) => !o.contractions.has(g));
+        const weighted = [...mapped].filter((g) => o.l1[g] > 0), dropped = [...mapped].filter((g) => o.l1[g] === 0);
+        // each weighted character folds to a mapped one of its weight, and the folded ones' bytes order as their weights,
+        // strictly: one weight, one character written
+        const written = [...new Set(weighted.map(folded))].sort((a, b) => a - b);
+        const ordered = weighted.every((g) => mapped.has(folded(g)) && o.l1[folded(g)] === o.l1[g])
+            && written.every((g, i) => i === 0 || o.l1[written[i - 1]] < o.l1[g]);
+        if (plain.length > 0 && written.length > 0 && ordered) {
+            bound = {
+                // the longest run of mapped characters ending in one whose weight is known there; and the whole text mapped
+                run: `^${charClass([...mapped])}*${charClass(plain)}`,
+                whole: `${charClass([...mapped])}*`,
+                dropped: dropped.length === 0 ? null : charClass(dropped),
+            };
+        }
+    }
+    BOUNDS.set(o, bound);
+    return bound;
+}
+/** the run's weights of `x`, ending in `E` where x ended in the run: x of at most `BOUND_CHARS` bytes, each a mapped
+ *  character (a mapped one is one byte), is whole in it; otherwise the run is the longest prefix of x's first
+ *  `BOUND_CHARS + 1` characters ending in a plain one. No mark is ever read from x */
+function weightsOf(b, x) {
+    const head = () => call("left", [structuredClone(x), int(BOUND_CHARS + 1)]);
+    const whole = and([op("<=", call("strlen", [structuredClone(x)]), int(BOUND_CHARS)), call("regexp_full_match", [head(), str(b.whole)])]);
+    const run = call("lower", [{ CaseExpr: { args: [{ CaseWhen: { expr: whole, result: op("||", structuredClone(x), chr(E)) } }],
+                defresult: call("regexp_extract", [head(), str(b.run)]) } }]);
+    return b.dropped === null ? run : call("regexp_replace", [run, str(b.dropped), str(""), str("g")]);
+}
+/** `lo(x)` and `hi(x)` (the module's header), each its own copy of `x` */
+export const lowBound = (b, x) => op("||", weightsOf(b, structuredClone(x)), chr(L));
+export const highBound = (b, x) => op("||", weightsOf(b, structuredClone(x)), chr(H));
+// ------------------------------------------------------------------------------------------------ the statement
+const isObject = (v) => typeof v === "object" && v !== null && !Array.isArray(v);
+const kindOf = (n) => isObject(n) ? Object.keys(n)[0] ?? null : null;
+const bodyOf = (n) => n[kindOf(n)];
+const isNodeOf = (n, kind) => kindOf(n) === kind && isObject(bodyOf(n));
+const nameOf = (list) => Array.isArray(list) ? list.map((x) => String(bodyOf(x).sval ?? "")).join(".") : "";
+const present = (v) => v !== undefined && v !== null && !(Array.isArray(v) && v.length === 0);
+/** an integer constant, or null (a 0 is `ival: {}`) */
+const constInt = (n) => {
+    if (!isNodeOf(n, "A_Const"))
+        return null;
+    const v = bodyOf(n).ival;
+    return isObject(v) ? Number(v.ival ?? 0) : null;
+};
+/** a libc key the rewrite wrote (`__mason_libc_key(<session>, x, '<collate>')`): x and the collcollate, or null */
+function libcKeyOf(n) {
+    if (!isNodeOf(n, "FuncCall") || nameOf(bodyOf(n).funcname) !== "__mason_libc_key")
+        return null;
+    const args = bodyOf(n).args;
+    const collate = args?.at(-1), x = args?.at(-2);
+    if (x === undefined || !isNodeOf(collate, "A_Const") || !isObject(bodyOf(collate).sval))
+        return null;
+    return { x, collate: String(bodyOf(collate).sval.sval ?? "") };
+}
+/** every SELECT body inside `v`, at any depth (a subquery's, a derived table's, a WITH query's) */
+function selectsIn(v, out = []) {
+    if (Array.isArray(v)) {
+        for (const x of v)
+            selectsIn(x, out);
+        return out;
+    }
+    if (!isObject(v))
+        return out;
+    if (isNodeOf(v, "SelectStmt"))
+        out.push(bodyOf(v));
+    for (const x of Object.values(v))
+        selectsIn(x, out);
+    return out;
+}
+/** `v` holds a node of one of `kinds` */
+function holds(v, kinds) {
+    if (Array.isArray(v))
+        return v.some((x) => holds(x, kinds));
+    if (!isObject(v))
+        return false;
+    const k = kindOf(v);
+    if (k !== null && kinds.has(k) && isObject(v[k]))
+        return true;
+    return Object.values(v).some((x) => holds(x, kinds));
+}
+/** a role's barrier (src/duck/row-filter.ts `BARRIER`): a LIMIT every row passes */
+const BARRIER_COUNT = /\d+/.exec(BARRIER)[0];
+const isBarrier = (s) => {
+    const c = s.limitCount;
+    return isNodeOf(c, "A_Const") && String(bodyOf(c).fval?.fval ?? bodyOf(c).ival?.ival ?? "") === BARRIER_COUNT && (s.limitOffset === undefined || constInt(s.limitOffset) === 0);
+};
+/** what is cloned reads the same rows twice: no LIMIT or OFFSET but a barrier, no DISTINCT ON, no sample, in any SELECT
+ *  inside it (nothing volatile, no window and no set-returning list in the statement: `statementReads`) */
+function readsAlike(v) {
+    if (holds(v, new Set(["RangeTableSample"])))
+        return false;
+    return selectsIn(v).every((s) => ((s.limitCount === undefined && s.limitOffset === undefined) || isBarrier(s))
+        && !(Array.isArray(s.distinctClause) && s.distinctClause.some((d) => kindOf(d) !== null)));
+}
+/** the WITH queries of the statement, by name (any depth: a name read in FROM may be any of them) */
+function withQueries(tree) {
+    const out = new Map();
+    for (const s of selectsIn(tree)) {
+        for (const c of (s.withClause?.ctes ?? [])) {
+            if (!isNodeOf(c, "CommonTableExpr"))
+                continue;
+            const b = bodyOf(c);
+            out.set(String(b.ctename), [...out.get(String(b.ctename)) ?? [], b.ctequery]);
+        }
+    }
+    return out;
+}
+/** the one FROM item of a SELECT that may be read twice: a table, a derived table, a WITH query, as `readsAlike` says */
+function oneItem(s, ctes) {
+    const from = s.fromClause;
+    if (from?.length !== 1)
+        return null;
+    const item = from[0];
+    if (isNodeOf(item, "RangeVar")) {
+        const b = bodyOf(item);
+        const named = b.schemaname === undefined ? ctes.get(String(b.relname)) : undefined;
+        return named === undefined || named.every(readsAlike) ? item : null;
+    }
+    if (isNodeOf(item, "RangeSubselect") && bodyOf(item).lateral !== true && readsAlike(item))
+        return item;
+    return null;
+}
+function statementReads(tree) {
+    let aggregates = 0, windowsOrSets = false;
+    const visit = (v) => {
+        if (Array.isArray(v)) {
+            for (const x of v)
+                visit(x);
+            return;
+        }
+        if (v === null || v === undefined || typeof v !== "object")
+            return;
+        const n = v;
+        if (n._node === "QUERY") {
+            if (n.hasAggs !== "false")
+                aggregates++;
+            if (n.hasWindowFuncs !== "false" || n.hasTargetSRFs !== "false")
+                windowsOrSets = true;
+        }
+        for (const [k, x] of Object.entries(n))
+            if (k !== "_node")
+                visit(x);
+    };
+    visit(tree);
+    return { aggregates, windowsOrSets };
+}
+/** a SELECT that is no set operation, lock or INTO, and whose list holds no window */
+const plainSelect = (s) => (s.op === undefined || s.op === "SETOP_NONE") && !present(s.windowClause) && !present(s.lockingClause)
+    && !present(s.intoClause) && !present(s.valuesLists) && !windowed(s.targetList);
+/** a call OVER a window in `v`, outside its subqueries */
+function windowed(v) {
+    if (Array.isArray(v))
+        return v.some(windowed);
+    if (!isObject(v) || isNodeOf(v, "SubLink"))
+        return false;
+    if (isNodeOf(v, "FuncCall") && bodyOf(v).over !== undefined)
+        return true;
+    return Object.values(v).some(windowed);
+}
+/** the threshold's subquery over the SELECT's own rows: `SELECT <aggregate> FROM (SELECT <value> AS __mason_bound FROM
+ *  <its FROM item> WHERE <its WHERE> AND x IS NOT NULL [ORDER BY 1 [DESC] LIMIT n]) AS __mason_bounds`; a top N's is its
+ *  `everyRow` where fewer than n rows have a value (`CASE WHEN count(*) >= n THEN … ELSE <everyRow> END`) */
+function threshold(s, item, x, value, aggregate, top) {
+    const where = and([...(s.whereClause !== undefined ? [structuredClone(s.whereClause)] : []), nullTest(structuredClone(x), "IS_NOT_NULL")]);
+    const inner = {
+        targetList: [{ ResTarget: { name: "__mason_bound", val: value } }],
+        fromClause: [structuredClone(item)], whereClause: where,
+        ...(top === null ? {} : {
+            sortClause: [{ SortBy: { node: int(1), sortby_dir: top.desc ? "SORTBY_DESC" : "SORTBY_DEFAULT", sortby_nulls: "SORTBY_NULLS_DEFAULT" } }],
+            limitCount: int(top.n), limitOption: "LIMIT_OPTION_COUNT",
+        }),
+        op: "SETOP_NONE",
+    };
+    const outer = {
+        targetList: [{ ResTarget: { val: ((v) => top === null ? v : { CaseExpr: {
+                            args: [{ CaseWhen: { expr: op(">=", { FuncCall: { funcname: [S("count")], agg_star: true, funcformat: "COERCE_EXPLICIT_CALL" } }, int(top.n)), result: v } }], defresult: top.everyRow
+                        } })(call(aggregate, [{ ColumnRef: { fields: [S("__mason_bounds"), S("__mason_bound")] } }])) } }],
+        fromClause: [{ RangeSubselect: { subquery: { SelectStmt: inner }, alias: { aliasname: "__mason_bounds" } } }],
+        op: "SETOP_NONE",
+    };
+    return { SubLink: { subLinkType: "EXPR_SUBLINK", subselect: { SelectStmt: outer } } };
+}
+/** a role's walled reads in `v`: each `SELECT * FROM (… LIMIT <barrier>) AS <wall>` */
+const wallsIn = (v, alias) => selectsIn(v).filter((s) => {
+    const from = s.fromClause;
+    return from?.length === 1 && isNodeOf(from[0], "RangeSubselect") && bodyOf(from[0]).alias?.aliasname === alias;
+}).length;
+/**
+ * The statement with each top N and each ungrouped min / max under a libc order keyed for its candidates alone (the
+ * module's header), or null where none is. `tree` is the rewrite's DuckDB tree (after src/duck/late.ts), `analysis`
+ * Postgres's of the statement it was written from.
+ */
+export async function libcBounded(tree, analysis, context) {
+    const tops = [], picks = [];
+    const ctes = withQueries(tree);
+    const boundFor = (collate) => { const o = context.order(collate); return o === undefined ? null : boundOf(o); };
+    for (const s of selectsIn(tree)) {
+        if (!plainSelect(s) || oneItem(s, ctes) === null || (s.whereClause !== undefined && !readsAlike(s.whereClause)))
+            continue;
+        // A TOP N: the first key a libc key, a constant LIMIT and OFFSET, no DISTINCT or GROUP BY
+        const sorts = (s.sortClause ?? []);
+        const first = sorts[0] !== undefined && isNodeOf(sorts[0], "SortBy") ? bodyOf(sorts[0]) : null;
+        const key = first === null || first.useOp !== undefined ? null : libcKeyOf(first.node);
+        const count = constInt(s.limitCount), offset = s.limitOffset === undefined ? 0 : constInt(s.limitOffset);
+        if (key !== null && count !== null && count > 0 && offset !== null && offset >= 0 && count + offset <= 2_147_483_647 && (s.limitOption === undefined || s.limitOption === "LIMIT_OPTION_COUNT" || s.limitOption === "LIMIT_OPTION_DEFAULT")
+            && !present(s.distinctClause) && !present(s.groupClause) && !present(s.havingClause)) {
+            const bound = boundFor(key.collate);
+            // (where NULL sorts: the engine's default_null_order is Postgres's, last ascending and first descending)
+            const desc = first.sortby_dir === "SORTBY_DESC";
+            const nullsFirst = first.sortby_nulls === "SORTBY_NULLS_FIRST" || (desc && first.sortby_nulls !== "SORTBY_NULLS_LAST");
+            if (bound !== null)
+                tops.push({ s, x: key.x, bound, desc, nullsFirst, n: count + offset });
+            continue;
+        }
+        // A MIN OR A MAX: arg_min / arg_max over the libc key's CASE (src/duck/rewrite.ts textMinMax), ungrouped, no FILTER
+        if (present(s.groupClause) || present(s.distinctClause) || present(s.sortClause) || s.limitCount !== undefined || s.limitOffset !== undefined)
+            continue;
+        const each = [];
+        const visit = (v) => {
+            if (Array.isArray(v)) {
+                v.forEach(visit);
+                return;
+            }
+            if (!isObject(v) || isNodeOf(v, "SubLink"))
+                return;
+            if (isNodeOf(v, "FuncCall")) {
+                const f = bodyOf(v), name = nameOf(f.funcname), args = f.args;
+                const arm = args?.[1] !== undefined && isNodeOf(args[1], "CaseExpr") ? bodyOf(args[1]) : null;
+                const when = arm !== null && arm.args.length === 1 && arm.defresult === undefined ? bodyOf(arm.args[0]) : null;
+                const result = when !== null && isNodeOf(when.result, "RowExpr") ? bodyOf(when.result).args[0] : undefined;
+                const k = libcKeyOf(result);
+                if ((name === "arg_min" || name === "arg_max") && k !== null && f.over === undefined && f.agg_filter === undefined && f.agg_distinct !== true) {
+                    const bound = boundFor(k.collate);
+                    if (bound !== null)
+                        each.push({ case: when, x: k.x, bound, min: name === "arg_min" });
+                    return;
+                }
+            }
+            Object.values(v).forEach(visit);
+        };
+        visit(s.targetList);
+        if (each.length > 0)
+            picks.push({ s, each });
+    }
+    if (tops.length === 0 && picks.length === 0)
+        return null;
+    // the statement reads its rows alike twice (the module's header): Postgres's analysis and pg_proc decide it
+    const reads = statementReads(analysis);
+    if (reads.windowsOrSets)
+        return null;
+    if ((await context.volatile([...functionsCalled(analysis).oids].filter((f) => f > 0).sort((a, b) => a - b))).size > 0)
+        return null;
+    // a top N in a statement of no aggregate; a min or max in the one query that aggregates
+    const chosen = { tops: reads.aggregates === 0 ? tops : [], picks: reads.aggregates === 1 ? picks : [] };
+    if (chosen.tops.length === 0 && chosen.picks.length === 0)
+        return null;
+    const counted = (n) => { if (context.wall !== null)
+        context.wall.added += wallsIn(n, context.wall.alias); return n; };
+    for (const t of chosen.tops) {
+        const item = oneItem(t.s, ctes);
+        // ASC: the n-th smallest hi, rows of lo <= it; DESC: the n-th largest lo, rows of hi >= it. Fewer than n values,
+        // a threshold every bound passes: `H` above every lo, '' below every hi. A NULL row's bound is one every threshold
+        // passes where NULL sorts first, and else the one only that threshold passes
+        const [everyRow, beforeEvery] = t.desc ? [str(""), chr(H)] : [chr(H), str("")];
+        const limit = counted(threshold(t.s, item, t.x, (t.desc ? lowBound : highBound)(t.bound, t.x), t.desc ? "min" : "max", { desc: t.desc, n: t.n, everyRow }));
+        const own = { CoalesceExpr: { args: [(t.desc ? highBound : lowBound)(t.bound, t.x), t.nullsFirst ? beforeEvery : structuredClone(everyRow)] } };
+        const kept = op(t.desc ? ">=" : "<=", own, limit);
+        t.s.whereClause = t.s.whereClause === undefined ? kept : and([t.s.whereClause, kept]);
+    }
+    for (const { s, each } of chosen.picks) {
+        const item = oneItem(s, ctes);
+        for (const p of each) {
+            // min: the least hi, rows of lo <= it; max: the greatest lo, rows of hi >= it
+            const limit = counted(threshold(s, item, p.x, (p.min ? highBound : lowBound)(p.bound, p.x), p.min ? "min" : "max", null));
+            p.case.expr = and([p.case.expr, op(p.min ? "<=" : ">=", (p.min ? lowBound : highBound)(p.bound, p.x), limit)]);
+        }
+    }
+    return tree;
+}