Optimize loadTreeForQuery by filtering eagerly in ParamsBuilder (#2890)

* nit: simplify constructions of RunCaseSpecific

* Optimize loadTreeForQuery by filtering eagerly in ParamsBuilder

Previously, if you load a query like
`webgpu:api,validation,encoding,cmds,render,draw:vertex_buffer_OOB:type="draw";VBSize="exile";IBSize="exile";VStride0=false;IStride0=true;AStride="zero";offset=1`
or
`webgpu:api,validation,encoding,cmds,render,draw:vertex_buffer_OOB:type="draw";VBSize="exile";IBSize="exile";*`
loadTreeForQuery would iterate *all* of the cases for that test
`webgpu:api,validation,encoding,cmds,render,draw:*`
and filter out the wrong ones at the very end.

This changes the ParamsBuilder iteration to be filter-aware, so it
doesn't even start generating any of the subspaces of the combinatorial
space that it can already see have a mismatch.

For the case above, this improves the runtime of loadTreeForQuery from
205ms to ~2ms (100x faster!). This saves tons of time if tests are run
in a way where loadTreeForQuery is called for many sub-test or
single-case queries (as is done in Chromium).

Bug: https://crbug.com/1470849

Author: kainino0x

src/common/framework/params_builder.ts

@@ -1,7 +1,10 @@
-import { Merged, assertMergedWithoutOverlap, mergeParams } from '../internal/params_utils.js';
+import { Merged, mergeParams, mergeParamsChecked } from '../internal/params_utils.js';
+import { comparePublicParamsPaths, Ordering } from '../internal/query/compare.js';
 import { stringifyPublicParams } from '../internal/query/stringify_params.js';
 import { assert, mapLazy } from '../util/util.js';
 
+import { TestParams } from './fixture.js';
+
 // ================================================================
 // "Public" ParamsBuilder API / Documentation
 // ================================================================
@@ -102,27 +105,32 @@ export type CaseSubcaseIterable<CaseP, SubcaseP> = Iterable<
  * Base class for `CaseParamsBuilder` and `SubcaseParamsBuilder`.
  */
 export abstract class ParamsBuilderBase<CaseP extends {}, SubcaseP extends {}> {
-  protected readonly cases: () => Generator<CaseP>;
+  protected readonly cases: (caseFilter: TestParams | null) => Generator<CaseP>;
 
-  constructor(cases: () => Generator<CaseP>) {
+  constructor(cases: (caseFilter: TestParams | null) => Generator<CaseP>) {
     this.cases = cases;
   }
 
   /**
    * Hidden from test files. Use `builderIterateCasesWithSubcases` to access this.
    */
-  protected abstract iterateCasesWithSubcases(): CaseSubcaseIterable<CaseP, SubcaseP>;
+  protected abstract iterateCasesWithSubcases(
+    caseFilter: TestParams | null
+  ): CaseSubcaseIterable<CaseP, SubcaseP>;
 }
 
 /**
  * Calls the (normally hidden) `iterateCasesWithSubcases()` method.
  */
-export function builderIterateCasesWithSubcases(builder: ParamsBuilderBase<{}, {}>) {
+export function builderIterateCasesWithSubcases(
+  builder: ParamsBuilderBase<{}, {}>,
+  caseFilter: TestParams | null
+) {
   interface IterableParamsBuilder {
-    iterateCasesWithSubcases(): CaseSubcaseIterable<{}, {}>;
+    iterateCasesWithSubcases(caseFilter: TestParams | null): CaseSubcaseIterable<{}, {}>;
   }
 
-  return ((builder as unknown) as IterableParamsBuilder).iterateCasesWithSubcases();
+  return ((builder as unknown) as IterableParamsBuilder).iterateCasesWithSubcases(caseFilter);
 }
 
 /**
@@ -136,31 +144,66 @@ export function builderIterateCasesWithSubcases(builder: ParamsBuilderBase<{}, {
 export class CaseParamsBuilder<CaseP extends {}>
   extends ParamsBuilderBase<CaseP, {}>
   implements Iterable<CaseP>, ParamsBuilder {
-  *iterateCasesWithSubcases(): CaseSubcaseIterable<CaseP, {}> {
-    for (const a of this.cases()) {
-      yield [a, undefined];
+  *iterateCasesWithSubcases(caseFilter: TestParams | null): CaseSubcaseIterable<CaseP, {}> {
+    for (const caseP of this.cases(caseFilter)) {
+      if (caseFilter) {
+        // this.cases() only filters out cases which conflict with caseFilter. Now that we have
+        // the final caseP, filter out cases which are missing keys that caseFilter requires.
+        const ordering = comparePublicParamsPaths(caseP, caseFilter);
+        if (ordering === Ordering.StrictSuperset || ordering === Ordering.Unordered) {
+          continue;
+        }
+      }
+
+      yield [caseP, undefined];
     }
   }
 
   [Symbol.iterator](): Iterator<CaseP> {
-    return this.cases();
+    return this.cases(null);
   }
 
   /** @inheritDoc */
   expandWithParams<NewP extends {}>(
-    expander: (_: Merged<{}, CaseP>) => Iterable<NewP>
+    expander: (_: CaseP) => Iterable<NewP>
   ): CaseParamsBuilder<Merged<CaseP, NewP>> {
-    const newGenerator = genExpandWithParams(this.cases, expander);
-    return new CaseParamsBuilder(() => newGenerator({}));
+    const baseGenerator = this.cases;
+    return new CaseParamsBuilder(function* (caseFilter) {
+      for (const a of baseGenerator(caseFilter)) {
+        for (const b of expander(a)) {
+          if (caseFilter) {
+            // If the expander generated any key-value pair that conflicts with caseFilter, skip.
+            if (Object.entries(b).some(([k, v]) => k in caseFilter && caseFilter[k] !== v)) {
+              continue;
+            }
+          }
+
+          yield mergeParamsChecked(a, b);
+        }
+      }
+    });
   }
 
   /** @inheritDoc */
   expand<NewPKey extends string, NewPValue>(
     key: NewPKey,
-    expander: (_: Merged<{}, CaseP>) => Iterable<NewPValue>
+    expander: (_: CaseP) => Iterable<NewPValue>
   ): CaseParamsBuilder<Merged<CaseP, { [name in NewPKey]: NewPValue }>> {
-    const newGenerator = genExpand(this.cases, key, expander);
-    return new CaseParamsBuilder(() => newGenerator({}));
+    const baseGenerator = this.cases;
+    return new CaseParamsBuilder(function* (caseFilter) {
+      for (const a of baseGenerator(caseFilter)) {
+        assert(!(key in a), `New key '${key}' already exists in ${JSON.stringify(a)}`);
+
+        const caseFilterV = caseFilter?.[key];
+        for (const v of expander(a)) {
+          // If the expander generated a value for this key that conflicts with caseFilter, skip.
+          if (caseFilter && (caseFilterV as {}) !== v) {
+            continue;
+          }
+          yield { ...a, [key]: v } as Merged<CaseP, { [name in NewPKey]: NewPValue }>;
+        }
+      }
+    });
   }
 
   /** @inheritDoc */
@@ -189,13 +232,17 @@ export class CaseParamsBuilder<CaseP extends {}>
   }
 
   /** @inheritDoc */
-  filter(pred: (_: Merged<{}, CaseP>) => boolean): CaseParamsBuilder<CaseP> {
-    const newGenerator = filterGenerator(this.cases, pred);
-    return new CaseParamsBuilder(() => newGenerator({}));
+  filter(pred: (_: CaseP) => boolean): CaseParamsBuilder<CaseP> {
+    const baseGenerator = this.cases;
+    return new CaseParamsBuilder(function* (caseFilter) {
+      for (const a of baseGenerator(caseFilter)) {
+        if (pred(a)) yield a;
+      }
+    });
   }
 
   /** @inheritDoc */
-  unless(pred: (_: Merged<{}, CaseP>) => boolean): CaseParamsBuilder<CaseP> {
+  unless(pred: (_: CaseP) => boolean): CaseParamsBuilder<CaseP> {
     return this.filter(x => !pred(x));
   }
 
@@ -205,12 +252,9 @@ export class CaseParamsBuilder<CaseP extends {}>
    * generate new subcases instead of new cases.
    */
   beginSubcases(): SubcaseParamsBuilder<CaseP, {}> {
-    return new SubcaseParamsBuilder(
-      () => this.cases(),
-      function* () {
-        yield {};
-      }
-    );
+    return new SubcaseParamsBuilder(this.cases, function* () {
+      yield {};
+    });
   }
 }
 
@@ -235,13 +279,25 @@ export class SubcaseParamsBuilder<CaseP extends {}, SubcaseP extends {}>
   implements ParamsBuilder {
   protected readonly subcases: (_: CaseP) => Generator<SubcaseP>;
 
-  constructor(cases: () => Generator<CaseP>, generator: (_: CaseP) => Generator<SubcaseP>) {
+  constructor(
+    cases: (caseFilter: TestParams | null) => Generator<CaseP>,
+    generator: (_: CaseP) => Generator<SubcaseP>
+  ) {
     super(cases);
     this.subcases = generator;
   }
 
-  *iterateCasesWithSubcases(): CaseSubcaseIterable<CaseP, SubcaseP> {
-    for (const caseP of this.cases()) {
+  *iterateCasesWithSubcases(caseFilter: TestParams | null): CaseSubcaseIterable<CaseP, SubcaseP> {
+    for (const caseP of this.cases(caseFilter)) {
+      if (caseFilter) {
+        // this.cases() only filters out cases which conflict with caseFilter. Now that we have
+        // the final caseP, filter out cases which are missing keys that caseFilter requires.
+        const ordering = comparePublicParamsPaths(caseP, caseFilter);
+        if (ordering === Ordering.StrictSuperset || ordering === Ordering.Unordered) {
+          continue;
+        }
+      }
+
       const subcases = Array.from(this.subcases(caseP));
       if (subcases.length) {
         yield [caseP, subcases];
@@ -253,15 +309,32 @@ export class SubcaseParamsBuilder<CaseP extends {}, SubcaseP extends {}>
   expandWithParams<NewP extends {}>(
     expander: (_: Merged<CaseP, SubcaseP>) => Iterable<NewP>
   ): SubcaseParamsBuilder<CaseP, Merged<SubcaseP, NewP>> {
-    return new SubcaseParamsBuilder(this.cases, genExpandWithParams(this.subcases, expander));
+    const baseGenerator = this.subcases;
+    return new SubcaseParamsBuilder(this.cases, function* (base) {
+      for (const a of baseGenerator(base)) {
+        for (const b of expander(mergeParams(base, a))) {
+          yield mergeParamsChecked(a, b);
+        }
+      }
+    });
   }
 
   /** @inheritDoc */
   expand<NewPKey extends string, NewPValue>(
     key: NewPKey,
     expander: (_: Merged<CaseP, SubcaseP>) => Iterable<NewPValue>
   ): SubcaseParamsBuilder<CaseP, Merged<SubcaseP, { [name in NewPKey]: NewPValue }>> {
-    return new SubcaseParamsBuilder(this.cases, genExpand(this.subcases, key, expander));
+    const baseGenerator = this.subcases;
+    return new SubcaseParamsBuilder(this.cases, function* (base) {
+      for (const a of baseGenerator(base)) {
+        const before = mergeParams(base, a);
+        assert(!(key in before), () => `Key '${key}' already exists in ${JSON.stringify(before)}`);
+
+        for (const v of expander(before)) {
+          yield { ...a, [key]: v } as Merged<SubcaseP, { [k in NewPKey]: NewPValue }>;
+        }
+      }
+    });
   }
 
   /** @inheritDoc */
@@ -283,7 +356,12 @@ export class SubcaseParamsBuilder<CaseP extends {}, SubcaseP extends {}>
 
   /** @inheritDoc */
   filter(pred: (_: Merged<CaseP, SubcaseP>) => boolean): SubcaseParamsBuilder<CaseP, SubcaseP> {
-    return new SubcaseParamsBuilder(this.cases, filterGenerator(this.subcases, pred));
+    const baseGenerator = this.subcases;
+    return new SubcaseParamsBuilder(this.cases, function* (base) {
+      for (const a of baseGenerator(base)) {
+        if (pred(mergeParams(base, a))) yield a;
+      }
+    });
   }
 
   /** @inheritDoc */
@@ -292,54 +370,6 @@ export class SubcaseParamsBuilder<CaseP extends {}, SubcaseP extends {}>
   }
 }
 
-/** Creates a generator function for expandWithParams() methods above. */
-function genExpandWithParams<Base, A, B>(
-  baseGenerator: (_: Base) => Generator<A>,
-  expander: (_: Merged<Base, A>) => Iterable<B>
-): (_: Base) => Generator<Merged<A, B>> {
-  return function* (base: Base) {
-    for (const a of baseGenerator(base)) {
-      for (const b of expander(mergeParams(base, a))) {
-        const merged = mergeParams(a, b);
-        assertMergedWithoutOverlap([a, b], merged);
-
-        yield merged;
-      }
-    }
-  };
-}
-
-/** Creates a generator function for expand() methods above. */
-function genExpand<Base, A, NewPKey extends string, NewPValue>(
-  baseGenerator: (_: Base) => Generator<A>,
-  key: NewPKey,
-  expander: (_: Merged<Base, A>) => Iterable<NewPValue>
-): (_: Base) => Generator<Merged<A, { [k in NewPKey]: NewPValue }>> {
-  return function* (base: Base) {
-    for (const a of baseGenerator(base)) {
-      const before = mergeParams(base, a);
-      assert(!(key in before), () => `Key '${key}' already exists in ${JSON.stringify(before)}`);
-
-      for (const v of expander(before)) {
-        yield { ...a, [key]: v } as Merged<A, { [k in NewPKey]: NewPValue }>;
-      }
-    }
-  };
-}
-
-function filterGenerator<Base, A>(
-  baseGenerator: (_: Base) => Generator<A>,
-  pred: (_: Merged<Base, A>) => boolean
-): (_: Base) => Generator<A> {
-  return function* (base: Base) {
-    for (const a of baseGenerator(base)) {
-      if (pred(mergeParams(base, a))) {
-        yield a;
-      }
-    }
-  };
-}
-
 /** Assert an object is not a Generator (a thing returned from a generator function). */
 function assertNotGenerator(x: object) {
   if ('constructor' in x) {

src/common/internal/params_utils.ts

@@ -124,10 +124,15 @@ export function mergeParams<A extends {}, B extends {}>(a: A, b: B): Merged<A, B
   return { ...a, ...b } as Merged<A, B>;
 }
 
-/** Asserts that the result of a mergeParams didn't have overlap. This is not extremely fast. */
-export function assertMergedWithoutOverlap([a, b]: [{}, {}], merged: {}): void {
+/**
+ * Merges two objects into one `{ ...a, ...b }` and asserts they had no overlapping keys.
+ * This is slower than {@link mergeParams}.
+ */
+export function mergeParamsChecked<A extends {}, B extends {}>(a: A, b: B): Merged<A, B> {
+  const merged = mergeParams(a, b);
   assert(
     Object.keys(merged).length === Object.keys(a).length + Object.keys(b).length,
     () => `Duplicate key between ${JSON.stringify(a)} and ${JSON.stringify(b)}`
   );
+  return merged;
 }