chore(role): cache element list by role (#29130)

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

@@ -16,9 +16,10 @@
 
 import type { SelectorEngine, SelectorRoot } from './selectorEngine';
 import { matchesAttributePart } from './selectorUtils';
-import { beginAriaCaches, endAriaCaches, getAriaChecked, getAriaDisabled, getAriaExpanded, getAriaLevel, getAriaPressed, getAriaRole, getAriaSelected, getElementAccessibleName, isElementHiddenForAria, kAriaCheckedRoles, kAriaExpandedRoles, kAriaLevelRoles, kAriaPressedRoles, kAriaSelectedRoles } from './roleUtils';
+import { beginAriaCaches, endAriaCaches, getAriaChecked, getAriaDisabled, getAriaExpanded, getAriaLevel, getAriaPressed, getAriaSelected, getElementAccessibleName, getElementsByRole, isElementHiddenForAria, kAriaCheckedRoles, kAriaExpandedRoles, kAriaLevelRoles, kAriaPressedRoles, kAriaSelectedRoles } from './roleUtils';
 import { parseAttributeSelector, type AttributeSelectorPart, type AttributeSelectorOperator } from '../../utils/isomorphic/selectorParser';
 import { normalizeWhiteSpace } from '../../utils/isomorphic/stringUtils';
+import { isInsideScope } from './domUtils';
 
 type RoleEngineOptions = {
   role: string;
@@ -125,26 +126,27 @@ function validateAttributes(attrs: AttributeSelectorPart[], role: string): RoleE
 }
 
 function queryRole(scope: SelectorRoot, options: RoleEngineOptions, internal: boolean): Element[] {
-  const result: Element[] = [];
-  const match = (element: Element) => {
-    if (getAriaRole(element) !== options.role)
-      return;
+  const doc = scope.nodeType === 9 /* Node.DOCUMENT_NODE */ ? scope as Document : scope.ownerDocument;
+  const elements = doc ? getElementsByRole(doc, options.role) : [];
+  return elements.filter(element => {
+    if (!isInsideScope(scope, element))
+      return false;
     if (options.selected !== undefined && getAriaSelected(element) !== options.selected)
-      return;
+      return false;
     if (options.checked !== undefined && getAriaChecked(element) !== options.checked)
-      return;
+      return false;
     if (options.pressed !== undefined && getAriaPressed(element) !== options.pressed)
-      return;
+      return false;
     if (options.expanded !== undefined && getAriaExpanded(element) !== options.expanded)
-      return;
+      return false;
     if (options.level !== undefined && getAriaLevel(element) !== options.level)
-      return;
+      return false;
     if (options.disabled !== undefined && getAriaDisabled(element) !== options.disabled)
-      return;
+      return false;
     if (!options.includeHidden) {
       const isHidden = isElementHiddenForAria(element);
       if (isHidden)
-        return;
+        return false;
     }
     if (options.name !== undefined) {
       // Always normalize whitespace in the accessible name.
@@ -155,25 +157,10 @@ function queryRole(scope: SelectorRoot, options: RoleEngineOptions, internal: bo
       if (internal && !options.exact && options.nameOp === '=')
         options.nameOp = '*=';
       if (!matchesAttributePart(accessibleName, { name: '', jsonPath: [], op: options.nameOp || '=', value: options.name, caseSensitive: !!options.exact }))
-        return;
+        return false;
     }
-    result.push(element);
-  };
-
-  const query = (root: Element | ShadowRoot | Document) => {
-    const shadows: ShadowRoot[] = [];
-    if ((root as Element).shadowRoot)
-      shadows.push((root as Element).shadowRoot!);
-    for (const element of root.querySelectorAll('*')) {
-      match(element);
-      if (element.shadowRoot)
-        shadows.push(element.shadowRoot);
-    }
-    shadows.forEach(query);
-  };
-
-  query(scope);
-  return result;
+    return true;
+  });
 }
 
 export function createRoleEngine(internal: boolean): SelectorEngine {

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

@@ -845,11 +845,51 @@ function getAccessibleNameFromAssociatedLabels(labels: Iterable<HTMLLabelElement
   })).filter(accessibleName => !!accessibleName).join(' ');
 }
 
+export function getElementsByRole(document: Document, role: string): Element[] {
+  if (document === cacheElementsByRoleDocument)
+    return cacheElementsByRole!.get(role) || [];
+  const map = calculateElementsByRoleMap(document);
+  if (cachesCounter) {
+    cacheElementsByRoleDocument = document;
+    cacheElementsByRole = map;
+  }
+  return map.get(role) || [];
+}
+
+function calculateElementsByRoleMap(document: Document) {
+  const result = new Map<string, Element[]>();
+
+  const visit = (root: Element | ShadowRoot | Document) => {
+    const shadows: ShadowRoot[] = [];
+    if ((root as Element).shadowRoot)
+      shadows.push((root as Element).shadowRoot!);
+    for (const element of root.querySelectorAll('*')) {
+      const role = getAriaRole(element);
+      if (role) {
+        let list = result.get(role);
+        if (!list) {
+          list = [];
+          result.set(role, list);
+        }
+        list.push(element);
+      }
+      if (element.shadowRoot)
+        shadows.push(element.shadowRoot);
+    }
+    shadows.forEach(visit);
+  };
+  visit(document);
+
+  return result;
+}
+
 let cacheAccessibleName: Map<Element, string> | undefined;
 let cacheAccessibleNameHidden: Map<Element, string> | undefined;
 let cacheIsHidden: Map<Element, boolean> | undefined;
 let cachePseudoContentBefore: Map<Element, string> | undefined;
 let cachePseudoContentAfter: Map<Element, string> | undefined;
+let cacheElementsByRole: Map<string, Element[]> | undefined;
+let cacheElementsByRoleDocument: Document | undefined;
 let cachesCounter = 0;
 
 export function beginAriaCaches() {
@@ -868,5 +908,7 @@ export function endAriaCaches() {
     cacheIsHidden = undefined;
     cachePseudoContentBefore = undefined;
     cachePseudoContentAfter = undefined;
+    cacheElementsByRole = undefined;
+    cacheElementsByRoleDocument = undefined;
   }
 }