From 3c0a6091fe96044e9dd84fbe5db3343339a88986 Mon Sep 17 00:00:00 2001
From: yoho <907415276@qq.com>
Date: Thu, 3 Mar 2022 21:46:25 +0800
Subject: [PATCH] fix: execute classic worker in dev mode (#7099)

---
 .prettierignore                               |  1 +
 .../worker/__tests__/worker.spec.ts           |  7 +-
 packages/playground/worker/classic-worker.js  | 29 ++++++
 packages/playground/worker/index.html         | 23 +++--
 .../worker/newUrl/classic-shared-worker.js    |  6 ++
 .../worker/newUrl/classic-worker.js           |  5 ++
 packages/playground/worker/public/classic.js  |  1 +
 packages/vite/src/node/importGlob.ts          | 52 ++++-------
 .../src/node/plugins/workerImportMetaUrl.ts   | 88 +++++++++++++++++--
 packages/vite/src/node/utils.ts               |  2 +
 10 files changed, 164 insertions(+), 50 deletions(-)
 create mode 100644 packages/playground/worker/classic-worker.js
 create mode 100644 packages/playground/worker/newUrl/classic-shared-worker.js
 create mode 100644 packages/playground/worker/newUrl/classic-worker.js
 create mode 100644 packages/playground/worker/public/classic.js

diff --git a/.prettierignore b/.prettierignore
index 1692b9d26cfa20..c624a3a21eecfe 100644
--- a/.prettierignore
+++ b/.prettierignore
@@ -10,3 +10,4 @@ pnpm-lock.yaml
 pnpm-workspace.yaml
 packages/playground/tsconfig-json-load-error/has-error/tsconfig.json
 packages/playground/html/invalid.html
+packages/playground/worker/classic-worker.js
diff --git a/packages/playground/worker/__tests__/worker.spec.ts b/packages/playground/worker/__tests__/worker.spec.ts
index 5992a37f933f2a..1c33e6ec68daf9 100644
--- a/packages/playground/worker/__tests__/worker.spec.ts
+++ b/packages/playground/worker/__tests__/worker.spec.ts
@@ -56,7 +56,7 @@ if (isBuild) {
   // assert correct files
   test('inlined code generation', async () => {
     const files = fs.readdirSync(assetsDir)
-    expect(files.length).toBe(6)
+    expect(files.length).toBe(8)
     const index = files.find((f) => f.includes('index'))
     const content = fs.readFileSync(path.resolve(assetsDir, index), 'utf-8')
     const worker = files.find((f) => f.includes('my-worker'))
@@ -88,3 +88,8 @@ if (isBuild) {
       })
   })
 }
+
+test('classic worker is run', async () => {
+  expect(await page.textContent('.classic-worker')).toMatch('A classic')
+  expect(await page.textContent('.classic-shared-worker')).toMatch('A classic')
+})
diff --git a/packages/playground/worker/classic-worker.js b/packages/playground/worker/classic-worker.js
new file mode 100644
index 00000000000000..bb6f9c3f49fc84
--- /dev/null
+++ b/packages/playground/worker/classic-worker.js
@@ -0,0 +1,29 @@
+// prettier-ignore
+function text(el, text) {
+  document.querySelector(el).textContent = text
+}
+
+const classicWorker = new Worker(
+  new URL('./newUrl/classic-worker.js', import.meta.url) /* , */  ,
+  // test comment
+
+)
+
+classicWorker.addEventListener('message', ({ data }) => {
+  text('.classic-worker', data)
+})
+classicWorker.postMessage('ping')
+
+const classicSharedWorker = new SharedWorker(
+  new URL('./newUrl/classic-shared-worker.js', import.meta.url),
+  {
+    type: 'classic'
+  }
+)
+classicSharedWorker.port.addEventListener('message', (ev) => {
+  text(
+    '.classic-shared-worker',
+    ev.data
+  )
+})
+classicSharedWorker.port.start()
diff --git a/packages/playground/worker/index.html b/packages/playground/worker/index.html
index be0b3becb5f94f..10207453712aa2 100644
--- a/packages/playground/worker/index.html
+++ b/packages/playground/worker/index.html
@@ -20,18 +20,25 @@
   <span class="tick-count">0</span>
 </div>
 
-<p>new Worker(new Url('path', import.meta.url))</p>
+<p>new Worker(new Url('path', import.meta.url), { type: 'module' })</p>
 <div class="worker-import-meta-url"></div>
 
-<p>new SharedWorker(new Url('path', import.meta.url))</p>
+<p>new SharedWorker(new Url('path', import.meta.url), { type: 'module' })</p>
 <div class="shared-worker-import-meta-url"></div>
 
+<p>new Worker(new Url('path', import.meta.url))</p>
+<div class="classic-worker"></div>
+
+<p>new Worker(new Url('path', import.meta.url), { type: 'classic' })</p>
+<div class="classic-shared-worker"></div>
+
 <script type="module">
   import myWorker from './my-worker?worker'
   import InlineWorker from './my-worker?worker&inline'
   import mySharedWorker from './my-shared-worker?sharedworker&name=shared'
   import TSOutputWorker from './possible-ts-output-worker?worker'
   import { mode } from './workerImport'
+  import './classic-worker'
 
   document.querySelector('.mode-true').textContent = mode
 
@@ -78,11 +85,12 @@
   function text(el, text) {
     document.querySelector(el).textContent = text
   }
-
+  const workerOptions = { type: 'module' }
   // url import worker
-  const w = new Worker(new URL('./newUrl/url-worker.js', import.meta.url), {
-    type: 'module'
-  })
+  const w = new Worker(
+    new URL('./newUrl/url-worker.js', import.meta.url),
+    /* @vite-ignore */ workerOptions
+  )
   w.addEventListener('message', (ev) =>
     text(
       '.worker-import-meta-url',
@@ -90,9 +98,12 @@
     )
   )
 
+  const genWorkerName = () => 'module'
   const w2 = new SharedWorker(
     new URL('./newUrl/url-shared-worker.js', import.meta.url),
     {
+      /* @vite-ignore */
+      name: genWorkerName(),
       type: 'module'
     }
   )
diff --git a/packages/playground/worker/newUrl/classic-shared-worker.js b/packages/playground/worker/newUrl/classic-shared-worker.js
new file mode 100644
index 00000000000000..462e49dfa8847f
--- /dev/null
+++ b/packages/playground/worker/newUrl/classic-shared-worker.js
@@ -0,0 +1,6 @@
+importScripts('/classic.js')
+
+self.onconnect = (event) => {
+  const port = event.ports[0]
+  port.postMessage(self.constant)
+}
diff --git a/packages/playground/worker/newUrl/classic-worker.js b/packages/playground/worker/newUrl/classic-worker.js
new file mode 100644
index 00000000000000..865810c76fbf85
--- /dev/null
+++ b/packages/playground/worker/newUrl/classic-worker.js
@@ -0,0 +1,5 @@
+importScripts('/classic.js')
+
+self.addEventListener('message', () => {
+  self.postMessage(self.constant)
+})
diff --git a/packages/playground/worker/public/classic.js b/packages/playground/worker/public/classic.js
new file mode 100644
index 00000000000000..c074cc050c2587
--- /dev/null
+++ b/packages/playground/worker/public/classic.js
@@ -0,0 +1 @@
+self.constant = 'A classic'
diff --git a/packages/vite/src/node/importGlob.ts b/packages/vite/src/node/importGlob.ts
index faa65da8cdafd9..41bbb90a1f8a37 100644
--- a/packages/vite/src/node/importGlob.ts
+++ b/packages/vite/src/node/importGlob.ts
@@ -8,7 +8,12 @@ import {
   preloadMarker
 } from './plugins/importAnalysisBuild'
 import { isCSSRequest } from './plugins/css'
-import { cleanUrl } from './utils'
+import {
+  cleanUrl,
+  singlelineCommentsRE,
+  multilineCommentsRE,
+  blankReplacer
+} from './utils'
 import type { RollupError } from 'rollup'
 
 export interface AssertOptions {
@@ -182,45 +187,20 @@ function lexGlobPattern(
         throw new Error('unknown import.meta.glob lexer state')
     }
   }
+  const noCommentCode = code
+    .slice(i + 1)
+    .replace(singlelineCommentsRE, blankReplacer)
+    .replace(multilineCommentsRE, blankReplacer)
+
+  const endIndex = noCommentCode.indexOf(')')
+  const options = noCommentCode.substring(0, endIndex)
+  const commaIndex = options.indexOf(',')
 
-  const endIndex = getEndIndex(code, i)
-  const options = code.substring(i + 1, endIndex)
-  const commaIndex = options.indexOf(`,`)
   let assert = {}
   if (commaIndex > -1) {
-    assert = JSON5.parse(options.substr(commaIndex + 1))
-  }
-  return [pattern, assert, endIndex + 1]
-}
-
-// reg without the 'g' option, only matches the first match
-const multilineCommentsRE = /\/\*(.|[\r\n])*?\*\//m
-const singlelineCommentsRE = /\/\/.*/
-
-function getEndIndex(code: string, i: number): number {
-  const findStart = i
-  const endIndex = code.indexOf(`)`, findStart)
-  const subCode = code.substring(findStart)
-
-  const matched =
-    subCode.match(singlelineCommentsRE) ?? subCode.match(multilineCommentsRE)
-  if (!matched) {
-    return endIndex
-  }
-
-  const str = matched[0]
-  const index = matched.index
-  if (!index) {
-    return endIndex
-  }
-
-  const commentStart = findStart + index
-  const commentEnd = commentStart + str.length
-  if (endIndex > commentStart && endIndex < commentEnd) {
-    return getEndIndex(code, commentEnd)
-  } else {
-    return endIndex
+    assert = JSON5.parse(options.substring(commaIndex + 1))
   }
+  return [pattern, assert, endIndex + i + 2]
 }
 
 function error(pos: number) {
diff --git a/packages/vite/src/node/plugins/workerImportMetaUrl.ts b/packages/vite/src/node/plugins/workerImportMetaUrl.ts
index 13c0ef2bb3883b..023b567ab9fb45 100644
--- a/packages/vite/src/node/plugins/workerImportMetaUrl.ts
+++ b/packages/vite/src/node/plugins/workerImportMetaUrl.ts
@@ -1,7 +1,9 @@
+import JSON5 from 'json5'
 import type { ResolvedConfig } from '../config'
 import type { Plugin } from '../plugin'
 import { getAssetHash, fileToUrl } from './asset'
 import {
+  blankReplacer,
   cleanUrl,
   injectQuery,
   multilineCommentsRE,
@@ -10,22 +12,90 @@ import {
 import path from 'path'
 import { bundleWorkerEntry } from './worker'
 import { parseRequest } from '../utils'
-import { ENV_PUBLIC_PATH } from '../constants'
+import { ENV_ENTRY, ENV_PUBLIC_PATH } from '../constants'
 import MagicString from 'magic-string'
+import type { ViteDevServer } from '..'
+
+type WorkerType = 'classic' | 'module' | 'ignore'
 
 const WORKER_FILE_ID = 'worker_url_file'
 
+function getWorkerType(
+  code: string,
+  noCommentsCode: string,
+  i: number
+): WorkerType {
+  const commaIndex = noCommentsCode.indexOf(',', i)
+  if (commaIndex === -1) {
+    return 'classic'
+  }
+  const endIndex = noCommentsCode.indexOf(')', i)
+
+  // need to find in comment code
+  let workerOptsString = code.substring(commaIndex + 1, endIndex)
+  const hasViteIgnore = /\/\*\s*@vite-ignore\s*\*\//.test(workerOptsString)
+  if (hasViteIgnore) {
+    return 'ignore'
+  }
+
+  // need to find in no comment code
+  workerOptsString = noCommentsCode.substring(commaIndex + 1, endIndex)
+  if (!workerOptsString.trim().length) {
+    return 'classic'
+  }
+
+  let workerOpts: { type: WorkerType } = { type: 'classic' }
+  try {
+    workerOpts = JSON5.parse(workerOptsString)
+  } catch (e) {
+    // can't parse by JSON5, so the worker options had unexpect char.
+    throw new Error(
+      'Vite is unable to parse the worker options as the value is not static.' +
+        'To ignore this error, please use /* @vite-ignore */ in the worker options.'
+    )
+  }
+
+  if (['classic', 'module'].includes(workerOpts.type)) {
+    return workerOpts.type
+  }
+  return 'classic'
+}
+
 export function workerImportMetaUrlPlugin(config: ResolvedConfig): Plugin {
   const isBuild = config.command === 'build'
+  let server: ViteDevServer
 
   return {
     name: 'vite:worker-import-meta-url',
 
+    configureServer(_server) {
+      server = _server
+    },
+
     async transform(code, id, options) {
       const query = parseRequest(id)
-      if (query && query[WORKER_FILE_ID] != null) {
+      if (query && query[WORKER_FILE_ID] != null && query['type'] != null) {
+        const workerType = query['type'] as WorkerType
+        let injectEnv = ''
+
+        if (workerType === 'classic') {
+          injectEnv = `importScripts('${ENV_PUBLIC_PATH}')\n`
+        } else if (workerType === 'module') {
+          injectEnv = `import '${ENV_PUBLIC_PATH}'\n`
+        } else if (workerType === 'ignore') {
+          if (isBuild) {
+            injectEnv = ''
+          } else if (server) {
+            // dynamic worker type we can't know how import the env
+            // so we copy /@vite/env code of server transform result into file header
+            const { moduleGraph } = server
+            const module = moduleGraph.getModuleById(ENV_ENTRY)
+            injectEnv = module?.transformResult?.code || ''
+          }
+        }
+
         return {
-          code: `import '${ENV_PUBLIC_PATH}'\n` + code
+          code: injectEnv + code
         }
       }
       if (
@@ -36,13 +106,12 @@ export function workerImportMetaUrlPlugin(config: ResolvedConfig): Plugin {
         const importMetaUrlRE =
           /\bnew\s+(Worker|SharedWorker)\s*\(\s*(new\s+URL\s*\(\s*('[^']+'|"[^"]+"|`[^`]+`)\s*,\s*import\.meta\.url\s*\))/g
         const noCommentsCode = code
-          .replace(multilineCommentsRE, (m) => ' '.repeat(m.length))
-          .replace(singlelineCommentsRE, (m) => ' '.repeat(m.length))
+          .replace(multilineCommentsRE, blankReplacer)
+          .replace(singlelineCommentsRE, blankReplacer)
         let match: RegExpExecArray | null
         let s: MagicString | null = null
         while ((match = importMetaUrlRE.exec(noCommentsCode))) {
           const { 0: allExp, 2: exp, 3: rawUrl, index } = match
-
           const urlIndex = allExp.indexOf(exp) + index
 
           if (options?.ssr) {
@@ -61,7 +130,11 @@ export function workerImportMetaUrlPlugin(config: ResolvedConfig): Plugin {
           }
 
           s ||= new MagicString(code)
-
+          const workerType = getWorkerType(
+            code,
+            noCommentsCode,
+            index + allExp.length
+          )
           const file = path.resolve(path.dirname(id), rawUrl.slice(1, -1))
           let url: string
           if (isBuild) {
@@ -80,6 +153,7 @@ export function workerImportMetaUrlPlugin(config: ResolvedConfig): Plugin {
           } else {
             url = await fileToUrl(cleanUrl(file), config, this)
             url = injectQuery(url, WORKER_FILE_ID)
+            url = injectQuery(url, `type=${workerType}`)
           }
           s.overwrite(urlIndex, urlIndex + exp.length, JSON.stringify(url))
         }
diff --git a/packages/vite/src/node/utils.ts b/packages/vite/src/node/utils.ts
index fb479a511abbd4..f830944da879cc 100644
--- a/packages/vite/src/node/utils.ts
+++ b/packages/vite/src/node/utils.ts
@@ -696,3 +696,5 @@ export function parseRequest(id: string): Record<string, string> | null {
   }
   return Object.fromEntries(new URLSearchParams(search.slice(1)))
 }
+
+export const blankReplacer = (match: string) => ' '.repeat(match.length)