fix(selectors): refactor chaining logic (#13764)

This fixes a few issues:
- strict mode was producing false negatives if multiple query paths
  lead to the same element being picked;
- in some cases the number of intermediate items in the list was
  exponential and crashed quickly.

What changed:
- `visible` engine is a real engine now;
- `capture` selectors are transformed to `has=` selectors for
  easier implementation;
- chained querying switched from a list to a set to avoid
  exponential size.

packages/playwright-core/src/server/injected/injectedScript.ts

@@ -63,11 +63,6 @@ export interface SelectorEngineV2 {
   queryAll(root: SelectorRoot, body: any): Element[];
 }
 
-export type ElementMatch = {
-  element: Element;
-  capture: Element | undefined;
-};
-
 export type HitTargetInterceptionResult = {
   stop: () => 'done' | { hitTargetDescription: string };
 };
@@ -105,7 +100,7 @@ export class InjectedScript {
     this._engines.set('data-test:light', this._createAttributeEngine('data-test', false));
     this._engines.set('css', this._createCSSEngine());
     this._engines.set('nth', { queryAll: () => [] });
-    this._engines.set('visible', { queryAll: () => [] });
+    this._engines.set('visible', this._createVisibleEngine());
     this._engines.set('control', this._createControlEngine());
     this._engines.set('has', this._createHasEngine());
 
@@ -140,93 +135,70 @@ export class InjectedScript {
   }
 
   querySelector(selector: ParsedSelector, root: Node, strict: boolean): Element | undefined {
-    if (!(root as any)['querySelector'])
+    const result = this.querySelectorAll(selector, root);
-      throw this.createStacklessError('Node is not queryable.');
+    if (strict && result.length > 1)
-    this._evaluator.begin();
+      throw this.strictModeViolationError(selector, result);
-    try {
+    return result[0];
-      const result = this._querySelectorRecursively([{ element: root as Element, capture: undefined }], selector, 0, new Map());
-      if (strict && result.length > 1)
-        throw this.strictModeViolationError(selector, result.map(r => r.element));
-      return result[0]?.capture || result[0]?.element;
-    } finally {
-      this._evaluator.end();
-    }
   }
 
-  private _querySelectorRecursively(roots: ElementMatch[], selector: ParsedSelector, index: number, queryCache: Map<Element, Element[][]>): ElementMatch[] {
+  private _queryNth(roots: Set<Element>, part: ParsedSelectorPart): Set<Element> {
-    if (index === selector.parts.length)
+    const list = [...roots];
-      return roots;
+    let nth = +part.body;
-
+    if (nth === -1)
-    const part = selector.parts[index];
+      nth = list.length - 1;
-    if (part.name === 'nth') {
+    return new Set<Element>(list.slice(nth, nth + 1));
-      let filtered: ElementMatch[] = [];
-      if (part.body === '0') {
-        filtered = roots.slice(0, 1);
-      } else if (part.body === '-1') {
-        if (roots.length)
-          filtered = roots.slice(roots.length - 1);
-      } else {
-        if (typeof selector.capture === 'number')
-          throw this.createStacklessError(`Can't query n-th element in a request with the capture.`);
-        const nth = +part.body;
-        const set = new Set<Element>();
-        for (const root of roots) {
-          set.add(root.element);
-          if (nth + 1 === set.size)
-            filtered = [root];
-        }
-      }
-      return this._querySelectorRecursively(filtered, selector, index + 1, queryCache);
-    }
-
-    if (part.name === 'visible') {
-      const visible = Boolean(part.body);
-      const filtered = roots.filter(match => visible === isVisible(match.element));
-      return this._querySelectorRecursively(filtered, selector, index + 1, queryCache);
-    }
-
-    const result: ElementMatch[] = [];
-    for (const root of roots) {
-      const capture = index - 1 === selector.capture ? root.element : root.capture;
-
-      // Do not query engine twice for the same element.
-      let queryResults = queryCache.get(root.element);
-      if (!queryResults) {
-        queryResults = [];
-        queryCache.set(root.element, queryResults);
-      }
-      let all = queryResults[index];
-      if (!all) {
-        all = this._queryEngineAll(part, root.element);
-        queryResults[index] = all;
-      }
-
-      for (const element of all) {
-        if (!('nodeName' in element))
-          throw this.createStacklessError(`Expected a Node but got ${Object.prototype.toString.call(element)}`);
-        result.push({ element, capture });
-      }
-    }
-    return this._querySelectorRecursively(result, selector, index + 1, queryCache);
   }
 
   querySelectorAll(selector: ParsedSelector, root: Node): Element[] {
+    if (selector.capture !== undefined) {
+      if (selector.parts.some(part => part.name === 'nth'))
+        throw this.createStacklessError(`Can't query n-th element in a request with the capture.`);
+      const withHas: ParsedSelector = { parts: selector.parts.slice(0, selector.capture + 1) };
+      if (selector.capture < selector.parts.length - 1) {
+        const body = { parts: selector.parts.slice(selector.capture + 1) };
+        const has: ParsedSelectorPart = { name: 'has', body, source: stringifySelector(body) };
+        withHas.parts.push(has);
+      }
+      return this.querySelectorAll(withHas, root);
+    }
+
     if (!(root as any)['querySelectorAll'])
       throw this.createStacklessError('Node is not queryable.');
+
+    if (selector.capture !== undefined) {
+      // We should have handled the capture above.
+      throw this.createStacklessError('Internal error: there should not be a capture in the selector.');
+    }
+
     this._evaluator.begin();
     try {
-      const result = this._querySelectorRecursively([{ element: root as Element, capture: undefined }], selector, 0, new Map());
+      let roots = new Set<Element>([root as Element]);
-      const set = new Set<Element>();
+      for (const part of selector.parts) {
-      for (const r of result)
+        if (part.name === 'nth') {
-        set.add(r.capture || r.element);
+          roots = this._queryNth(roots, part);
-      return [...set];
+        } else {
+          const next = new Set<Element>();
+          for (const root of roots) {
+            const all = this._queryEngineAll(part, root);
+            for (const one of all)
+              next.add(one);
+          }
+          roots = next;
+        }
+      }
+      return [...roots];
     } finally {
       this._evaluator.end();
     }
   }
 
   private _queryEngineAll(part: ParsedSelectorPart, root: SelectorRoot): Element[] {
-    return this._engines.get(part.name)!.queryAll(root, part.body);
+    const result = this._engines.get(part.name)!.queryAll(root, part.body);
+    for (const element of result) {
+      if (!('nodeName' in element))
+        throw this.createStacklessError(`Expected a Node but got ${Object.prototype.toString.call(element)}`);
+    }
+    return result;
   }
 
   private _createAttributeEngine(attribute: string, shadow: boolean): SelectorEngine {
@@ -304,6 +276,15 @@ export class InjectedScript {
     return { queryAll };
   }
 
+  private _createVisibleEngine(): SelectorEngineV2 {
+    const queryAll = (root: SelectorRoot, body: string) => {
+      if (root.nodeType !== 1 /* Node.ELEMENT_NODE */)
+        return [];
+      return isVisible(root as Element) === Boolean(body) ? [root as Element] : [];
+    };
+    return { queryAll };
+  }
+
   extend(source: string, params: any): any {
     const constrFunction = globalThis.eval(`
     (() => {

tests/page/selectors-misc.spec.ts

@@ -135,6 +135,31 @@ it('should work with nth=', async ({ page }) => {
   });
   const element = await promise;
   expect(await element.evaluate(e => e.id)).toBe('target3');
+
+  await page.setContent(`
+    <div>
+      <div>
+        <div>
+          <span>hi</span>
+          <span>hello</span>
+        </div>
+      </div>
+    </div>
+  `);
+  expect(await page.locator('div >> div >> span >> nth=1').textContent()).toBe('hello');
+});
+
+it('should work with strict mode and chaining', async ({ page }) => {
+  await page.setContent(`
+    <div>
+      <div>
+        <div>
+          <span>hi</span>
+        </div>
+      </div>
+    </div>
+  `);
+  expect(await page.locator('div >> div >> span').textContent()).toBe('hi');
 });
 
 it('should work with position selectors', async ({ page }) => {
@@ -373,3 +398,35 @@ it('should work with has=', async ({ page, server }) => {
   const error4 = await page.$(`div >> has="span!"`).catch(e => e);
   expect(error4.message).toContain('Unexpected token "!" while parsing selector "span!"');
 });
+
+it('chaining should work with large DOM @smoke', async ({ page, server }) => {
+  await page.evaluate(() => {
+    let last = document.body;
+    for (let i = 0; i < 100; i++) {
+      const e = document.createElement('div');
+      last.appendChild(e);
+      last = e;
+    }
+    const target = document.createElement('span');
+    target.textContent = 'Found me!';
+    last.appendChild(target);
+  });
+
+  // Naive implementation generates C(100, 9) ~= 1.9*10^12 entries.
+  const selectors = [
+    'div >> div >> div >> div >> div >> div >> div >> div >> span',
+    'div div div div div div div div span',
+    'div div >> div div >> div div >> div div >> span',
+  ];
+
+  const counts = [];
+  const times = [];
+  for (const selector of selectors) {
+    const time = Date.now();
+    counts.push(await page.$$eval(selector, els => els.length));
+    times.push({ selector, time: Date.now() - time });
+  }
+  expect(counts).toEqual([1, 1, 1]);
+  // Uncomment to see performance results.
+  // console.log(times);
+});