From f58c321c074ab6ac92e5d332d883e31e5ce9daba Mon Sep 17 00:00:00 2001 From: Chris Tate Date: Wed, 21 Oct 2026 08:31:35 -0500 Subject: [PATCH] Sort by several keys without calling a chained comparator A comparator that chains key differences with `||`, such as `(b.score ?? 0) - (a.score ?? 0) || a.id - b.id`, ran through the comparator sort: TimSort called the closure once per comparison. Each term orders one number field, ascending or descending, and a term that is 0 or NaN falls through to the next, so the chain orders elements by their keys with the first deciding. With no key NaN or -0 that order is consistent, and a stable sort by the keys gives the comparator sort's permutation. The recognizer now accepts such chains of up to eight terms over number fields of records and class instances, and a key may read an optional number field with a number default (`x.k ?? 0`), also as a single key. The runtime reads and checks every key of every element before it moves a slot, then runs stable keyed passes from the last key to the first. A NaN or -0 key, or a chain with any other term, keeps the comparator sort. --- .../lowering/containers/array-callbacks.ts | 40 +++-- .../src/frontend/lowering/containers/bytes.ts | 2 +- .../lowering/lower-array-sort.test.ts | 147 ++++++++++++++-- .../src/frontend/lowering/lower-array-sort.ts | 158 ++++++++++++++---- packages/compiler/src/ir/validate.test.ts | 16 +- packages/compiler/src/ir/validate.ts | 31 +++- .../corpus/array-sort-keyed-chains.ts.json | 6 + packages/runtime/src/scr_array.c | 80 ++++++++- packages/runtime/src/scr_array.h | 6 +- packages/runtime/test/test_array.c | 89 ++++++++++ tests/corpus/array-sort-keyed-chains.ts | 149 +++++++++++++++++ .../array-sort-keyed-chains.ts.json | 1 + 12 files changed, 655 insertions(+), 70 deletions(-) create mode 100644 packages/compiler/test/ts7/baselines/order-parity/tests/corpus/array-sort-keyed-chains.ts.json create mode 100644 tests/corpus/array-sort-keyed-chains.ts create mode 100644 tests/harness/corpus-costs/array-sort-keyed-chains.ts.json diff --git a/packages/compiler/src/frontend/lowering/containers/array-callbacks.ts b/packages/compiler/src/frontend/lowering/containers/array-callbacks.ts index 2cec8bc24..eb99592bb 100644 --- a/packages/compiler/src/frontend/lowering/containers/array-callbacks.ts +++ b/packages/compiler/src/frontend/lowering/containers/array-callbacks.ts @@ -28,9 +28,12 @@ import { import { isJsSourceFile, locOf } from "../../program.js"; import { type SortKeyOrder, + type SortKeyRead, buildArraySortFn, buildSortKeyFn, closureKeyOrder, + sortKeysId, + sortKeysOrder, } from "../lower-array-sort.js"; import { arrayIndexPresent, @@ -2112,8 +2115,10 @@ export function lowerArraySortCall( elem, elem.kind === "union" ? lowerer.unions.get(elem.unionId)?.arms : undefined, ); - const keyed = native ? recognizedSortKey(lowerer, fnArg, elem) : null; - const key = `${method}:${typeKey(elem)}:${arity}${native ? ":native" : ""}${keyed ? `:keyed${keyed.order}${keyed.read ? ":field" : ""}` : ""}`; + const keyed = native ? recognizedSortKeys(lowerer, fnArg, elem) : null; + const reads = keyed?.every((k) => k.read !== null) ? keyed.map((k) => k.read!) : null; + const order = keyed ? sortKeysOrder(keyed) : 0; + const key = `${method}:${typeKey(elem)}:${arity}${native ? ":native" : ""}${keyed ? `:keyed${order}${reads ? `:field${reads.length}` : ""}` : ""}`; let helper = lowerer.arrHofHelpers.get(key); if (!helper) { helper = `%arr.${method}.${lowerer.arrHofHelpers.size}`; @@ -2128,41 +2133,46 @@ export function lowerArraySortCall( undefinedTag >= 0 ? undefinedTag : null, loc, native, - keyed && { order: keyed.order, key: keyed.read !== null }, + keyed && { order, key: reads !== null, keys: keyed.length }, ), ); } const args = [receiver, fnArg]; - if (keyed?.read) args.push(sortKeyReader(lowerer, elem, keyed.read, loc)); + if (reads) args.push(sortKeyReader(lowerer, elem, reads, loc)); return { kind: "call", callee: helper, args, type: arrT, loc }; } -/** The key order of a capture-free comparator closure the runtime can sort - * by without calling it: over number elements, or over one number field - * of record or class elements. */ -function recognizedSortKey(lowerer: Lowerer, fnArg: IrExpr, elem: IrType): SortKeyOrder | null { +/** The keys of a capture-free comparator closure the runtime can sort by + * without calling it: number elements themselves, or number fields of + * record or class elements (one or several). */ +function recognizedSortKeys( + lowerer: Lowerer, + fnArg: IrExpr, + elem: IrType, +): readonly SortKeyOrder[] | null { const keyed = closureKeyOrder(lowerer.liftedFns, fnArg); if (!keyed) return null; - if (keyed.read === null) return elem.kind === "f64" ? keyed : null; + if (keyed[0]!.read === null) return elem.kind === "f64" ? keyed : null; return elem.kind === "record" || elem.kind === "object" ? keyed : null; } -/** A zero-capture closure over the shared `(x) => x.field` key reader. */ +/** A zero-capture closure over the shared key reader of these reads + * (buildSortKeyFn): `(x) => x.field`, or `(x, j) => key j` for several. */ function sortKeyReader( lowerer: Lowerer, elem: IrType, - read: NonNullable, + reads: readonly SortKeyRead[], loc: SrcLoc, ): IrExpr { - const owner = read.kind === "recordGet" ? read.shapeId : read.className; - const cacheKey = `sortKey:${typeKey(elem)}:${read.kind}:${owner}:${read.field}`; + const cacheKey = `sortKey:${typeKey(elem)}:${sortKeysId(reads.map((read) => ({ order: 1, read })))}`; let reader = lowerer.arrHofHelpers.get(cacheKey); if (!reader) { reader = `%arr.sortKey.${lowerer.arrHofHelpers.size}`; lowerer.arrHofHelpers.set(cacheKey, reader); - lowerer.liftedFns.push(buildSortKeyFn(reader, elem, read, loc)); + lowerer.liftedFns.push(buildSortKeyFn(reader, elem, reads, loc)); } - return { kind: "closure", fnName: reader, captures: [], type: funcOf([elem], F64), loc }; + const params = reads.length > 1 ? [elem, F64] : [elem]; + return { kind: "closure", fnName: reader, captures: [], type: funcOf(params, F64), loc }; } /** `d.filter(pred)` on a dyn receiver: the receiver must BE a dyn array diff --git a/packages/compiler/src/frontend/lowering/containers/bytes.ts b/packages/compiler/src/frontend/lowering/containers/bytes.ts index 79154e67f..8d1d17821 100644 --- a/packages/compiler/src/frontend/lowering/containers/bytes.ts +++ b/packages/compiler/src/frontend/lowering/containers/bytes.ts @@ -146,7 +146,7 @@ function lowerBytesSortCall( } const arity = fnArg.type.params.length; const keyed = closureKeyOrder(lowerer.liftedFns, fnArg); - const order = keyed !== null && keyed.read === null ? keyed.order : null; + const order = keyed?.length === 1 && keyed[0]!.read === null ? keyed[0]!.order : null; const helper = helperFor(`${arity}${order === null ? "" : `:keyed${order}`}`, arity, true, order); return { kind: "call", diff --git a/packages/compiler/src/frontend/lowering/lower-array-sort.test.ts b/packages/compiler/src/frontend/lowering/lower-array-sort.test.ts index be73e45d9..11b17acd4 100644 --- a/packages/compiler/src/frontend/lowering/lower-array-sort.test.ts +++ b/packages/compiler/src/frontend/lowering/lower-array-sort.test.ts @@ -15,9 +15,11 @@ import { buildBytesSortFn, buildSortKeyFn, recognizeKeyComparator, + sortKeysOrder, } from "./lower-array-sort.js"; const loc = { file: "sort.ts", start: 0, end: 0 }; +const OPTIONAL: IrType = { kind: "union", unionId: "u0" }; const ROW: IrType = { kind: "record", shapeId: "r0" }; const a = (type: IrType = F64): IrExpr => varRef("a.0", type, loc); const b = (type: IrType = F64): IrExpr => varRef("b.0", type, loc); @@ -86,9 +88,40 @@ test("sort helpers give every IR site its own node", () => { buildArraySortFn("%arr.toSorted.1", F64, 1, true, null, loc), buildArraySortFn("%arr.sort.2", DYN, 2, false, null, loc), buildArraySortFn("%arr.sort.3", union, 2, true, 1, loc), - buildArraySortFn("%arr.sort.6", F64, 2, false, null, loc, true, { order: -1, key: false }), - buildArraySortFn("%arr.toSorted.7", ROW, 2, true, null, loc, true, { order: 1, key: true }), - buildSortKeyFn("%arr.sortKey.8", ROW, field(varRef("a.0", ROW, loc)), loc), + buildArraySortFn("%arr.sort.6", F64, 2, false, null, loc, true, { + order: -1, + key: false, + keys: 1, + }), + buildArraySortFn("%arr.toSorted.7", ROW, 2, true, null, loc, true, { + order: 1, + key: true, + keys: 1, + }), + buildArraySortFn("%arr.sort.11", ROW, 2, false, null, loc, true, { + order: 0x202, + key: true, + keys: 2, + }), + buildSortKeyFn( + "%arr.sortKey.8", + ROW, + [{ field: field(varRef("a.0", ROW, loc)), fallback: null }], + loc, + ), + buildSortKeyFn( + "%arr.sortKey.12", + ROW, + [ + { + field: field(varRef("a.0", ROW, loc), "k", OPTIONAL), + fallback: numLit(0, loc) as IrExpr & { kind: "numLit" }, + }, + { field: field(varRef("a.0", ROW, loc), "j"), fallback: null }, + { field: field(varRef("a.0", ROW, loc), "i"), fallback: null }, + ], + loc, + ), buildBytesSortFn("%bytes.sort.4", { kind: "bytes", elem: "u8" }, 2, true, true, null, loc), buildBytesSortFn("%bytes.sort.5", { kind: "bytes", elem: "u8" }, 0, false, true, null, loc), buildBytesSortFn("%bytes.sort.9", { kind: "bytes", elem: "f64" }, 2, true, false, -1, loc), @@ -98,16 +131,20 @@ test("sort helpers give every IR site its own node", () => { }); test("recognizes subtraction comparators over numbers", () => { - expect(recognizeKeyComparator(comparator(sub(a(), b())))).toEqual({ order: 1, read: null }); - expect(recognizeKeyComparator(comparator(sub(b(), a())))).toEqual({ order: -1, read: null }); + expect(recognizeKeyComparator(comparator(sub(a(), b())))).toEqual([{ order: 1, read: null }]); + expect(recognizeKeyComparator(comparator(sub(b(), a())))).toEqual([{ order: -1, read: null }]); }); test("recognizes subtraction comparators over one field", () => { const ascending = recognizeKeyComparator(comparator(sub(field(a(ROW)), field(b(ROW))), ROW)); - expect(ascending?.order).toBe(1); - expect(ascending?.read).toMatchObject({ kind: "recordGet", field: "k" }); + expect(ascending).toHaveLength(1); + expect(ascending?.[0]?.order).toBe(1); + expect(ascending?.[0]?.read).toMatchObject({ + field: { kind: "recordGet", field: "k" }, + fallback: null, + }); const descending = recognizeKeyComparator(comparator(sub(field(b(ROW)), field(a(ROW))), ROW)); - expect(descending?.order).toBe(-1); + expect(descending?.[0]?.order).toBe(-1); }); test("recognizes three-way ternaries in either test order and direction", () => { @@ -117,12 +154,12 @@ test("recognizes three-way ternaries in either test order and direction", () => numLit(firstSign, loc), pick(cmp(first === "<" ? ">" : "<", x, y), numLit(-firstSign, loc), numLit(0, loc)), ); - expect(recognizeKeyComparator(comparator(three("<", -1, a(), b())))?.order).toBe(1); - expect(recognizeKeyComparator(comparator(three(">", 1, a(), b())))?.order).toBe(1); - expect(recognizeKeyComparator(comparator(three("<", 1, a(), b())))?.order).toBe(-1); - expect(recognizeKeyComparator(comparator(three("<", -1, b(), a())))?.order).toBe(-1); + expect(recognizeKeyComparator(comparator(three("<", -1, a(), b())))?.[0]?.order).toBe(1); + expect(recognizeKeyComparator(comparator(three(">", 1, a(), b())))?.[0]?.order).toBe(1); + expect(recognizeKeyComparator(comparator(three("<", 1, a(), b())))?.[0]?.order).toBe(-1); + expect(recognizeKeyComparator(comparator(three("<", -1, b(), a())))?.[0]?.order).toBe(-1); const fields = three("<", -1, field(a(ROW)), field(b(ROW))); - expect(recognizeKeyComparator(comparator(fields, ROW))?.order).toBe(1); + expect(recognizeKeyComparator(comparator(fields, ROW))?.[0]?.order).toBe(1); }); test("rejects comparators that are not one consistent numeric key", () => { @@ -172,3 +209,87 @@ test("rejects comparators with captures or other bodies", () => { ]; expect(recognizeKeyComparator({ ...comparator(sub(a(), b())), body })).toBeNull(); }); + +const or = (left: IrExpr, right: IrExpr): IrExpr => ({ + kind: "logical", + op: "||", + left, + right, + type: F64, + loc, +}); +const orZero = (read: IrExpr & { kind: "recordGet" }, value = 0): IrExpr => ({ + kind: "nullish", + left: read, + right: numLit(value, loc), + type: F64, + loc, +}); + +test("recognizes an optional field with a number default as a key", () => { + const keys = recognizeKeyComparator( + comparator( + sub(orZero(field(b(ROW), "k", OPTIONAL)), orZero(field(a(ROW), "k", OPTIONAL))), + ROW, + ), + ); + expect(keys).toHaveLength(1); + expect(keys?.[0]?.order).toBe(-1); + expect(keys?.[0]?.read?.fallback).toMatchObject({ kind: "numLit", value: 0 }); +}); + +test("recognizes chains of key terms joined by ||", () => { + const keys = recognizeKeyComparator( + comparator( + or( + or( + sub(orZero(field(b(ROW), "k", OPTIONAL)), orZero(field(a(ROW), "k", OPTIONAL))), + sub(field(a(ROW), "j"), field(b(ROW), "j")), + ), + sub(field(b(ROW), "i"), field(a(ROW), "i")), + ), + ROW, + ), + ); + expect(keys?.map((key) => [key.order, key.read?.field.field])).toEqual([ + [-1, "k"], + [1, "j"], + [-1, "i"], + ]); + // Three keys, the first and third descending. + expect(sortKeysOrder(keys!)).toBe(3 | (1 << 8) | (1 << 10)); +}); + +test("rejects key chains that are not all field keys", () => { + const rejected: IrExpr[] = [ + // The elements themselves cannot be one of several keys. + or(sub(a(), b()), sub(b(), a())), + // Each operand pair must read the same key with the same default. + or(sub(field(a(ROW)), field(b(ROW))), sub(field(a(ROW), "j"), field(b(ROW), "i"))), + or( + sub(orZero(field(a(ROW), "k", OPTIONAL)), orZero(field(b(ROW), "k", OPTIONAL), 1)), + sub(field(a(ROW), "j"), field(b(ROW), "j")), + ), + // `&&` does not fall through on equal keys. + { + kind: "logical", + op: "&&", + left: sub(field(a(ROW)), field(b(ROW))), + right: sub(field(a(ROW), "j"), field(b(ROW), "j")), + type: F64, + loc, + }, + ]; + for (const value of rejected) expect(recognizeKeyComparator(comparator(value, ROW))).toBeNull(); + const nine = Array.from({ length: 9 }, (_, i) => + sub(field(a(ROW), `f${i}`), field(b(ROW), `f${i}`)), + ); + expect( + recognizeKeyComparator( + comparator( + nine.reduce((left, right) => or(left, right)), + ROW, + ), + ), + ).toBeNull(); +}); diff --git a/packages/compiler/src/frontend/lowering/lower-array-sort.ts b/packages/compiler/src/frontend/lowering/lower-array-sort.ts index 82d91a375..5521d30ef 100644 --- a/packages/compiler/src/frontend/lowering/lower-array-sort.ts +++ b/packages/compiler/src/frontend/lowering/lower-array-sort.ts @@ -699,45 +699,87 @@ function timSortBody(b: SortIr, src: SortLocal, n: IrExpr, less: Less): IrStmt[] ]; } -type SortKeyRead = IrExpr & { kind: "recordGet" | "fieldGet" }; +type SortKeyField = IrExpr & { kind: "recordGet" | "fieldGet" }; -/** A comparator that orders by one number: the elements themselves or one - * field of each, ascending or descending. */ +/** A key read over the comparator's first parameter: one number field, or + * an optional number field with a number default (`x.k ?? 0`). */ +export interface SortKeyRead { + field: SortKeyField; + fallback: (IrExpr & { kind: "numLit" }) | null; +} + +/** One key of a comparator: the elements themselves or one field of each, + * ascending or descending. */ export interface SortKeyOrder { /** 1 for ascending, -1 for descending. */ order: 1 | -1; - /** The field read over the comparator's first parameter, or null when - * the elements are the numbers. */ + /** The key read, or null when the elements are the numbers. */ read: SortKeyRead | null; } +/** Most keys a comparator chain is recognized with (the runtime's + * SCR_SORT_KEYS_MAX). */ +const MAX_SORT_KEYS = 8; + /** Recognizes the lifted body of a capture-free numeric comparator: * `(a, b) => a - b`, `b - a`, `a < b ? -1 : a > b ? 1 : 0` (either test * first, either sign), or any of them over one number field of both - * parameters (`a.k - b.k`). Such a comparator runs no user code, cannot - * throw, and is consistent over numbers other than NaN, so the runtime may - * order the elements by their keys without calling it. */ -export function recognizeKeyComparator(fn: IrFunction): SortKeyOrder | null { + * parameters (`a.k - b.k`, or `(a.k ?? 0) - (b.k ?? 0)` over an optional + * field). Such a comparator runs no user code, cannot throw, and is + * consistent over numbers other than NaN, so the runtime may order the + * elements by their keys without calling it. A chain `t0 || t1 || ...` of + * such terms over fields orders by several keys, the first deciding: a term + * that is 0 or NaN falls through to the next. Returns the keys in order. */ +export function recognizeKeyComparator(fn: IrFunction): readonly SortKeyOrder[] | null { if (fn.params.length !== 2 || fn.async || fn.generator || (fn.captures?.length ?? 0) > 0) return null; - const [first, second] = fn.params; const ret = fn.body.length === 1 ? fn.body[0]! : null; if (ret?.kind !== "return" || !ret.value) return null; + const terms: IrExpr[] = []; + const flatten = (e: IrExpr): void => { + if (e.kind === "logical" && e.op === "||" && e.type.kind === "f64") { + flatten(e.left); + flatten(e.right); + } else terms.push(e); + }; + flatten(ret.value); + if (terms.length > MAX_SORT_KEYS) return null; + const keys: SortKeyOrder[] = []; + for (const term of terms) { + const key = recognizeKeyTerm(fn, term); + if (!key) return null; + keys.push(key); + } + // Several keys order elements by their fields. + if (keys.length > 1 && keys.some((key) => key.read === null)) return null; + return keys; +} + +/** One term of a key comparator over one key (recognizeKeyComparator). */ +function recognizeKeyTerm(fn: IrFunction, value: IrExpr): SortKeyOrder | null { + const [first, second] = fn.params; let read: SortKeyRead | null | undefined; // Which parameter an operand reads; every operand must read the same key. const operand = (e: IrExpr): 0 | 1 | null => { if (e.type.kind !== "f64") return null; let ref: IrExpr = e; - let field: SortKeyRead | null = null; - if (e.kind === "recordGet" || e.kind === "fieldGet") { - field = e; + let key: SortKeyRead | null = null; + if ( + e.kind === "nullish" && + (e.left.kind === "recordGet" || e.left.kind === "fieldGet") && + e.right.kind === "numLit" + ) { + key = { field: e.left, fallback: e.right }; + ref = e.left.obj; + } else if (e.kind === "recordGet" || e.kind === "fieldGet") { + key = { field: e, fallback: null }; ref = e.obj; } if (ref.kind !== "varRef") return null; const param = ref.localId === first!.localId ? 0 : ref.localId === second!.localId ? 1 : null; if (param === null) return null; - if (read === undefined) read = field; - else if (!sameKeyRead(read, field)) return null; + if (read === undefined) read = key; + else if (!sameKeyRead(read, key)) return null; return param; }; // `x < y` or `x > y` over both parameters, as the first one's relation. @@ -750,7 +792,6 @@ export function recognizeKeyComparator(fn: IrFunction): SortKeyOrder | null { }; const sign = (e: IrExpr): number | null => e.kind === "numLit" && (e.value === 1 || e.value === -1 || e.value === 0) ? e.value : null; - const value = ret.value; if (value.kind === "bin" && value.op === "-") { const left = operand(value.left), right = operand(value.right); @@ -778,12 +819,12 @@ export function recognizeKeyComparator(fn: IrFunction): SortKeyOrder | null { return null; } -/** The key order of a comparator closure: a zero-capture closure (not a +/** The keys of a comparator closure: a zero-capture closure (not a * signature adapter) over a lifted function recognizeKeyComparator accepts. */ export function closureKeyOrder( liftedFns: readonly IrFunction[], closure: IrExpr, -): SortKeyOrder | null { +): readonly SortKeyOrder[] | null { if (closure.kind !== "closure" || closure.captures.length > 0 || closure.adapts) return null; const fn = liftedFns.find((lifted) => lifted.name === closure.fnName); return fn ? recognizeKeyComparator(fn) : null; @@ -791,29 +832,87 @@ export function closureKeyOrder( function sameKeyRead(left: SortKeyRead | null, right: SortKeyRead | null): boolean { if (left === null || right === null) return left === right; - if (left.kind === "recordGet" && right.kind === "recordGet") - return left.shapeId === right.shapeId && left.field === right.field; - if (left.kind === "fieldGet" && right.kind === "fieldGet") - return left.className === right.className && left.field === right.field; + if (!Object.is(left.fallback?.value, right.fallback?.value)) return false; + const a = left.field, + b = right.field; + if (a.kind === "recordGet" && b.kind === "recordGet") + return a.shapeId === b.shapeId && a.field === b.field; + if (a.kind === "fieldGet" && b.kind === "fieldGet") + return a.className === b.className && a.field === b.field; return false; } -/** The key reader of a recognized field comparator: `(x) => x.k`, the same - * field read as the comparator's. */ +/** The order argument of a keyed sort: 1 or -1 for one key; for several, + * the key count in the low byte and, from bit 8, one bit per descending + * key (the runtime's SCR_SORT_KEYS encoding). */ +export function sortKeysOrder(keys: readonly SortKeyOrder[]): number { + if (keys.length === 1) return keys[0]!.order; + return keys.reduce((code, key, j) => (key.order < 0 ? code | (1 << (8 + j)) : code), keys.length); +} + +/** A cache key that identifies a recognized comparator's key reads. */ +export function sortKeysId(keys: readonly SortKeyOrder[]): string { + return keys + .map((key) => { + if (!key.read) return "self"; + const field = key.read.field; + const owner = field.kind === "recordGet" ? field.shapeId : field.className; + const fallback = key.read.fallback + ? `??${Object.is(key.read.fallback.value, -0) ? "-0" : key.read.fallback.value}` + : ""; + return `${field.kind}:${owner}:${field.field}${fallback}`; + }) + .join("|"); +} + +/** The key reader of a recognized field comparator: `(x) => x.k` (or + * `x.k ?? d`), the same read as the comparator's. For several keys it takes + * the key's index too: `(x, j) => j === 0 ? key0 : j === 1 ? key1 : ...`. */ export function buildSortKeyFn( name: string, elem: IrType, - read: SortKeyRead, + reads: readonly SortKeyRead[], loc: SrcLoc, ): IrFunction { const b = new SortIr(loc); const x = b.local("x", elem); + const read = (key: SortKeyRead): IrExpr => { + const field: IrExpr = { ...key.field, obj: x, loc }; + return key.fallback + ? { kind: "nullish", left: field, right: key.fallback, type: F64, loc } + : field; + }; + if (reads.length === 1) { + return { + name, + params: [{ localId: x.localId, name: "x", type: elem }], + returnType: F64, + locals: b.locals, + body: distinctNodes([{ kind: "return", value: read(reads[0]!), loc }]), + loc, + }; + } + const j = b.local("j"); + let value = read(reads[reads.length - 1]!); + for (let index = reads.length - 2; index >= 0; index--) { + value = { + kind: "ternary", + cond: b.eq(j, b.num(index)), + then: read(reads[index]!), + else_: value, + type: F64, + loc, + }; + } return { name, - params: [{ localId: x.localId, name: "x", type: elem }], + params: [ + { localId: x.localId, name: "x", type: elem }, + { localId: j.localId, name: "j", type: F64 }, + ], returnType: F64, locals: b.locals, - body: [{ kind: "return", value: { ...read, obj: x, loc }, loc }], + body: distinctNodes([{ kind: "return", value, loc }]), loc, }; } @@ -826,14 +925,15 @@ export function buildArraySortFn( undefinedTag: number | null, loc: SrcLoc, native = false, - keyed: { order: 1 | -1; key: boolean } | null = null, + keyed: { order: number; key: boolean; keys: number } | null = null, ): IrFunction { const b = new SortIr(loc); const arrT = arrayOf(elem), fnT = funcOf([elem, elem].slice(0, arity), F64); const a = b.local("a", arrT), f = b.local("f", fnT); - const keyT = funcOf([elem], F64); + // A reader of several keys takes the key's index too (buildSortKeyFn). + const keyT = funcOf(keyed && keyed.keys > 1 ? [elem, F64] : [elem], F64); const key = keyed?.key ? b.local("key", keyT) : null; // A recognized comparator lets the runtime order by keys; it still // receives the comparator for inputs whose keys it does not order. diff --git a/packages/compiler/src/ir/validate.test.ts b/packages/compiler/src/ir/validate.test.ts index 059391590..83e0f9a31 100644 --- a/packages/compiler/src/ir/validate.test.ts +++ b/packages/compiler/src/ir/validate.test.ts @@ -2106,6 +2106,7 @@ test("runtime comparator sort validates a recognized comparator's key order", () mod.functions.push( fn("cmp", [elem, elem], F64), fn("key", [row], F64), + fn("keys", [row, F64], F64), fn("text", [row], STRING), ); return mod; @@ -2115,12 +2116,25 @@ test("runtime comparator sort validates a recognized comparator's key order", () keyedModule(F64, [order(-1)], true), keyedModule(row, [order(1), closure("key", [row])]), keyedModule(row, [order(-1), closure("key", [row])], true), + // Two keys, the second descending. + keyedModule(row, [order(2 | (1 << 9)), closure("keys", [row, F64])]), ]) { expect(validateModule(mod)).toEqual([]); expect(deserializeModule(serializeModule(mod))).toEqual(mod); } for (const [mod, message] of [ - [keyedModule(F64, [order(2)]), "key order must be the literal 1 or -1"], + [keyedModule(F64, [order(0.5)]), "key order must be 1, -1 or a key-chain order"], + [keyedModule(row, [order(9), closure("keys", [row, F64])]), "key-chain order"], + [keyedModule(row, [order(2 | (1 << 10)), closure("keys", [row, F64])]), "key-chain order"], + [keyedModule(F64, [order(2)]), "key order on f64"], + [ + keyedModule(row, [order(2), closure("key", [row])]), + "key reader must be (elem, number) => number", + ], + [ + keyedModule(row, [order(1), closure("keys", [row, F64])]), + "key reader must be (elem) => number", + ], [keyedModule(F64, [order(1), closure("key", [row])]), "key order on f64"], [keyedModule(STRING, [order(1)], true), "key order on string"], [keyedModule(row, [order(1)]), "key reader must be (elem) => number"], diff --git a/packages/compiler/src/ir/validate.ts b/packages/compiler/src/ir/validate.ts index a581a6ff9..9635ceb91 100644 --- a/packages/compiler/src/ir/validate.ts +++ b/packages/compiler/src/ir/validate.ts @@ -2388,27 +2388,46 @@ function validateFunction( if (e.method === "sortDense" && e.args[1]?.kind !== "boolLit") err("arrIntrinsic sortDense owned flag must be a boolean literal", e.loc); // A recognized comparator's key order: [order] for numbers, - // [order, key reader] for record and class elements. + // [order, key reader] for record and class elements. An order + // from 2 carries several keys (sortKeysOrder): the key count + // (2..8) in the low byte, a descending bit per key from bit 8, and + // a reader that also takes the key's index. const order = e.args[2]; if (order) { - if (order.kind !== "numLit" || (order.value !== 1 && order.value !== -1)) - err(`arrIntrinsic ${e.method} key order must be the literal 1 or -1`, e.loc); + const count = order.kind === "numLit" && order.value >= 2 ? order.value & 0xff : 1; + if ( + order.kind !== "numLit" || + (order.value !== 1 && + order.value !== -1 && + !( + Number.isInteger(order.value) && + count >= 2 && + count <= 8 && + (order.value >> 8) >> count === 0 + )) + ) + err(`arrIntrinsic ${e.method} key order must be 1, -1 or a key-chain order`, e.loc); const key = e.args[3]?.type; if ( elem.kind === "f64" - ? key !== undefined + ? key !== undefined || count > 1 : elem.kind !== "record" && elem.kind !== "object" ) err(`arrIntrinsic ${e.method} key order on ${elem.kind} elements`, e.loc); + const params = count > 1 ? 2 : 1; if ( elem.kind !== "f64" && (!key || key.kind !== "func" || - key.params.length !== 1 || + key.params.length !== params || !typeEquals(key.params[0]!, elem) || + (params === 2 && key.params[1]!.kind !== "f64") || key.ret.kind !== "f64") ) - err(`arrIntrinsic ${e.method} key reader must be (elem) => number`, e.loc); + err( + `arrIntrinsic ${e.method} key reader must be (elem${params === 2 ? ", number" : ""}) => number`, + e.loc, + ); } } if (e.method === "getNumber" && elem.kind !== "f64") { diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/array-sort-keyed-chains.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/array-sort-keyed-chains.ts.json new file mode 100644 index 000000000..78b80d02b --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/array-sort-keyed-chains.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/array-sort-keyed-chains.ts" + ], + "diags": [] +} diff --git a/packages/runtime/src/scr_array.c b/packages/runtime/src/scr_array.c index 0dcfc318c..4887f23ff 100644 --- a/packages/runtime/src/scr_array.c +++ b/packages/runtime/src/scr_array.c @@ -2739,7 +2739,8 @@ typedef struct { uint32_t arity; bool numbers; bool borrowing; /* enter through f->fn_borrowed, lending the slots */ - int32_t order; /* 1 or -1: f orders number keys ascending or descending */ + int32_t order; /* 1 or -1: f orders number keys ascending or descending; + * from 2, the keys of a multi-key comparator (SCR_SORT_KEYS) */ ScrClosure *key; /* the key reader of reference elements under `order` */ ScrSortLess less; /* a caller comparison in place of `f` (scr_sort_slots_with) */ void *less_ctx; @@ -2803,11 +2804,24 @@ static bool scr_sort_greater(const ScrSortCompare *c, uint64_t left, uint64_t ri #define SCR_SORT_PREFETCH(c, buf, index, limit) \ SCR_SORT_PREFETCH_AT(c, buf, (index) + SCR_SORT_PREFETCH_DISTANCE, limit) +/* A multi-key comparator (`k0 || k1 || ...` over field keys of reference + * elements) passes `order` as its key count in the low byte and, from bit + * 8, one bit per key that orders descending. Its reader closure takes the + * element and the key's index. */ +#define SCR_SORT_KEYS_MAX 8 +static bool scr_sort_keys_valid(int32_t order) { + uint32_t count = (uint32_t)order & 0xff; + return order > 0 && count >= 2 && count <= SCR_SORT_KEYS_MAX && + ((uint32_t)order >> 8) >> count == 0; +} + static ScrSortCompare scr_sort_compare_for(const ScrArr *a, ScrClosure *f, uint32_t arity, int32_t order, ScrClosure *key) { ScrSortCompare c = {f, NULL, arity, a->elem == SCR_ELEM_F64, false, order, key}; - if (order != 0 && (arity != 2 || (order != 1 && order != -1) || - (c.numbers ? key != NULL : a->elem != SCR_ELEM_REF || key == NULL))) + bool multi = order >= 2; + if (order != 0 && + (arity != 2 || (order != 1 && order != -1 && !scr_sort_keys_valid(order)) || + (c.numbers ? key != NULL || multi : a->elem != SCR_ELEM_REF || key == NULL))) scr_trap("scriptc: invalid keyed sort\n"); if (!c.numbers) { c.borrowing = f->fn_borrowed != NULL; @@ -3179,6 +3193,63 @@ static bool scr_sort_keyed_refs(uint64_t *values, size_t n, const ScrSortCompare return true; } +/* Key `j` of a reference element through the borrowed multi-key reader. */ +static double scr_sort_read_key_at(const ScrSortCompare *c, uint64_t slot, uint32_t j) { + ScrClosure *reader = c->key; + void *x = scr_slot_to_ptr(slot); + if (reader->fn_borrowed) + return ((double (*)(ScrClosure *, void *, double))reader->fn_borrowed)(reader, x, (double)j); + return ((double (*)(ScrClosure *, void *, double))reader->fn)(reader, x ? c->retain(x) : x, + (double)j); +} + +/* Sorts reference slots by several keys, the first deciding. A comparator + * `d0 || d1 || ...` whose terms each order one key orders lexicographically: + * a term that is 0 or NaN (equal keys, or Infinity minus Infinity) falls + * through to the next. With no key NaN or -0 that order is consistent, so + * a stable sort by the keys gives the comparator sort's permutation. Every + * key is read and checked before any slot moves; stable passes from the + * last key to the first then leave the slots in key order. */ +static bool scr_sort_keyed_multi(uint64_t *values, size_t n, const ScrSortCompare *c) { + uint32_t count = (uint32_t)c->order & 0xff, descending = (uint32_t)c->order >> 8; + if (n > SIZE_MAX / count / sizeof(uint64_t)) scr_arr_oom(); + uint64_t *keys = malloc(n * count * sizeof(*keys)); + if (!keys) scr_arr_oom(); + for (size_t i = 0; i < n; i++) + for (uint32_t j = 0; j < count; j++) { + uint64_t flip = (descending >> j) & 1 ? ~UINT64_C(0) : 0; + if (!scr_sort_key_encode(scr_sort_read_key_at(c, values[i], j), flip, + &keys[i * count + j])) { + free(keys); + return false; + } + } + /* Each item's slot field holds the element's input position. */ + ScrSortItem *items = malloc(n * sizeof(*items)); + ScrSortItem *tmp = n > SCR_SORT_KEYED_INSERTION ? malloc(n * sizeof(*tmp)) : NULL; + if (!items || (n > SCR_SORT_KEYED_INSERTION && !tmp)) scr_arr_oom(); + for (size_t i = 0; i < n; i++) items[i] = (ScrSortItem){keys[i * count + count - 1], i}; + ScrSortItem *buf = items, *other = tmp; + for (uint32_t j = count; j-- > 0;) { + if (j + 1 < count) + for (size_t i = 0; i < n; i++) buf[i].key = keys[buf[i].slot * count + j]; + ScrSortItem *sorted = buf; + SCR_SORT_KEYED_RECORDS(ScrSortItem, SCR_SORT_ITEM_KEY, 64, buf, other, n, false, &sorted); + if (sorted != buf) { + other = buf; + buf = sorted; + } + } + /* The key buffer holds at least n words: it keeps the input slots while + * they are written back in order. */ + memcpy(keys, values, n * sizeof(*values)); + for (size_t i = 0; i < n; i++) values[i] = keys[buf[i].slot]; + free(keys); + free(items); + free(tmp); + return true; +} + /* Orders `n >= 2` slots by key when the comparator allows it. Returns false, * leaving the slots as they were, when a key is NaN or -0: the comparator * sort then runs instead. */ @@ -3186,7 +3257,8 @@ static bool scr_sort_keyed(uint64_t *values, size_t n, const ScrSortCompare *c) uint64_t flip = c->order < 0 ? ~UINT64_C(0) : 0; if (c->numbers) return scr_sort_keyed_numbers(values, n, flip); scr_dyn_this_push_dyn(NULL); - bool sorted = scr_sort_keyed_refs(values, n, c, flip); + bool sorted = c->order >= 2 ? scr_sort_keyed_multi(values, n, c) + : scr_sort_keyed_refs(values, n, c, flip); scr_dyn_this_pop(); return sorted; } diff --git a/packages/runtime/src/scr_array.h b/packages/runtime/src/scr_array.h index 909f9ae8c..1678d1568 100644 --- a/packages/runtime/src/scr_array.h +++ b/packages/runtime/src/scr_array.h @@ -398,7 +398,11 @@ struct ScrClosure; * (-1) without side effects: f64 elements are their own keys, and reference * elements read theirs through the borrowed `key` closure (element -> * number). Keys that are all numbers other than NaN and -0 then sort without - * calling `f`, with the same stable result. */ + * calling `f`, with the same stable result. An `order` from 2 says `f` + * orders reference elements by several such keys, the first deciding: the + * low byte is the key count (2..8), each bit from bit 8 marks a key that + * orders descending, and `key` reads key j of an element (element, j -> + * number). */ void scr_arr_sort_values(ScrArr *a, double count, struct ScrClosure *f, uint32_t arity, int32_t order, struct ScrClosure *key); /* Comparator sort of a dense array in place: `a.sort(f)`, or with `owned` diff --git a/packages/runtime/test/test_array.c b/packages/runtime/test/test_array.c index a3ad4735b..3eea84dcf 100644 --- a/packages/runtime/test/test_array.c +++ b/packages/runtime/test/test_array.c @@ -1720,6 +1720,94 @@ static void test_keyed_sort(void) { check(mock_live == live0, "keyed record sorts balance every key read"); } +/* Multi-key sorts: a record's value v has keys v / 100 (major), v / 10 % 10 + * (middle) and v % 10 (minor). Value -1 reads key 1 as -0 and value -2 reads + * it as NaN; the comparators treat both as the comparator expressions do. */ +static double multi_key_of(int value, uint32_t j) { + if (j == 1 && value == -1) return -0.0; + if (j == 1 && value == -2) return NAN; + if (value < 0) value = 0; + return j == 0 ? value / 100 : j == 1 ? value / 10 % 10 : value % 10; +} + +static double multi_term(double d) { return d != d ? 0 : d; } + +/* `a.major - b.major || b.middle - a.middle` */ +static double multi_two(ScrClosure *self, void *a, void *b) { + (void)self; + sort_calls++; + int x = ((MockRec *)a)->value, y = ((MockRec *)b)->value; + mock_release(a); + mock_release(b); + double d = multi_term(multi_key_of(x, 0) - multi_key_of(y, 0)); + return d != 0 ? d : multi_key_of(y, 1) - multi_key_of(x, 1); +} + +/* `b.major - a.major || a.middle - b.middle || b.minor - a.minor` */ +static double multi_three(ScrClosure *self, void *a, void *b) { + (void)self; + sort_calls++; + int x = ((MockRec *)a)->value, y = ((MockRec *)b)->value; + mock_release(a); + mock_release(b); + double d = multi_term(multi_key_of(y, 0) - multi_key_of(x, 0)); + if (d == 0) d = multi_term(multi_key_of(x, 1) - multi_key_of(y, 1)); + return d != 0 ? d : multi_key_of(y, 2) - multi_key_of(x, 2); +} + +static double multi_key(ScrClosure *self, void *a, double j) { + (void)self; + key_reads++; + int value = ((MockRec *)a)->value; + mock_release(a); + return multi_key_of(value, (uint32_t)j); +} + +static void test_multi_key_sort(void) { + ScrClosure two = {SIZE_MAX, (void *)multi_two, 0, NULL, 0}; + ScrClosure three = {SIZE_MAX, (void *)multi_three, 0, NULL, 0}; + ScrClosure key = {SIZE_MAX, (void *)multi_key, 0, NULL, 0}; + /* Two keys with the second descending; three with the first and third. */ + const int32_t two_order = 2 | 1 << 9, three_order = 3 | 1 << 8 | 1 << 10; + const size_t sizes[] = {2, 5, 40, 49, 1000, 70000, 140000}; + long live0 = mock_live; + bool identical = true, direct = true, fallback = true; + for (size_t s = 0; s < sizeof(sizes) / sizeof(*sizes); s++) { + for (int variant = 0; variant < 6; variant++) { + /* Two or three keys; then no special key, a -0 or a NaN in key 1. */ + bool use_three = variant & 1; + int special = variant / 2; + size_t n = sizes[s]; + ScrArr *expected = scr_arr_new_ops(&mock_ops, 0); + ScrArr *actual = scr_arr_new_ops(&mock_ops, 0); + for (size_t i = 0; i < n; i++) { + int value = (int)(keyed_random() % (n < 100 ? 300 : 3000)); + if (special && i == n / 2) value = -special; + MockRec *r = mock_new(value); + scr_arr_push_ref(expected, r); + scr_arr_push_ref(actual, mock_retain(r)); + } + ScrClosure *cmp = use_three ? &three : &two; + int32_t order = use_three ? three_order : two_order; + scr_arr_sort_values(expected, (double)n, cmp, 2, 0, NULL); + sort_calls = 0; + key_reads = 0; + if (s % 2) scr_arr_sort_values(actual, (double)n, cmp, 2, order, &key); + else scr_arr_sort_dense(actual, cmp, 2, s % 4 == 2, order, &key); + for (size_t i = 0; i < n; i++) + identical &= scr_arr_peek_ref(expected, (double)i) == scr_arr_peek_ref(actual, (double)i); + if (special) fallback &= sort_calls > 0; + else direct &= sort_calls == 0 && key_reads == (long)n * (use_three ? 3 : 2); + scr_arr_release(expected); + scr_arr_release(actual); + } + } + check(identical, "multi-key record sorts keep the comparator sort's stable permutation"); + check(direct, "multi-key record sorts read each key once and call no comparator"); + check(fallback, "a -0 or NaN in a later key sorts through the comparator"); + check(mock_live == live0, "multi-key record sorts balance every key read"); +} + /* Typed-array sorts: the default numeric order and recognized comparator * orders against the comparator sort, which follows Node's merge schedule. */ static double typed_mod_ascending(ScrClosure *self, double a, double b) { @@ -2046,6 +2134,7 @@ int main(int argc, char **argv) { test_primitive_sort(); test_comparator_sort(); test_keyed_sort(); + test_multi_key_sort(); test_dense_slots(); test_typed_sort(); test_dense_comparator_sort(); diff --git a/tests/corpus/array-sort-keyed-chains.ts b/tests/corpus/array-sort-keyed-chains.ts new file mode 100644 index 000000000..e058cddf7 --- /dev/null +++ b/tests/corpus/array-sort-keyed-chains.ts @@ -0,0 +1,149 @@ +// Comparators that chain number-field terms with `||` (`a.k - b.k || a.j - +// b.j`), optionally over optional fields with a number default (`(a.k ?? 0) +// - (b.k ?? 0)`), sort records and class instances by several keys: equal +// keys fall through to the next term, ties keep their order, and NaN and -0 +// keys, holes and comparators with other terms match Node. +interface Row { + id: number; + group: number; + score?: number; + weight: number; + label: string; +} + +class Entry { + id: number; + major: number; + minor: number; + constructor(id: number, major: number, minor: number) { + this.id = id; + this.major = major; + this.minor = minor; + } +} + +let seed = 11; +function random(): number { + seed = (Math.imul(seed, 1103515245) + 12345) >>> 0; + return seed / 4294967296; +} + +function ids(rows: { id: number }[]): string { + let hash = 0; + for (const row of rows) hash = (Math.imul(hash, 31) + row.id) | 0; + return rows.length + ":" + hash; +} + +function makeRows(n: number, groups: number, scored: number): Row[] { + const rows: Row[] = []; + for (let i = 0; i < n; i++) { + const row: Row = { + id: i, + group: Math.floor(random() * groups), + weight: Math.floor(random() * 4) / 2, + label: "r" + i, + }; + if (random() < scored) row.score = Math.floor(random() * 20) - 5; + rows.push(row); + } + return rows; +} + +for (const [groups, scored] of [ + [1, 1], + [3, 0.5], + [10, 0.9], + [1000, 0.2], +]) { + const parts: string[] = []; + for (const n of [0, 1, 2, 5, 16, 47, 48, 49, 200, 5000]) { + const rows = makeRows(n, groups!, scored!); + const two = rows.toSorted((a, b) => a.group - b.group || a.weight - b.weight); + const mixed = rows.toSorted((a, b) => b.group - a.group || a.weight - b.weight); + const three = rows.toSorted( + (a, b) => (b.score ?? 0) - (a.score ?? 0) || a.group - b.group || b.weight - a.weight, + ); + const ternary = rows.toSorted( + (a, b) => (a.group < b.group ? -1 : a.group > b.group ? 1 : 0) || b.id - a.id, + ); + const missingLast = rows.toSorted((a, b) => (a.score ?? 99) - (b.score ?? 99) || a.id - b.id); + const single = rows.toSorted((a, b) => (b.score ?? -1) - (a.score ?? -1)); + const inPlace = rows.slice(); + inPlace.sort((p, q) => q.weight - p.weight || p.group - q.group); + parts.push( + [two, mixed, three, ternary, missingLast, single, inPlace].map((sorted) => ids(sorted)).join("/"), + ); + } + console.log(groups, scored, parts.join(" ")); +} + +// Ties on every key keep the input order. +const tied = makeRows(24, 2, 0.5); +console.log( + tied + .toSorted((a, b) => a.group - b.group || (a.score ?? 0) - (b.score ?? 0)) + .map((r) => r.label) + .join(","), +); + +// Class instances sort by their fields. +const entries: Entry[] = []; +for (let i = 0; i < 300; i++) entries.push(new Entry(i, i % 7, (i * 13) % 5)); +console.log(ids(entries.toSorted((a, b) => a.major - b.major || b.minor - a.minor))); +console.log( + entries + .toSorted((a, b) => b.major - a.major || a.minor - b.minor || a.id - b.id) + .slice(0, 12) + .map((e) => `${e.id}:${e.major}:${e.minor}`) + .join(" "), +); + +// A NaN or -0 key in a later term, and infinities in the first, sort +// through the comparator. +function keyed(values: [number, number][]): Row[] { + return values.map(([group, weight], id) => ({ id, group, weight, label: "k" + id })); +} +const nan = keyed([ + [1, 3], + [0, NaN], + [1, 1], + [0, 2], + [1, NaN], + [0, 0], + [2, 5], + [1, 2], +]); +console.log(nan.toSorted((a, b) => a.group - b.group || a.weight - b.weight).map((r) => r.id).join(",")); +console.log(nan.toSorted((a, b) => b.group - a.group || b.weight - a.weight).map((r) => r.id).join(",")); +const zeros = keyed(Array.from({ length: 30 }, (_, i): [number, number] => [i % 3, i % 4 === 0 ? -0 : (i % 4) - 2])); +console.log(zeros.toSorted((a, b) => a.group - b.group || a.weight - b.weight).map((r) => r.id).join(",")); +const infinite = keyed([ + [Infinity, 2], + [-Infinity, 1], + [Infinity, 1], + [0, 3], + [-Infinity, 0], + [Infinity, 0], +]); +console.log( + infinite.toSorted((a, b) => a.group - b.group || a.weight - b.weight).map((r) => r.id).join(","), +); + +// A default that differs per operand is not one key; the comparator runs. +const uneven = makeRows(40, 3, 0.5); +console.log(ids(uneven.toSorted((a, b) => (a.score ?? 0) - (b.score ?? 1) || a.id - b.id))); + +// A chain with a term that is not a key difference calls the comparator. +const labelled = makeRows(60, 4, 1); +let calls = 0; +const byLabel = labelled.toSorted((a, b) => { + calls++; + return a.group - b.group || (a.label < b.label ? -1 : a.label > b.label ? 1 : 0); +}); +console.log(ids(byLabel), calls > 0); + +// Holes and a sparse tail take the snapshot sort. +const holes: Row[] = makeRows(20, 3, 0.5); +holes.length = 26; +holes.sort((a, b) => a.group - b.group || (b.score ?? 0) - (a.score ?? 0)); +console.log(holes.length, holes.map((r) => (r === undefined ? "-" : r.id)).join(",")); diff --git a/tests/harness/corpus-costs/array-sort-keyed-chains.ts.json b/tests/harness/corpus-costs/array-sort-keyed-chains.ts.json new file mode 100644 index 000000000..1d0513983 --- /dev/null +++ b/tests/harness/corpus-costs/array-sort-keyed-chains.ts.json @@ -0,0 +1 @@ +{ "plain": 3900, "san": 5400 }