feat: add wasm support for test and SSR

packages/vite/src/node/plugins/wasm.ts

@@ -1,11 +1,20 @@
+import { fileURLToPath, pathToFileURL } from 'node:url'
+import { readFile } from 'node:fs/promises'
+import MagicString from 'magic-string'
 import { exactRegex } from '@rolldown/pluginutils'
+import { createToImportMetaURLBasedRelativeRuntime } from '../build'
 import type { Plugin } from '../plugin'
-import { fileToUrl } from './asset'
+import { fsPathFromId } from '../utils'
+import { FS_PREFIX } from '../constants'
+import { cleanUrl } from '../../shared/utils'
+import { assetUrlRE, fileToUrl } from './asset'
 
 const wasmHelperId = '\0vite/wasm-helper.js'
 
 const wasmInitRE = /(?<![?#].*)\.wasm\?init/
 
+const wasmInitUrlRE: RegExp = /__VITE_WASM_INIT__([\w$]+)__/g
+
 const wasmHelper = async (opts = {}, url: string) => {
   let result
   if (url.startsWith('data:')) {
@@ -26,28 +35,47 @@ const wasmHelper = async (opts = {}, url: string) => {
     }
     result = await WebAssembly.instantiate(bytes, opts)
   } else {
-    // https://github.com/mdn/webassembly-examples/issues/5
-    // WebAssembly.instantiateStreaming requires the server to provide the
-    // correct MIME type for .wasm files, which unfortunately doesn't work for
-    // a lot of static file servers, so we just work around it by getting the
-    // raw buffer.
-    const response = await fetch(url)
-    const contentType = response.headers.get('Content-Type') || ''
-    if (
-      'instantiateStreaming' in WebAssembly &&
-      contentType.startsWith('application/wasm')
-    ) {
-      result = await WebAssembly.instantiateStreaming(response, opts)
-    } else {
-      const buffer = await response.arrayBuffer()
-      result = await WebAssembly.instantiate(buffer, opts)
-    }
+    result = await instantiateFromUrl(url, opts)
   }
   return result.instance
 }
 
 const wasmHelperCode = wasmHelper.toString()
 
+const instantiateFromUrl = async (url: string, opts?: WebAssembly.Imports) => {
+  // https://github.com/mdn/webassembly-examples/issues/5
+  // WebAssembly.instantiateStreaming requires the server to provide the
+  // correct MIME type for .wasm files, which unfortunately doesn't work for
+  // a lot of static file servers, so we just work around it by getting the
+  // raw buffer.
+  const response = await fetch(url)
+  const contentType = response.headers.get('Content-Type') || ''
+  if (
+    'instantiateStreaming' in WebAssembly &&
+    contentType.startsWith('application/wasm')
+  ) {
+    return WebAssembly.instantiateStreaming(response, opts)
+  } else {
+    const buffer = await response.arrayBuffer()
+    return WebAssembly.instantiate(buffer, opts)
+  }
+}
+
+const instantiateFromUrlCode = instantiateFromUrl.toString()
+
+const instantiateFromFile = async (url: string, opts?: WebAssembly.Imports) => {
+  let fsPath = url
+  if (url.startsWith('file:')) {
+    fsPath = fileURLToPath(url)
+  } else if (url.startsWith('/')) {
+    fsPath = url.slice(1)
+  }
+  const buffer = await readFile(fsPath)
+  return WebAssembly.instantiate(buffer, opts)
+}
+
+const instantiateFromFileCode = instantiateFromFile.toString()
+
 export const wasmHelperPlugin = (): Plugin => {
   return {
     name: 'vite:wasm-helper',
@@ -62,18 +90,74 @@ export const wasmHelperPlugin = (): Plugin => {
     load: {
       filter: { id: [exactRegex(wasmHelperId), wasmInitRE] },
       async handler(id) {
+        const isServer = this.environment.config.consumer === 'server'
+
         if (id === wasmHelperId) {
-          return `export default ${wasmHelperCode}`
+          if (isServer) {
+            return `
+import { readFile } from 'node:fs/promises'
+import { fileURLToPath } from 'node:url'
+const instantiateFromUrl = ${instantiateFromFileCode}
+export default ${wasmHelperCode}
+`
+          } else {
+            return `
+const instantiateFromUrl = ${instantiateFromUrlCode}
+export default ${wasmHelperCode}
+`
+          }
         }
 
-        const url = await fileToUrl(this, id)
-
+        id = id.split('?')[0]
+        let url = await fileToUrl(this, id)
+        if (isServer) {
+          if (url.startsWith(FS_PREFIX)) {
+            url = pathToFileURL(fsPathFromId(id)).href
+          } else if (assetUrlRE.test(url)) {
+            url = url.replace('__VITE_ASSET__', '__VITE_WASM_INIT__')
+          }
+        }
         return `
   import initWasm from "${wasmHelperId}"
   export default opts => initWasm(opts, ${JSON.stringify(url)})
   `
       },
     },
+
+    renderChunk(code, chunk, opts) {
+      if (this.environment.config.consumer !== 'server') {
+        return null
+      }
+
+      const toRelativeRuntime = createToImportMetaURLBasedRelativeRuntime(
+        opts.format,
+        this.environment.config.isWorker,
+      )
+
+      let match: RegExpExecArray | null
+      let s: MagicString | undefined
+
+      wasmInitUrlRE.lastIndex = 0
+      while ((match = wasmInitUrlRE.exec(code))) {
+        const [full, referenceId] = match
+        const file = this.getFileName(referenceId)
+        chunk.viteMetadata!.importedAssets.add(cleanUrl(file))
+        const { runtime } = toRelativeRuntime(file, chunk.fileName)
+        s ||= new MagicString(code)
+        s.update(match.index, match.index + full.length, `"+${runtime}+"`)
+      }
+
+      if (s) {
+        return {
+          code: s.toString(),
+          map: this.environment.config.build.sourcemap
+            ? s.generateMap({ hires: 'boundary' })
+            : null,
+        }
+      } else {
+        return null
+      }
+    },
   }
 }
 

playground/ssr-wasm/__tests__/serve.ts

@@ -0,0 +1,42 @@
+// this is automatically detected by playground/vitestSetup.ts and will replace
+// the default e2e test serve behavior
+
+import path from 'node:path'
+import kill from 'kill-port'
+import { build } from 'vite'
+import { hmrPorts, isBuild, ports, rootDir } from '~utils'
+
+export const port = ports['ssr-wasm']
+
+export async function preServe() {
+  if (isBuild) {
+    await build({ root: rootDir })
+  }
+}
+
+export async function serve(): Promise<{ close(): Promise<void> }> {
+  await kill(port)
+
+  const { createServer } = await import(path.resolve(rootDir, 'server.js'))
+  const { app, vite } = await createServer(rootDir, hmrPorts['ssr-wasm'])
+
+  return new Promise((resolve, reject) => {
+    try {
+      const server = app.listen(port, () => {
+        resolve({
+          // for test teardown
+          async close() {
+            await new Promise((resolve) => {
+              server.close(resolve)
+            })
+            if (vite) {
+              await vite.close()
+            }
+          },
+        })
+      })
+    } catch (e) {
+      reject(e)
+    }
+  })
+}

playground/ssr-wasm/__tests__/ssr-wasm.spec.ts

@@ -0,0 +1,37 @@
+import { expect, test } from 'vitest'
+import { port } from './serve'
+import { findAssetFile, isBuild, listAssets, page } from '~utils'
+
+const url = `http://localhost:${port}`
+
+test('should work when inlined', async () => {
+  await page.goto(`${url}/static-light`)
+  expect(await page.textContent('.static-light')).toMatch('42')
+})
+
+test('should work when output', async () => {
+  await page.goto(`${url}/static-heavy`)
+  expect(await page.textContent('.static-heavy')).toMatch('24')
+})
+
+test.runIf(isBuild)('should not contain wasm file when inlined', async () => {
+  const assets = await listAssets()
+  const lightWasm = assets.find((f) => /light-.+\.wasm$/.test(f))
+  expect(lightWasm).toBeUndefined()
+
+  const staticLight = await findAssetFile(/^static-light-.+\.js$/)
+  expect(staticLight).toContain('data:application/wasm;base64,')
+})
+
+test.runIf(isBuild)(
+  'should contain and reference wasm file when output',
+  async () => {
+    const assets = await listAssets()
+    const heavyWasm = assets.find((f) => /heavy-.+\.wasm$/.test(f))
+    expect(heavyWasm).toBeDefined()
+
+    const staticHeavy = await findAssetFile(/^static-heavy-.+\.js$/)
+    expect(staticHeavy).toContain(heavyWasm)
+    expect(staticHeavy).not.toContain('data:application/wasm;base64,')
+  },
+)

playground/ssr-wasm/package.json

@@ -0,0 +1,15 @@
+{
+  "name": "@vitejs/test-ssr-wasm",
+  "private": true,
+  "version": "0.0.0",
+  "type": "module",
+  "scripts": {
+    "dev": "node server",
+    "build": "vite build",
+    "preview": "NODE_ENV=production node server"
+  },
+  "dependencies": {},
+  "devDependencies": {
+    "express": "^5.1.0"
+  }
+}

playground/ssr-wasm/server.js

@@ -0,0 +1,60 @@
+import express from 'express'
+
+const isTest = process.env.VITEST
+const isProduction = process.env.NODE_ENV === 'production'
+
+export async function createServer(root = process.cwd(), hmrPort) {
+  const app = express()
+
+  /** @type {import('vite').ViteDevServer} */
+  let vite
+  if (!isProduction) {
+    vite = await (
+      await import('vite')
+    ).createServer({
+      root,
+      logLevel: isTest ? 'error' : 'info',
+      server: {
+        middlewareMode: true,
+        watch: {
+          // During tests we edit the files too fast and sometimes chokidar
+          // misses change events, so enforce polling for consistency
+          usePolling: true,
+          interval: 100,
+        },
+        hmr: {
+          port: hmrPort,
+        },
+      },
+      appType: 'custom',
+    })
+    // use vite's connect instance as middleware
+    app.use(vite.middlewares)
+  }
+
+  app.use('*all', async (req, res, next) => {
+    try {
+      const url = req.originalUrl
+      const render = isProduction
+        ? (await import('./dist/app.js')).render
+        : (await vite.ssrLoadModule('/src/app.js')).render
+      const html = await render(url)
+      res.status(200).set({ 'Content-Type': 'text/html' }).end(html)
+    } catch (e) {
+      vite?.ssrFixStacktrace(e)
+      if (isTest) throw e
+      console.log(e.stack)
+      res.status(500).end(e.stack)
+    }
+  })
+
+  return { app, vite }
+}
+
+if (!isTest) {
+  createServer().then(({ app }) =>
+    app.listen(5173, () => {
+      console.log('http://localhost:5173')
+    }),
+  )
+}

playground/ssr-wasm/src/app.js

@@ -0,0 +1,8 @@
+export async function render(url) {
+  switch (url) {
+    case '/static-light':
+      return (await import('./static-light')).render()
+    case '/static-heavy':
+      return (await import('./static-heavy')).render()
+  }
+}

playground/ssr-wasm/src/heavy.wasm

@@ -0,0 +1 @@
+../../wasm/heavy.wasm

playground/ssr-wasm/src/light.wasm

@@ -0,0 +1 @@
+../../wasm/light.wasm

playground/ssr-wasm/src/static-heavy.js

@@ -0,0 +1,12 @@
+import heavy from './heavy.wasm?init'
+
+export async function render() {
+  let result
+  const { exported_func } = await heavy({
+    imports: {
+      imported_func: (res) => (result = res),
+    },
+  }).then((i) => i.exports)
+  exported_func()
+  return `<div class="static-heavy">${result}</div>`
+}

playground/ssr-wasm/src/static-light.js

@@ -0,0 +1,12 @@
+import light from './light.wasm?init'
+
+export async function render() {
+  let result
+  const { exported_func } = await light({
+    imports: {
+      imported_func: (res) => (result = res),
+    },
+  }).then((i) => i.exports)
+  exported_func()
+  return `<div class="static-light">${result}</div>`
+}

playground/ssr-wasm/vite.config.ts

@@ -0,0 +1,9 @@
+import { defineConfig } from 'vite'
+export default defineConfig({
+  build: {
+    // make cannot emit light.wasm
+    assetsInlineLimit: 80,
+    ssr: './src/app.js',
+    ssrEmitAssets: true,
+  },
+})

playground/test-utils.ts

@@ -34,6 +34,7 @@ export const ports = {
   'ssr-html': 9602,
   'ssr-noexternal': 9603,
   'ssr-pug': 9604,
+  'ssr-wasm': 9608,
   'ssr-webworker': 9605,
   'proxy-bypass': 9606, // not imported but used in `proxy-hmr/vite.config.js`
   'proxy-bypass/non-existent-app': 9607, // not imported but used in `proxy-hmr/other-app/vite.config.js`
@@ -57,6 +58,7 @@ export const hmrPorts = {
   'ssr-html': 24683,
   'ssr-noexternal': 24684,
   'ssr-pug': 24685,
+  'ssr-wasm': 24691,
   'css/lightningcss-proxy': 24686,
   json: 24687,
   'ssr-conditions': 24688,