From 416f190aa92f06a06f3ded403fb1e4cb8729256d Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Thu, 11 Feb 2021 17:00:51 -0500
Subject: [PATCH] feat(plugin-legacy): inject import.meta.env.LEGACY

---
 .../legacy/__tests__/legacy.spec.ts           |  6 ++
 packages/playground/legacy/index.html         |  1 +
 packages/playground/legacy/main.js            |  4 +
 packages/playground/legacy/package.json       |  3 +-
 packages/plugin-legacy/README.md              |  2 +
 packages/plugin-legacy/index.js               | 75 ++++++++++++++++++-
 packages/plugin-legacy/package.json           |  1 +
 7 files changed, 88 insertions(+), 4 deletions(-)

diff --git a/packages/playground/legacy/__tests__/legacy.spec.ts b/packages/playground/legacy/__tests__/legacy.spec.ts
index 5ac67eb368105e..aff02741711964 100644
--- a/packages/playground/legacy/__tests__/legacy.spec.ts
+++ b/packages/playground/legacy/__tests__/legacy.spec.ts
@@ -1,3 +1,9 @@
+import { isBuild } from '../../testUtils'
+
 test('should work', async () => {
   expect(await page.textContent('#app')).toMatch('Hello')
 })
+
+test('import.meta.env.LEGACY', async () => {
+  expect(await page.textContent('#env')).toMatch(isBuild ? 'true' : 'false')
+})
diff --git a/packages/playground/legacy/index.html b/packages/playground/legacy/index.html
index 26cb9f4685e81d..aef4bccfc61f74 100644
--- a/packages/playground/legacy/index.html
+++ b/packages/playground/legacy/index.html
@@ -1,2 +1,3 @@
 <h1 id="app"></h1>
+<div id="env"></div>
 <script type="module" src="./main.js"></script>
diff --git a/packages/playground/legacy/main.js b/packages/playground/legacy/main.js
index 973cd65966ad09..84409b57793631 100644
--- a/packages/playground/legacy/main.js
+++ b/packages/playground/legacy/main.js
@@ -4,3 +4,7 @@ async function run() {
 }
 
 run()
+
+document.getElementById('env').textContent = `is legacy: ${
+  import.meta.env.LEGACY
+}`
diff --git a/packages/playground/legacy/package.json b/packages/playground/legacy/package.json
index 9c7c676b3d186e..89a5130f8e9eb9 100644
--- a/packages/playground/legacy/package.json
+++ b/packages/playground/legacy/package.json
@@ -10,6 +10,5 @@
   },
   "devDependencies": {
     "@vitejs/plugin-legacy": "^1.0.0"
-  },
-  "dependencies": {}
+  }
 }
diff --git a/packages/plugin-legacy/README.md b/packages/plugin-legacy/README.md
index 259550b48fb404..cc38130d184691 100644
--- a/packages/plugin-legacy/README.md
+++ b/packages/plugin-legacy/README.md
@@ -12,6 +12,8 @@ By default, this plugin will:
 
 - Inject `<script nomodule>` tags into generated HTML to conditionally load the polyfills and legacy bundle only in browsers without native ESM support.
 
+- Inject the `import.meta.env.LEGACY` env variable, which will only be `true` in the legacy production build, and `false` in all other cases. (requires `vite@^2.0.0-beta.69`)
+
 ## Usage
 
 ```js
diff --git a/packages/plugin-legacy/index.js b/packages/plugin-legacy/index.js
index 2707f64546301d..32f0a48f58ddcf 100644
--- a/packages/plugin-legacy/index.js
+++ b/packages/plugin-legacy/index.js
@@ -2,6 +2,7 @@
 const path = require('path')
 const { createHash } = require('crypto')
 const { build } = require('vite')
+const MagicString = require('magic-string').default
 
 // lazy load babel since it's not used during dev
 let babel
@@ -17,6 +18,8 @@ const safari10NoModuleFix = `!function(){var e=document,t=e.createElement("scrip
 const legacyEntryId = 'vite-legacy-entry'
 const systemJSInlineCode = `System.import(document.getElementById('${legacyEntryId}').getAttribute('data-src'))`
 
+const legacyEnvVarMarker = `__VITE_IS_LEGACY__`
+
 /**
  * @param {import('.').Options} options
  * @returns {import('vite').Plugin[]}
@@ -184,6 +187,28 @@ function viteLegacyPlugin(options = {}) {
           // analyze and record modern polyfills
           detectPolyfills(raw, { esmodules: true }, modernPolyfills)
         }
+
+        if (raw.includes(legacyEnvVarMarker)) {
+          const re = new RegExp(`"${legacyEnvVarMarker}"`, 'g')
+          if (config.build.sourcemap) {
+            const s = new MagicString(raw)
+            let match
+            while ((match = re.exec(raw))) {
+              s.overwrite(
+                match.index,
+                match.index + legacyEnvVarMarker.length + 2,
+                `false`
+              )
+            }
+            return {
+              code: s.toString(),
+              map: s.generateMap({ hires: true })
+            }
+          } else {
+            return raw.replace(re, `false`)
+          }
+        }
+
         return null
       }
 
@@ -210,7 +235,10 @@ function viteLegacyPlugin(options = {}) {
           // preset so we can catch the injected import statements...
           [
             () => ({
-              plugins: [recordAndRemovePolyfillBabelPlugin(legacyPolyfills)]
+              plugins: [
+                recordAndRemovePolyfillBabelPlugin(legacyPolyfills),
+                replaceLegacyEnvBabelPlugin()
+              ]
             })
           ],
           [
@@ -339,7 +367,37 @@ function viteLegacyPlugin(options = {}) {
     }
   }
 
-  return [legacyGenerateBundlePlugin, legacyPostPlugin]
+  let envInjectionFaled = false
+  /**
+   * @type {import('vite').Plugin}
+   */
+  const legacyEnvPlugin = {
+    name: 'legacy-env',
+
+    config(_, env) {
+      if (env) {
+        return {
+          define: {
+            'import.meta.env.LEGACY':
+              env.command === 'serve' ? false : legacyEnvVarMarker
+          }
+        }
+      } else {
+        envInjectionFaled = true
+      }
+    },
+
+    configResolved(config) {
+      if (envInjectionFaled) {
+        config.logger.warn(
+          `[@vitejs/plugin-legacy] import.meta.env.LEGACY was not injected due ` +
+            `to incompatible vite version (requires vite@^2.0.0-beta.69).`
+        )
+      }
+    }
+  }
+
+  return [legacyGenerateBundlePlugin, legacyPostPlugin, legacyEnvPlugin]
 }
 
 /**
@@ -488,6 +546,19 @@ function recordAndRemovePolyfillBabelPlugin(polyfills) {
   })
 }
 
+function replaceLegacyEnvBabelPlugin() {
+  return ({ types: t }) => ({
+    name: 'vite-replace-env-legacy',
+    visitor: {
+      StringLiteral(path) {
+        if (path.node.value === legacyEnvVarMarker) {
+          path.replaceWith(t.booleanLiteral(true))
+        }
+      }
+    }
+  })
+}
+
 module.exports = viteLegacyPlugin
 
 viteLegacyPlugin.default = viteLegacyPlugin
diff --git a/packages/plugin-legacy/package.json b/packages/plugin-legacy/package.json
index 56b502ca5487a3..9f35fba355e735 100644
--- a/packages/plugin-legacy/package.json
+++ b/packages/plugin-legacy/package.json
@@ -27,6 +27,7 @@
   "dependencies": {
     "@babel/standalone": "^7.12.12",
     "core-js": "^3.8.2",
+    "magic-string": "^0.25.7",
     "regenerator-runtime": "^0.13.7",
     "systemjs": "^6.8.3"
   },