chore: BFS nodes, simplify querying (#7861)

src/server/injected/injectedScript.ts

@@ -42,10 +42,14 @@ export type ElementStateWithoutStable = 'visible' | 'hidden' | 'enabled' | 'disa
 export type ElementState = ElementStateWithoutStable | 'stable';
 
 export interface SelectorEngineV2 {
-  query?(root: SelectorRoot, body: any): Element | undefined;
   queryAll(root: SelectorRoot, body: any): Element[];
 }
 
+export type ElementMatch = {
+  element: Element;
+  capture: Element | undefined;
+};
+
 export class InjectedScript {
   private _engines: Map<string, SelectorEngineV2>;
   _evaluator: SelectorEvaluatorImpl;
@@ -69,6 +73,8 @@ export class InjectedScript {
     this._engines.set('data-test', this._createAttributeEngine('data-test', true));
     this._engines.set('data-test:light', this._createAttributeEngine('data-test', false));
     this._engines.set('css', this._createCSSEngine());
+    this._engines.set('_first', { queryAll: () => [] });
+    this._engines.set('_visible', { queryAll: () => [] });
 
     for (const { name, engine } of customEngines)
       this._engines.set(name, engine);
@@ -91,30 +97,47 @@ export class InjectedScript {
       throw new Error('Node is not queryable.');
     this._evaluator.begin();
     try {
-      return this._querySelectorRecursively(root as SelectorRoot, selector, strict, 0);
+      const result = this._querySelectorRecursively([{ element: root as Element, capture: undefined }], selector, 0, new Map());
+      if (strict && result.length > 1)
+        throw new Error(`strict mode violation: selector resolved to ${result.length} elements.`);
+      return result[0]?.capture || result[0]?.element;
     } finally {
       this._evaluator.end();
     }
   }
 
-  private _querySelectorRecursively(root: SelectorRoot, selector: ParsedSelector, strict: boolean, index: number): Element | undefined {
-    const current = selector.parts[index];
-    if (index === selector.parts.length - 1) {
-      if (strict) {
-        const all = this._queryEngineAll(current, root);
-        if (all.length > 1)
-          throw new Error(`strict mode violation: selector resolved to ${all.length} elements.`);
-        return all[0];
-      } else {
-        return this._queryEngine(current, root);
+  private _querySelectorRecursively(roots: ElementMatch[], selector: ParsedSelector, index: number, queryCache: Map<Element, Element[][]>): ElementMatch[] {
+    if (index === selector.parts.length)
+      return roots;
+
+    if (selector.parts[index].name === '_first')
+      return roots.slice(0, 1);
+
+    if (selector.parts[index].name === '_visible') {
+      const visible = Boolean(selector.parts[index].body);
+      return roots.filter(match => visible === isVisible(match.element));
     }
+
+    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);
       }
-    const all = this._queryEngineAll(current, root);
-    for (const next of all) {
-      const result = this._querySelectorRecursively(next, selector, strict, index + 1);
-      if (result)
-        return selector.capture === index ? next : result;
+      let all = queryResults[index];
+      if (!all) {
+        all = this._queryEngineAll(selector.parts[index], root.element);
+        queryResults[index] = all;
       }
+
+      for (const element of all)
+        result.push({ element, capture });
+    }
+    return this._querySelectorRecursively(result, selector, index + 1, queryCache);
   }
 
   querySelectorAll(selector: ParsedSelector, root: Node): Element[] {
@@ -122,42 +145,16 @@ export class InjectedScript {
       throw new Error('Node is not queryable.');
     this._evaluator.begin();
     try {
-      const capture = selector.capture === undefined ? selector.parts.length - 1 : selector.capture;
-      // Query all elements up to the capture.
-      const partsToQueryAll = selector.parts.slice(0, capture + 1);
-      // Check they have a descendant matching everything after the capture.
-      const partsToCheckOne = selector.parts.slice(capture + 1);
-      let set = new Set<SelectorRoot>([ root as SelectorRoot ]);
-      for (const part of partsToQueryAll) {
-        const newSet = new Set<Element>();
-        for (const prev of set) {
-          for (const next of this._queryEngineAll(part, prev)) {
-            if (newSet.has(next))
-              continue;
-            newSet.add(next);
-          }
-        }
-        set = newSet;
-      }
-      let result = [...set] as Element[];
-      if (partsToCheckOne.length) {
-        const partial = { parts: partsToCheckOne };
-        result = result.filter(e => !!this._querySelectorRecursively(e, partial, false, 0));
-      }
-      return result;
+      const result = this._querySelectorRecursively([{ element: root as Element, capture: undefined }], selector, 0, new Map());
+      const set = new Set<Element>();
+      for (const r of result)
+        set.add(r.capture || r.element);
+      return [...set];
     } finally {
       this._evaluator.end();
     }
   }
 
-  private _queryEngine(part: ParsedSelectorPart, root: SelectorRoot): Element | undefined {
-    const engine = this._engines.get(part.name)!;
-    if (engine.query)
-      return engine.query(root, part.body);
-    else
-      return engine.queryAll(root, part.body)[0];
-  }
-
   private _queryEngineAll(part: ParsedSelectorPart, root: SelectorRoot): Element[] {
     return this._engines.get(part.name)!.queryAll(root, part.body);
   }
@@ -184,7 +181,7 @@ export class InjectedScript {
   }
 
   private _createTextEngine(shadow: boolean): SelectorEngine {
-    const queryList = (root: SelectorRoot, selector: string, single: boolean): Element[] => {
+    const queryList = (root: SelectorRoot, selector: string): Element[] => {
       const { matcher, kind } = createTextMatcher(selector);
       const result: Element[] = [];
       let lastDidNotMatchSelf: Element | null = null;
@@ -198,26 +195,19 @@ export class InjectedScript {
           lastDidNotMatchSelf = element;
         if (matches === 'self' || (matches === 'selfAndChildren' && kind === 'strict'))
           result.push(element);
-        return single && result.length > 0;
       };
 
-      if (root.nodeType === Node.ELEMENT_NODE && appendElement(root as Element))
-        return result;
+      if (root.nodeType === Node.ELEMENT_NODE)
+        appendElement(root as Element);
       const elements = this._evaluator._queryCSS({ scope: root as Document | Element, pierceShadow: shadow }, '*');
-      for (const element of elements) {
-        if (appendElement(element))
-          return result;
-      }
+      for (const element of elements)
+        appendElement(element);
       return result;
     };
 
     return {
-      query: (root: SelectorRoot, selector: string): Element | undefined => {
-        return queryList(root, selector, true)[0];
-      },
-
       queryAll: (root: SelectorRoot, selector: string): Element[] => {
-        return queryList(root, selector, false);
+        return queryList(root, selector);
       }
     };
   }

src/server/injected/selectorEngine.ts

@@ -17,7 +17,5 @@
 export type SelectorRoot = Element | ShadowRoot | Document;
 
 export interface SelectorEngine {
-  query?(root: SelectorRoot, selector: string): Element | undefined;
-
   queryAll(root: SelectorRoot, selector: string): Element[];
 }

src/server/injected/xpathSelectorEngine.ts

@@ -17,19 +17,6 @@
 import { SelectorEngine, SelectorRoot } from './selectorEngine';
 
 export const XPathEngine: SelectorEngine = {
-  query(root: SelectorRoot, selector: string): Element | undefined {
-    if (selector.startsWith('/'))
-      selector = '.' + selector;
-    const document = root instanceof Document ? root : root.ownerDocument;
-    if (!document)
-      return;
-    const it = document.evaluate(selector, root, null, XPathResult.ORDERED_NODE_ITERATOR_TYPE);
-    for (let node = it.iterateNext(); node; node = it.iterateNext()) {
-      if (node.nodeType === Node.ELEMENT_NODE)
-        return node as Element;
-    }
-  },
-
   queryAll(root: SelectorRoot, selector: string): Element[] {
     if (selector.startsWith('/'))
       selector = '.' + selector;

src/server/selectors.ts

@@ -125,14 +125,12 @@ export class Selectors {
     const parsed = parseSelector(selector);
     let needsMainWorld = false;
     for (const part of parsed.parts) {
-      if (!Array.isArray(part)) {
       const custom = this._engines.get(part.name);
       if (!custom && !this._builtinEngines.has(part.name))
         throw new Error(`Unknown engine "${part.name}" while parsing selector ${selector}`);
       if (custom && !custom.contentScript)
         needsMainWorld = true;
     }
-    }
     return {
       parsed,
       selector,

tests/page/selectors-css.spec.ts

@@ -143,6 +143,16 @@ it('should keep dom order with comma separated list', async ({page}) => {
   expect(await page.$$eval(`css=section >> *css=div,span >> css=y`, els => els.map(e => e.nodeName).join(','))).toBe('SPAN,DIV');
 });
 
+it('should return multiple captures for the same node', async ({page}) => {
+  await page.setContent(`<div><div><div><span></span></div></div></div>`);
+  expect(await page.$$eval(`*css=div >> span`, els => els.map(e => e.nodeName).join(','))).toBe('DIV,DIV,DIV');
+});
+
+it('should return multiple captures when going up the hierarchy', async ({page}) => {
+  await page.setContent(`<section>Hello<ul><li></li><li></li></ul></section>`);
+  expect(await page.$$eval(`*css=li >> ../.. >> text=Hello`, els => els.map(e => e.nodeName).join(','))).toBe('LI,LI');
+});
+
 it('should work with comma separated list in various positions', async ({page}) => {
   await page.setContent(`<section><span><div><x></x><y></y></div></span></section>`);
   expect(await page.$$eval(`css=span,div >> css=x,y`, els => els.map(e => e.nodeName).join(','))).toBe('X,Y');