From 7a4b2edbea41a8f2d6d805d3ac5f9cf8820f14d9 Mon Sep 17 00:00:00 2001
From: Ayaka Rizumu <rizumu@oqo.moe>
Date: Thu, 12 Jan 2023 20:53:18 +0800
Subject: [PATCH] feat(vite): introduce `hmrTopLevelAwait` option (#2066)

Co-authored-by: Frozen FIsh <76603360+sudongyuer@users.noreply.github.com>
Co-authored-by: Chris <1633711653@qq.com>
Co-authored-by: Anthony Fu <anthonyfu117@hotmail.com>
---
 packages/vite/src/modes/global/dev.ts | 29 ++++++++++++++++++---------
 packages/vite/src/types.ts            | 16 ++++++++++++++-
 playground/src/auto-imports.d.ts      |  1 +
 playground/vite.config.ts             |  6 ++++--
 4 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/packages/vite/src/modes/global/dev.ts b/packages/vite/src/modes/global/dev.ts
index b11845bc7b..cf6e29180e 100644
--- a/packages/vite/src/modes/global/dev.ts
+++ b/packages/vite/src/modes/global/dev.ts
@@ -1,13 +1,14 @@
 import type { Plugin, Update, ViteDevServer, ResolvedConfig as ViteResolvedConfig } from 'vite'
 import type { GenerateResult, UnocssPluginContext } from '@unocss/core'
 import { notNull } from '@unocss/core'
+import type { VitePluginConfig } from 'unocss/vite'
 import { LAYER_MARK_ALL, getHash, getPath, resolveId, resolveLayer } from '../../integration'
 
 const WARN_TIMEOUT = 20000
 const WS_EVENT_PREFIX = 'unocss:hmr'
 const HASH_LENGTH = 6
 
-export function GlobalModeDevPlugin({ uno, tokens, tasks, flushTasks, affectedModules, onInvalidate, extract, filter }: UnocssPluginContext): Plugin[] {
+export function GlobalModeDevPlugin({ uno, tokens, tasks, flushTasks, affectedModules, onInvalidate, extract, filter, getConfig }: UnocssPluginContext): Plugin[] {
   const servers: ViteDevServer[] = []
   let base = ''
   const entries = new Set<string>()
@@ -162,18 +163,28 @@ export function GlobalModeDevPlugin({ uno, tokens, tasks, flushTasks, affectedMo
         return env.command === 'serve' && !config.build?.ssr
       },
       enforce: 'post',
-      transform(code, id) {
+      async transform(code, id) {
         const layer = resolveLayer(getPath(id))
 
         // inject css modules to send callback on css load
         if (layer && code.includes('import.meta.hot')) {
-          return `${code}
-if (import.meta.hot) {
-  try { await import.meta.hot.send('${WS_EVENT_PREFIX}', ['${layer}', __vite__css.slice(2,${2 + HASH_LENGTH})]); }
-  catch (e) { console.warn('[unocss-hmr]', e) }
-  if (!import.meta.url.includes('?'))
-    await new Promise(resolve => setTimeout(resolve, 100))
-}`
+          let hmr = `
+try {
+  await import.meta.hot.send('${WS_EVENT_PREFIX}', ['${layer}', __vite__css.slice(2,${2 + HASH_LENGTH})]);
+} catch (e) {
+  console.warn('[unocss-hmr]', e)
+}
+if (!import.meta.url.includes('?'))
+  await new Promise(resolve => setTimeout(resolve, 100))`
+
+          const config = await getConfig() as VitePluginConfig
+
+          if (config.hmrTopLevelAwait === false)
+            hmr = `;(async function() {${hmr}\n})()`
+
+          hmr = `\nif (import.meta.hot) {${hmr}}`
+
+          return code + hmr
         }
       },
     },
diff --git a/packages/vite/src/types.ts b/packages/vite/src/types.ts
index ec2ca46b41..becf5bcf6d 100644
--- a/packages/vite/src/types.ts
+++ b/packages/vite/src/types.ts
@@ -21,6 +21,7 @@ export interface VitePluginConfig<Theme extends {} = {}> extends UserConfig<Them
    * @default 'global'
    */
   mode?: 'global' | 'per-module' | 'vue-scoped' | 'svelte-scoped' | 'dist-chunk' | 'shadow-dom'
+
   /**
    * Transform CSS for `@apply` directive
    *
@@ -28,10 +29,23 @@ export interface VitePluginConfig<Theme extends {} = {}> extends UserConfig<Them
    * @default false
    */
   transformCSS?: boolean | 'pre' | 'post'
+
   /**
+   * Make the generated css processed by postcss (https://vitejs.dev/guide/features.html#postcss)
    *
-   * make the generated css processed by postcss (https://vitejs.dev/guide/features.html#postcss)
    * @default true
    */
   postcss?: boolean
+
+  /**
+   * Use top level await in HMR code to avoid FOUC on dev time.
+   *
+   * You usually don't need to disable this, unless you are developing on
+   * a browser that does not support top level await.
+   *
+   * This will only affect on dev time.
+   *
+   * @default true
+   */
+  hmrTopLevelAwait?: boolean
 }
diff --git a/playground/src/auto-imports.d.ts b/playground/src/auto-imports.d.ts
index 6894241fec..010a81256c 100644
--- a/playground/src/auto-imports.d.ts
+++ b/playground/src/auto-imports.d.ts
@@ -111,6 +111,7 @@ declare global {
   const useArrayMap: typeof import('@vueuse/core')['useArrayMap']
   const useArrayReduce: typeof import('@vueuse/core')['useArrayReduce']
   const useArraySome: typeof import('@vueuse/core')['useArraySome']
+  const useArrayUnique: typeof import('@vueuse/core')['useArrayUnique']
   const useAsyncQueue: typeof import('@vueuse/core')['useAsyncQueue']
   const useAsyncState: typeof import('@vueuse/core')['useAsyncState']
   const useAttrs: typeof import('vue')['useAttrs']
diff --git a/playground/vite.config.ts b/playground/vite.config.ts
index 626376592f..d9e69303c3 100644
--- a/playground/vite.config.ts
+++ b/playground/vite.config.ts
@@ -3,7 +3,7 @@ import Vue from '@vitejs/plugin-vue'
 import Inspect from 'vite-plugin-inspect'
 import Components from 'unplugin-vue-components/vite'
 import AutoImport from 'unplugin-auto-import/vite'
-import Unocss from '@unocss/vite'
+import UnoCSS from '@unocss/vite'
 import { alias } from '../alias'
 
 export default defineConfig({
@@ -13,7 +13,9 @@ export default defineConfig({
   },
   plugins: [
     Vue(),
-    Unocss(),
+    UnoCSS({
+      // hmrTopLevelAwait: false,
+    }),
     Inspect(),
     Components({
       dirs: [