fix: hmr server

Makefile

@@ -0,0 +1,20 @@
+install:
+	@echo "Setup pnpm package manager..."
+	@corepack enable
+	pnpm install
+
+build:
+	@echo "Building..."
+	@pnpm exec rollup --config rollup.config.ts --configPlugin swc3
+
+dev:
+	@echo "Starting development server..."
+	@pnpm exec rollup --config rollup.config.ts --configPlugin swc3 --watch
+
+test:
+	@echo "Running tests..."
+	@pnpm exec vitest --dir __tests__
+
+end-to-end-test:
+	@echo "Running end-to-end tests..."
+	@pnpm exec vitest --dir e2e

src/context.ts

@@ -107,8 +107,11 @@ export class PluginContext {
     let pass = false
     for (const stmt of this.#stmts) {
       const { n } = stmt
-      if (n && this.importSources.some(i => !path.isAbsolute(n) && n.includes(typeof i === 'string' ? i : i.from))) {
-        pass = true
+      if (n) {
+        if (n.endsWith('.css')) continue
+        if (this.importSources.some(i => !path.isAbsolute(n) && n.includes(typeof i === 'string' ? i : i.from))) {
+          pass = true
+        }
       }
     }
     return pass

src/plugins/index.ts

@@ -1,19 +1,21 @@
-import type { Plugin } from 'vite'
+import type { HookHandler, Plugin } from 'vite'
 import { PluginContext } from '../context'
-import { error, searchForPackageRoot, unique } from '../shared'
+import { error, hijackHook, searchForPackageRoot, unique } from '../shared'
 import type { AdapterContext } from '../interface'
 import { stylexBuild } from './build'
 import { CONSTANTS, stylexServer } from './server'
+import type { CssHooks } from './server'
 
 export function createForViteServer(ctx: PluginContext, extend: (c: PluginContext) => Plugin) {
   const cssPlugins: Plugin[] = []
+  const cssHooks: CssHooks = new Map()
 
   return (plugin: Plugin) => {
     plugin.configResolved = function configResolved(conf) {
       const adapterContext: AdapterContext = {
         // eslint-disable-next-line prefer-spread
         produceCSS: (...rest: any) => ctx.produceCSS.apply(ctx, rest),
-        transform: plugin.transform as any,
+        transform: plugin.transform as HookHandler<Plugin['transform']>,
         vite: { cssPlugins, config: conf },
         env: ctx.env,
         rules: ctx.styleRules,
@@ -41,14 +43,19 @@ export function createForViteServer(ctx: PluginContext, extend: (c: PluginContex
           ? [...conf.ssr.noExternal, ...optimizedDeps]
           : conf.ssr.noExternal
       }
-      cssPlugins.push(...conf.plugins.filter(p => CONSTANTS.CSS_PLUGINS.includes(p.name)))
-      cssPlugins.sort((a, b) => a.name.length < b.name.length ? -1 : 1)
+
+      conf.plugins.forEach(p => {
+        if (CONSTANTS.CSS_PLUGINS.includes(p.name)) {
+          cssHooks.set(p.name, hijackHook(p, 'transform', (fn, c, args) => fn.apply(c, args), true))
+        }
+      })
+
       const pos = conf.plugins.findIndex(p => p.name === 'stylex')
       if (Object.keys(ctx.stylexExtendOptions).length) {
         // @ts-expect-error
         conf.plugins.splice(pos, 0, extend(ctx))
       }
-      ctx.env === 'build' ? stylexBuild(plugin, ctx, cssPlugins) : stylexServer(plugin, ctx, cssPlugins, conf)
+      ctx.env === 'build' ? stylexBuild(plugin, ctx, cssPlugins) : stylexServer(plugin, { ctx, cssHooks, config: conf })
       if (typeof ctx.stylexOptions.adapter === 'function') {
         const adapter = ctx.stylexOptions.adapter()
         if (!adapter.name) {

src/plugins/server.ts

@@ -3,16 +3,27 @@ import type { HookHandler, Plugin, ResolvedConfig, Update, ViteDevServer } from
 import { PluginContext, parseRequest } from '../context'
 import { hash, hijackHook } from '../shared'
 
+export type CssHooks = Map<'vite:css' | 'vite:css-post' | string & ({}), HookHandler<Plugin['transform']>>
+
+export interface InternalConfig {
+  config: ResolvedConfig
+  cssHooks: CssHooks
+  ctx: PluginContext
+}
+
 export const CONSTANTS = {
   WS_EVENT: 'stylex:hmr',
   VIRTUAL_STYLEX_ID: 'virtual:stylex.css',
   RESOLVED_ID_WITH_QUERY_REG: /[/\\]__stylex(_.*?)?\.css(\?.*)?$/,
   RESOLVED_ID_REG: /[/\\]__stylex(?:_(.*?))?\.css$/,
   VIRTUAL_ENTRY_ALIAS: [/^(?:virtual:)?stylex(?::(.+))?\.css(\?.*)?$/],
+  STYLEX_CSS: '/__stylex.css',
   STYLEX_BUNDLE_MARK: '@stylex__bundle__marker;',
   HASH_LENGTH: 6,
   CSS_PLUGINS: ['vite:css', 'vite:css-post'],
-  WELL_KNOW_LIBRARIES: ['@stylexjs/open-props']
+  WELL_KNOW_LIBRARIES: ['@stylexjs/open-props'],
+  STYLEX_START_COMMENT: '#--stylex-dev-start--#',
+  STYLEX_END_COMMENT: '#--stylex-dev-end--#'
 }
 
 export function resolveId(id: string) {
@@ -40,128 +51,136 @@ try {
 } catch (e) {
   console.warn('[vite-plugin0-stylex]', e)
 }
-if (!import.meta.url.include('?')) {
+if(!import.meta.url.includes('?')) {
    await new Promise(r => setTimeout(r, 100))
 }
 `
 
 hmr = `\nif (import.meta.hot) { ${hmr} }`
 
-export function stylexServer(plugin: Plugin, ctx: PluginContext, cssPlugins: Plugin[], conf: ResolvedConfig) {
-  const cssHooks = new Map<'vite:css' | 'vite:css-post' | string & ({}), HookHandler<Plugin['transform']>>()
-  let viteDevServer: ViteDevServer | null = null
-  let lastServerTime = Date.now()
-  const modules = new Set<string>()
-
-  const generateCSS = () => {
+export function createHasteCSS(ctx: PluginContext, effects: Set<string>) {
+  return (handler?: (css: string) => void) => {
     let css = ''
-    const expect = modules.size
     for (;;) {
       css = ctx.produceCSS()
-      if (expect === ctx.styleRules.size) {
+      if (effects.size === ctx.styleRules.size) {
         break
       }
     }
-    lastServerTime = Date.now()
-    return { css, hash: hash(css) }
+
+    return css
   }
+}
+
+export function stylexServer(self: Plugin, options: InternalConfig) {
+  // let viteDevServer: ViteDevServer | null = null
+  // let lastServerTime = Date.now()
+  // const modules = new Set<string>()
+
+  const { ctx, config, cssHooks } = options
+  let viteServer: ViteDevServer | null = null
+  const { isManuallyControlCSS, controlCSSByManually = { id, symbol } } = ctx
+  const effects = new Set<string>()
+  const entries = new Set<string>()
+  // const diffs = new Set<string>()
+  let lastHMRTime = Date.now()
+  let invalidateTimer: NodeJS.Timeout | null
+  const generateCSS = createHasteCSS(ctx, effects)
 
   const update = (ids: Set<string>) => {
-    if (!viteDevServer) return
-    viteDevServer.ws.send({
+    viteServer?.ws.send({
       type: 'update',
       updates: Array.from(ids).map(id => {
-        const mod = viteDevServer?.moduleGraph.getModuleById(id)
+        const mod = viteServer?.moduleGraph.getModuleById(id)
         if (!mod) return null
         return {
           acceptedPath: id,
           path: mod.url,
-          timestamp: lastServerTime,
+          timestamp: lastHMRTime,
           type: 'js-update'
-        } as Update
+        } satisfies Update
       }).filter((s) => s !== null) as Update[]
     })
   }
-
-  let invalidateTimer: any
-
-  const invalidate = (ids: Set<string>) => {
+  const onInvalidate = (ids: Set<string>) => {
     for (const id of ids) {
-      const mod = viteDevServer?.moduleGraph.getModuleById(id)
+      const mod = viteServer?.moduleGraph.getModuleById(id)
       if (!mod) continue
-      viteDevServer?.moduleGraph.invalidateModule(mod)
+      viteServer?.moduleGraph.invalidateModule(mod)
     }
-    clearTimeout(invalidateTimer)
+    invalidateTimer && clearTimeout(invalidateTimer)
     invalidateTimer = setTimeout(() => {
       update(ids)
-    }, 10)
+    }, 20)
   }
 
-  const entries = new Set<string>()
-
-  const schedule = <Partial<Plugin>> {
+  const scan = {
+    name: 'stylex:server-scan',
     enforce: 'pre',
-    name: 'stylex:server',
     apply: 'serve',
-    resolveId(id: string) {
-      const entry = resolveId(id)
-      if (entry) {
-        entries.add(entry)
-        return entry
+    configureServer(server) {
+      viteServer = server
+      server.ws.on(CONSTANTS.WS_EVENT, () => {
+        lastHMRTime = Date.now()
+        update(entries)
+      })
+    },
+    resolveId(id) {
+      if (isManuallyControlCSS && id === controlCSSByManually.id) {
+        entries.add(CONSTANTS.STYLEX_CSS)
+        return CONSTANTS.STYLEX_CSS
+      }
+      if (resolveId(id)) {
+        entries.add(CONSTANTS.STYLEX_CSS)
+        return CONSTANTS.STYLEX_CSS
       }
     },
-    load(id: string) {
+    load(id) {
       const { original } = parseRequest(id)
-      const matched = original.match(CONSTANTS.RESOLVED_ID_REG)
-      if (matched) {
-        const { hash, css } = generateCSS()
+      if (original.match(CONSTANTS.RESOLVED_ID_REG) && !isManuallyControlCSS) {
+        let uuid = ''
+        const css = generateCSS((css) => {
+          uuid = hash(css)
+        })
+        lastHMRTime = Date.now()
         return {
-          code: `${css}__stylex_hash_${hash}{--:'';}`,
+          code: CONSTANTS.STYLEX_START_COMMENT + `${css}__css_hash_${uuid}{--:'';}` + CONSTANTS.STYLEX_END_COMMENT,
           map: { mappings: '' }
         }
-      } else {
-        if (ctx.isManuallyControlCSS && original === ctx.controlCSSByManually.id) {
-          entries.add(id)
-        }
       }
-    },
-    configureServer(server) {
-      viteDevServer = server
-      viteDevServer.ws.on(CONSTANTS.WS_EVENT, () => {
-        update(entries)
-      })
     }
-  }
+  } satisfies Plugin
 
-  const pos = conf.plugins.findIndex(p => p.name === 'stylex')
-  // @ts-expect-error
-  conf.plugins.splice(pos, 0, schedule)
-  //
-  hijackHook(plugin, 'transform', async (fn, c, args) => {
-    const id = args[1]
-    const { original } = parseRequest(id)
-    const result = await fn.apply(c, args)
-    if (ctx.styleRules.has(original)) {
-      // record affect module.
-      modules.add(original)
-      invalidate(new Set([...entries, ...modules]))
-    }
+  const schedule = {
+    name: 'stylex:server-schedule',
+    enforce: 'post',
 
-    if (ctx.isManuallyControlCSS && ctx.controlCSSByManually.id === original) {
-      //  FIXME
-      // Find a better way pipe to vite's internal processer
-      if (!cssHooks.size) {
-        cssPlugins.forEach((p) => {
-          cssHooks.set(p.name, hijackHook(p, 'transform', (fn, c, args) => fn.apply(c, args), true))
-        })
+    async transform(code, id) {
+      // force blocking css generation
+      if (id === CONSTANTS.STYLEX_CSS && code.includes('import.meta.hot')) {
+        return {
+          code: code + hmr,
+          map: { mappings: '' }
+        }
       }
-      const css = fs.readFileSync(ctx.controlCSSByManually.id, 'utf8').replace(ctx.controlCSSByManually.symbol!, generateCSS().css)
-      return cssHooks.get('vite:css')?.apply(c, [css, id, args[2]])
     }
+  } satisfies Plugin
 
-    if (original.match(CONSTANTS.RESOLVED_ID_REG) && args[0].includes('import.meta.hot')) {
-      const code = args[0] + hmr
-      return { code, map: { mappings: '' } }
+  const pos = config.plugins.findIndex(p => p.name === 'vite:css')
+  // @ts-expect-error
+  config.plugins.splice(pos, 0, scan)
+  const cssPostPos = config.plugins.findIndex(p => p.name === 'vite:css-post')
+  // @ts-expect-error
+  config.plugins.splice(cssPostPos + 1, 0, schedule)
+
+  hijackHook(self, 'transform', async (fn, c, args) => {
+    const result = await fn.apply(c, args)
+    const { original } = parseRequest(args[1])
+    if (result && typeof result === 'object') {
+      if (result.meta && Reflect.has(result.meta, 'stylex') && result.meta.stylex.length) {
+        effects.add(original)
+      }
+      onInvalidate(new Set([args[1], ...entries]))
     }
     return result
   })