Skip to content

Commit

Permalink
feat(codegen): improve selector generation (#5364)
Browse files Browse the repository at this point in the history
- Snap to buttons, inputs, selects, etc.
- Try `<label>` selector in addition to the element.
- Use parent selectors when needed.
- Remove xpath fallback as it should be covered with css.
  • Loading branch information
dgozman committed Feb 9, 2021
1 parent b50c363 commit 0871a9c
Show file tree
Hide file tree
Showing 5 changed files with 411 additions and 344 deletions.
57 changes: 32 additions & 25 deletions src/server/injected/injectedScript.ts
Expand Up @@ -40,7 +40,7 @@ export type InjectedScriptPoll<T> = {

export class InjectedScript {
private _enginesV1: Map<string, SelectorEngine>;
private _evaluator: SelectorEvaluatorImpl;
_evaluator: SelectorEvaluatorImpl;

constructor(customEngines: { name: string, engine: SelectorEngine}[]) {
this._enginesV1 = new Map();
Expand Down Expand Up @@ -75,9 +75,12 @@ export class InjectedScript {
querySelector(selector: ParsedSelector, root: Node): Element | undefined {
if (!(root as any)['querySelector'])
throw new Error('Node is not queryable.');
const result = this._querySelectorRecursively(root as SelectorRoot, selector, 0);
this._evaluator.clearCaches();
return result;
this._evaluator.begin();
try {
return this._querySelectorRecursively(root as SelectorRoot, selector, 0);
} finally {
this._evaluator.end();
}
}

private _querySelectorRecursively(root: SelectorRoot, selector: ParsedSelector, index: number): Element | undefined {
Expand All @@ -95,30 +98,34 @@ export class InjectedScript {
querySelectorAll(selector: ParsedSelector, root: Node): Element[] {
if (!(root as any)['querySelectorAll'])
throw new Error('Node is not queryable.');
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);
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;
}
set = newSet;
}
let result = Array.from(set) as Element[];
if (partsToCheckOne.length) {
const partial = { parts: partsToCheckOne };
result = result.filter(e => !!this._querySelectorRecursively(e, partial, 0));
let result = Array.from(set) as Element[];
if (partsToCheckOne.length) {
const partial = { parts: partsToCheckOne };
result = result.filter(e => !!this._querySelectorRecursively(e, partial, 0));
}
return result;
} finally {
this._evaluator.end();
}
this._evaluator.clearCaches();
return result;
}

private _queryEngine(part: ParsedSelectorPart, root: SelectorRoot): Element | undefined {
Expand Down
104 changes: 61 additions & 43 deletions src/server/injected/selectorEvaluator.ts
Expand Up @@ -45,6 +45,7 @@ export class SelectorEvaluatorImpl implements SelectorEvaluator {
private _cacheQuerySimple: QueryCache = new Map();
_cacheText = new Map<Element | ShadowRoot, string>();
private _scoreMap: Map<Element, number> | undefined;
private _retainCacheCounter = 0;

constructor(extraEngines: Map<string, SelectorEngine>) {
for (const [name, engine] of extraEngines)
Expand Down Expand Up @@ -75,16 +76,23 @@ export class SelectorEvaluatorImpl implements SelectorEvaluator {
throw new Error(`Please keep customCSSNames in sync with evaluator engines`);
}

clearCaches() {
this._cacheQueryCSS.clear();
this._cacheMatches.clear();
this._cacheQuery.clear();
this._cacheMatchesSimple.clear();
this._cacheMatchesParents.clear();
this._cacheCallMatches.clear();
this._cacheCallQuery.clear();
this._cacheQuerySimple.clear();
this._cacheText.clear();
begin() {
++this._retainCacheCounter;
}

end() {
--this._retainCacheCounter;
if (!this._retainCacheCounter) {
this._cacheQueryCSS.clear();
this._cacheMatches.clear();
this._cacheQuery.clear();
this._cacheMatchesSimple.clear();
this._cacheMatchesParents.clear();
this._cacheCallMatches.clear();
this._cacheCallQuery.clear();
this._cacheQuerySimple.clear();
this._cacheText.clear();
}
}

private _cached<T>(cache: QueryCache, main: any, rest: any[], cb: () => T): T {
Expand All @@ -109,43 +117,53 @@ export class SelectorEvaluatorImpl implements SelectorEvaluator {

matches(element: Element, s: Selector, context: QueryContext): boolean {
const selector = this._checkSelector(s);
return this._cached<boolean>(this._cacheMatches, element, [selector, context], () => {
if (Array.isArray(selector))
return this._matchesEngine(isEngine, element, selector, context);
if (!this._matchesSimple(element, selector.simples[selector.simples.length - 1].selector, context))
return false;
return this._matchesParents(element, selector, selector.simples.length - 2, context);
});
this.begin();
try {
return this._cached<boolean>(this._cacheMatches, element, [selector, context], () => {
if (Array.isArray(selector))
return this._matchesEngine(isEngine, element, selector, context);
if (!this._matchesSimple(element, selector.simples[selector.simples.length - 1].selector, context))
return false;
return this._matchesParents(element, selector, selector.simples.length - 2, context);
});
} finally {
this.end();
}
}

query(context: QueryContext, s: any): Element[] {
const selector = this._checkSelector(s);
return this._cached<Element[]>(this._cacheQuery, selector, [context], () => {
if (Array.isArray(selector))
return this._queryEngine(isEngine, context, selector);

// query() recursively calls itself, so we set up a new map for this particular query() call.
const previousScoreMap = this._scoreMap;
this._scoreMap = new Map();
let elements = this._querySimple(context, selector.simples[selector.simples.length - 1].selector);
elements = elements.filter(element => this._matchesParents(element, selector, selector.simples.length - 2, context));
if (this._scoreMap.size) {
elements.sort((a, b) => {
const aScore = this._scoreMap!.get(a);
const bScore = this._scoreMap!.get(b);
if (aScore === bScore)
return 0;
if (aScore === undefined)
return 1;
if (bScore === undefined)
return -1;
return aScore - bScore;
});
}
this._scoreMap = previousScoreMap;
this.begin();
try {
return this._cached<Element[]>(this._cacheQuery, selector, [context], () => {
if (Array.isArray(selector))
return this._queryEngine(isEngine, context, selector);

// query() recursively calls itself, so we set up a new map for this particular query() call.
const previousScoreMap = this._scoreMap;
this._scoreMap = new Map();
let elements = this._querySimple(context, selector.simples[selector.simples.length - 1].selector);
elements = elements.filter(element => this._matchesParents(element, selector, selector.simples.length - 2, context));
if (this._scoreMap.size) {
elements.sort((a, b) => {
const aScore = this._scoreMap!.get(a);
const bScore = this._scoreMap!.get(b);
if (aScore === bScore)
return 0;
if (aScore === undefined)
return 1;
if (bScore === undefined)
return -1;
return aScore - bScore;
});
}
this._scoreMap = previousScoreMap;

return elements;
});
return elements;
});
} finally {
this.end();
}
}

_markScore(element: Element, score: number) {
Expand Down Expand Up @@ -458,7 +476,7 @@ function shouldSkipForTextMatching(element: Element | ShadowRoot) {
return element.nodeName === 'SCRIPT' || element.nodeName === 'STYLE' || document.head && document.head.contains(element);
}

function elementText(evaluator: SelectorEvaluatorImpl, root: Element | ShadowRoot): string {
export function elementText(evaluator: SelectorEvaluatorImpl, root: Element | ShadowRoot): string {
let value = evaluator._cacheText.get(root);
if (value === undefined) {
value = '';
Expand Down

0 comments on commit 0871a9c

Please sign in to comment.