fix($$): use utility context when possible (#3870)

This avoids the typical issue of overridden bulitins,
trading it for performance of one by one node adoptions.

src/server/dom.ts

@@ -552,7 +552,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
   }
 
   async $$(selector: string): Promise<ElementHandle<Element>[]> {
-    return this._page.selectors._queryAll(this._context.frame, selector, this);
+    return this._page.selectors._queryAll(this._context.frame, selector, this, true /* adoptToMain */);
   }
 
   async _$evalExpression(selector: string, expression: string, isFunction: boolean, arg: any): Promise<any> {

src/server/firefox/ffPage.ts

@@ -488,7 +488,7 @@ export class FFPage implements PageDelegate {
     const parent = frame.parentFrame();
     if (!parent)
       throw new Error('Frame has been detached.');
-    const handles = await this._page.selectors._queryAll(parent, 'iframe', undefined, true /* allowUtilityContext */);
+    const handles = await this._page.selectors._queryAll(parent, 'iframe', undefined);
     const items = await Promise.all(handles.map(async handle => {
       const frame = await handle.contentFrame().catch(e => null);
       return { handle, frame };

src/server/frames.ts

@@ -590,7 +590,7 @@ export class Frame extends EventEmitter {
   }
 
   async $$(selector: string): Promise<dom.ElementHandle<Element>[]> {
-    return this._page.selectors._queryAll(this, selector);
+    return this._page.selectors._queryAll(this, selector, undefined, true /* adoptToMain */);
   }
 
   async content(): Promise<string> {

src/server/selectors.ts

@@ -68,11 +68,7 @@ export class Selectors {
       return null;
     }
     const mainContext = await frame._mainContext();
-    if (elementHandle._context === mainContext)
+    return this._adoptIfNeeded(elementHandle, mainContext);
-      return elementHandle;
-    const adopted = frame._page._delegate.adoptElementHandle(elementHandle, mainContext);
-    elementHandle.dispose();
-    return adopted;
   }
 
   async _queryArray(frame: frames.Frame, selector: string, scope?: dom.ElementHandle): Promise<js.JSHandle<Element[]>> {
@@ -85,9 +81,9 @@ export class Selectors {
     return arrayHandle;
   }
 
-  async _queryAll(frame: frames.Frame, selector: string, scope?: dom.ElementHandle, allowUtilityContext?: boolean): Promise<dom.ElementHandle<Element>[]> {
+  async _queryAll(frame: frames.Frame, selector: string, scope?: dom.ElementHandle, adoptToMain?: boolean): Promise<dom.ElementHandle<Element>[]> {
     const info = this._parseSelector(selector);
-    const context = await frame._context(allowUtilityContext ? info.world : 'main');
+    const context = await frame._context(info.world);
     const injectedScript = await context.injectedScript();
     const arrayHandle = await injectedScript.evaluateHandle((injected, { parsed, scope }) => {
       return injected.querySelectorAll(parsed, scope || document);
@@ -95,15 +91,27 @@ export class Selectors {
 
     const properties = await arrayHandle.getProperties();
     arrayHandle.dispose();
-    const result: dom.ElementHandle<Element>[] = [];
+
+    // Note: adopting elements one by one may be slow. If we encounter the issue here,
+    // we might introduce 'useMainContext' option or similar to speed things up.
+    const targetContext = adoptToMain ? await frame._mainContext() : context;
+    const result: Promise<dom.ElementHandle<Element>>[] = [];
     for (const property of properties.values()) {
       const elementHandle = property.asElement() as dom.ElementHandle<Element>;
       if (elementHandle)
-        result.push(elementHandle);
+        result.push(this._adoptIfNeeded(elementHandle, targetContext));
       else
         property.dispose();
     }
-    return result;
+    return Promise.all(result);
+  }
+
+  private async _adoptIfNeeded<T extends Node>(handle: dom.ElementHandle<T>, context: dom.FrameExecutionContext): Promise<dom.ElementHandle<T>> {
+    if (handle._context === context)
+      return handle;
+    const adopted = handle._page._delegate.adoptElementHandle(handle, context);
+    handle.dispose();
+    return adopted;
   }
 
   async _createSelector(name: string, handle: dom.ElementHandle<Element>): Promise<string | undefined> {

src/server/webkit/wkPage.ts

@@ -854,7 +854,7 @@ export class WKPage implements PageDelegate {
     const parent = frame.parentFrame();
     if (!parent)
       throw new Error('Frame has been detached.');
-    const handles = await this._page.selectors._queryAll(parent, 'iframe', undefined, true /* allowUtilityContext */);
+    const handles = await this._page.selectors._queryAll(parent, 'iframe', undefined);
     const items = await Promise.all(handles.map(async handle => {
       const frame = await handle.contentFrame().catch(e => null);
       return { handle, frame };

test/queryselector.spec.ts

@@ -114,3 +114,15 @@ it('xpath should return multiple elements', async ({page, server}) => {
   const elements = await page.$$('xpath=/html/body/div');
   expect(elements.length).toBe(2);
 });
+
+it('$$ should work with bogus Array.from', async ({page, server}) => {
+  await page.setContent('<div>hello</div><div></div>');
+  const div1 = await page.evaluateHandle(() => {
+    Array.from = () => [];
+    return document.querySelector('div');
+  });
+  const elements = await page.$$('div');
+  expect(elements.length).toBe(2);
+  // Check that element handle is functional and belongs to the main world.
+  expect(await elements[0].evaluate((div, div1) => div === div1, div1)).toBe(true);
+});