From d45585fe3bc44bd19099e0d7dea56d24f9fafced Mon Sep 17 00:00:00 2001
From: Kiko Beats <josefrancisco.verdu@gmail.com>
Date: Sat, 14 Sep 2024 15:38:58 +0200
Subject: [PATCH] perf(function): avoid page template if it's not used (#593)

---
 packages/function/package.json     |  4 ++-
 packages/function/src/function.js  | 15 ++-------
 packages/function/src/template.js  | 41 +++++++++++++++++++++++
 packages/function/test/template.js | 53 ++++++++++++++++++++++++++++++
 4 files changed, 99 insertions(+), 14 deletions(-)
 create mode 100644 packages/function/src/template.js
 create mode 100644 packages/function/test/template.js

diff --git a/packages/function/package.json b/packages/function/package.json
index 953a6f5fb..1f63fe858 100644
--- a/packages/function/package.json
+++ b/packages/function/package.json
@@ -35,12 +35,14 @@
   },
   "devDependencies": {
     "@browserless/test": "^10.6.0",
+    "acorn": "~8.12.1",
+    "acorn-walk": "~8.3.4",
     "ava": "5",
     "browserless": "^10.6.0",
     "lodash": "latest"
   },
   "engines": {
-    "node": ">= 12"
+    "node": ">= 20"
   },
   "files": [
     "src"
diff --git a/packages/function/src/function.js b/packages/function/src/function.js
index 1d0c03ddb..8e223e008 100644
--- a/packages/function/src/function.js
+++ b/packages/function/src/function.js
@@ -1,21 +1,10 @@
 'use strict'
 
 const isolatedFunction = require('isolated-function')
-
-const createFn = code => `
-async (url, browserWSEndpoint, opts) => {
-  const puppeteer = require('@cloudflare/puppeteer')
-  const browser = await puppeteer.connect({ browserWSEndpoint })
-  const page = (await browser.pages())[1]
-  try {
-    return await (${code})({ page, ...opts })
-  } finally {
-    await browser.disconnect()
-  }
-}`
+const template = require('./template')
 
 module.exports = async ({ url, code, vmOpts, browserWSEndpoint, ...opts }) => {
-  const [fn, teardown] = isolatedFunction(createFn(code), { ...vmOpts, throwError: false })
+  const [fn, teardown] = isolatedFunction(template(code), { ...vmOpts, throwError: false })
   const result = await fn(url, browserWSEndpoint, opts)
   await teardown()
   return result
diff --git a/packages/function/src/template.js b/packages/function/src/template.js
new file mode 100644
index 000000000..9c7a296b3
--- /dev/null
+++ b/packages/function/src/template.js
@@ -0,0 +1,41 @@
+'use strict'
+
+const walk = require('acorn-walk')
+const acorn = require('acorn')
+
+module.exports = code => {
+  const ast = acorn.parse(code, { ecmaVersion: 2023, sourceType: 'module' })
+
+  let isUsingPage = false
+
+  walk.simple(ast, {
+    ObjectPattern (node) {
+      node.properties.forEach(prop => {
+        if (prop.type === 'Property' && prop.key.name === 'page') {
+          isUsingPage = true
+        }
+        if (prop.type === 'RestElement' && prop.argument.name === 'page') {
+          isUsingPage = true
+        }
+      })
+    },
+    MemberExpression (node) {
+      if (node.property.name === 'page' || node.property.value === 'page') {
+        isUsingPage = true
+      }
+    }
+  })
+
+  if (!isUsingPage) return `async (url, _, opts) => (${code})(opts)`
+  return `
+    async (url, browserWSEndpoint, opts) => {
+      const puppeteer = require('@cloudflare/puppeteer')
+      const browser = await puppeteer.connect({ browserWSEndpoint })
+      const page = (await browser.pages())[1]
+      try {
+        return await (${code})({ page, ...opts })
+      } finally {
+        await browser.disconnect()
+      }
+    }`
+}
diff --git a/packages/function/test/template.js b/packages/function/test/template.js
new file mode 100644
index 000000000..f2c8b1eee
--- /dev/null
+++ b/packages/function/test/template.js
@@ -0,0 +1,53 @@
+'use strict'
+
+const test = require('ava')
+
+const template = require('../src/template')
+
+test('use a simplified template if page is not used', t => {
+  {
+    const code = () => {
+      function isStrict () {
+        return !this
+      }
+      return isStrict()
+    }
+
+    t.is(template(code.toString()).includes('puppeteer'), false)
+  }
+  {
+    const code = () => ({ page: 'title' })
+    t.is(template(code.toString()).includes('puppeteer'), false)
+  }
+  {
+    const code = () => 'page'
+    t.is(template(code.toString()).includes('puppeteer'), false)
+  }
+})
+
+test('require puppeteer if page is used', t => {
+  {
+    const code = ({ page }) => page.title()
+    t.is(template(code.toString()).includes('puppeteer'), true)
+  }
+  {
+    const code = ({ page: p }) => p.title()
+    t.is(template(code.toString()).includes('puppeteer'), true)
+  }
+  {
+    const code = obj => obj.page.title()
+    t.is(template(code.toString()).includes('puppeteer'), true)
+  }
+  {
+    const code = obj => (() => obj.page.title())()
+    t.is(template(code.toString()).includes('puppeteer'), true)
+  }
+  {
+    const code = ({ ...page }) => page.title
+    t.is(template(code.toString()).includes('puppeteer'), true)
+  }
+  {
+    const code = obj => obj.page.title()
+    t.is(template(code.toString()).includes('puppeteer'), true)
+  }
+})