feat(wasm): add SSR support for .wasm?init

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

@@ -1,5 +1,6 @@
 import { exactRegex } from '@rolldown/pluginutils'
 import type { Plugin } from '../plugin'
+import { fsPathFromId } from '../utils'
 import { fileToUrl } from './asset'
 
 const wasmHelperId = '\0vite/wasm-helper.js'
@@ -26,28 +27,45 @@ 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 (
+  fsPath: string,
+  opts?: WebAssembly.Imports,
+) => {
+  const fs = await import('node:fs/promises')
+  const buffer = await fs.readFile(fsPath)
+  return WebAssembly.instantiate(buffer, opts)
+}
+
+const instantiateFromFileCode = instantiateFromFile.toString()
+
 export const wasmHelperPlugin = (): Plugin => {
   return {
     name: 'vite:wasm-helper',
@@ -62,11 +80,20 @@ 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}`
+          const instantiateFromUrl = isServer
+            ? instantiateFromFileCode
+            : instantiateFromUrlCode
+          return `
+const instantiateFromUrl = ${instantiateFromUrl}
+export default ${wasmHelperCode}
+`
         }
 
-        const url = await fileToUrl(this, id)
+        id = id.split('?')[0]
+        const url = isServer ? fsPathFromId(id) : await fileToUrl(this, id)
 
         return `
   import initWasm from "${wasmHelperId}"

playground/ssr-wasm/__tests__/serve.ts

@@ -0,0 +1,35 @@
+// 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 { hmrPorts, ports, rootDir } from '~utils'
+
+export const port = ports['ssr-wasm']
+
+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,15 @@
+import { expect, test } from 'vitest'
+import { port } from './serve'
+import { 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')
+})

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",
+    "serve": "NODE_ENV=production node server",
+    "debug": "node --inspect-brk server"
+  },
+  "dependencies": {},
+  "devDependencies": {
+    "express": "^5.1.0"
+  }
+}

playground/ssr-wasm/server.js

@@ -0,0 +1,56 @@
+import express from 'express'
+
+const isTest = process.env.VITEST
+
+export async function createServer(root = process.cwd(), hmrPort) {
+  const app = express()
+
+  /**
+   * @type {import('vite').ViteDevServer}
+   */
+  const 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 = (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 && 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,7 @@
+import { defineConfig } from 'vite'
+export default defineConfig({
+  build: {
+    // make cannot emit light.wasm
+    assetsInlineLimit: 80,
+  },
+})

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,

playground/wasm/__tests__/wasm.spec.ts

@@ -1,33 +1,67 @@
-import { expect, test } from 'vitest'
+import { describe, expect, test } from 'vitest'
 import { isBuild, page } from '~utils'
 
-test('should work when inlined', async () => {
-  await page.click('.inline-wasm .run')
-  await expect
-    .poll(() => page.textContent('.inline-wasm .result'))
-    .toMatch('42')
-})
+describe('WASM in browser', () => {
+  test('should work when inlined', async () => {
+    await page.click('.inline-wasm .run')
+    await expect
+      .poll(() => page.textContent('.inline-wasm .result'))
+      .toMatch('42')
+  })
 
-test('should work when output', async () => {
-  await page.click('.output-wasm .run')
-  await expect
-    .poll(() => page.textContent('.output-wasm .result'))
-    .toMatch('24')
-})
+  test('should work when output', async () => {
+    await page.click('.output-wasm .run')
+    await expect
+      .poll(() => page.textContent('.output-wasm .result'))
+      .toMatch('24')
+  })
 
-test('init function returns WebAssembly.Instance', async () => {
-  await page.click('.init-returns-instance .run')
-  await expect
-    .poll(() => page.textContent('.init-returns-instance .result'))
-    .toMatch('true')
-})
+  test('init function returns WebAssembly.Instance', async () => {
+    await page.click('.init-returns-instance .run')
+    await expect
+      .poll(() => page.textContent('.init-returns-instance .result'))
+      .toMatch('true')
+  })
 
-test('?url', async () => {
-  expect(await page.textContent('.url')).toMatch(
-    isBuild ? 'data:application/wasm' : '/light.wasm',
-  )
+  test('?url', async () => {
+    expect(await page.textContent('.url')).toMatch(
+      isBuild ? 'data:application/wasm' : '/light.wasm',
+    )
+  })
+
+  test('should work when wasm in worker', async () => {
+    await expect
+      .poll(() => page.textContent('.worker-wasm .result'))
+      .toMatch('3')
+  })
 })
 
-test('should work when wasm in worker', async () => {
-  await expect.poll(() => page.textContent('.worker-wasm .result')).toMatch('3')
+describe('WASM in vitest', () => {
+  test('should work when inlined', async () => {
+    const { default: light } = await import('../light.wasm?init')
+    let result
+    const instance = await light({
+      imports: {
+        imported_func: (res) => (result = res),
+      },
+    })
+    expect(instance instanceof WebAssembly.Instance).toBe(true)
+    // @ts-ignore
+    instance.exports.exported_func()
+    expect(result).toBe(42)
+  })
+
+  test('should work when output', async () => {
+    const { default: heavy } = await import('../heavy.wasm?init')
+    let result
+    const instance = await heavy({
+      imports: {
+        imported_func: (res) => (result = res),
+      },
+    })
+    expect(instance instanceof WebAssembly.Instance).toBe(true)
+    // @ts-ignore
+    instance.exports.exported_func()
+    expect(result).toBe(24)
+  })
 })