feat(next/mdx): support experimental mdx-rs loader (#41919)

<!--
Thanks for opening a PR! Your contribution is much appreciated.
To make sure your PR is handled as smoothly as possible we request that
you follow the checklist sections below.
Choose the right checklist for the change that you're making:
-->

This PR implements a configuration option in `next.config.js` to support
for rust-based MDX compiler (https://github.com/wooorm/mdxjs-rs).
Currently it is marked as experimental, as mdx-rs and loader both may
need some iteration to fix regressions.

When a new experimental flag is set (`mdxRs: true`), instead of using
`@mdx/js` loader `next/mdx` loads a new webpack loader instead.
Internally, it mostly mimics mdx/js's loader behavior, however actual
compilation will be delegated into next/swc's binding to call mdx-rs's
compiler instead.

I choose to use next-swc to expose its binding function, as mdx-rs
internally uses swc already and creating new binary can double up
bytesize (with all the complex processes of publishing new package).

## Bug

- [ ] Related issues linked using `fixes #number`
- [ ] Integration tests added
- [ ] Errors have a helpful link attached, see `contributing.md`

## Feature

- [ ] Implements an existing feature request or RFC. Make sure the
feature request has been accepted for implementation before opening a
PR.
- [ ] Related issues linked using `fixes #number`
- [ ] Integration tests added
- [ ] Documentation added
- [ ] Telemetry added. In case of a feature if it's used or not.
- [ ] Errors have a helpful link attached, see `contributing.md`

## Documentation / Examples

- [ ] Make sure the linting passes by running `pnpm lint`
- [ ] The "examples guidelines" are followed from [our contributing
doc](https://github.com/vercel/next.js/blob/canary/contributing/examples/adding-examples.md)

docs/advanced-features/using-mdx.md

@@ -207,6 +207,19 @@ export default function Post(props) {
 
 If you use it across the site you may want to add the provider to `_app.js` so all MDX pages pick up the custom element config.
 
+## Using rust based MDX compiler (experimental)
+
+Next.js supports a new MDX compiler written in Rust. This compiler is still experimental and is not recommended for production use. To use the new compiler, you need to configure `next.config.js` when you pass it to `withMDX`:
+
+```js
+// next.config.js
+module.exports = withMDX({
+  experimental: {
+    mdxRs: true,
+  },
+})
+```
+
 ## Helpful Links
 
 - [MDX](https://mdxjs.com)

package.json

@@ -56,7 +56,7 @@
     "@babel/preset-react": "7.14.5",
     "@edge-runtime/jest-environment": "2.0.0",
     "@fullhuman/postcss-purgecss": "1.3.0",
-    "@mdx-js/loader": "0.18.0",
+    "@mdx-js/loader": "^1.5.1",
     "@next/bundle-analyzer": "workspace:*",
     "@next/env": "workspace:*",
     "@next/eslint-plugin-next": "workspace:*",

packages/next-mdx/index.js

@@ -3,20 +3,32 @@ module.exports =
   (nextConfig = {}) => {
     const extension = pluginOptions.extension || /\.mdx$/
 
+    const loader = nextConfig?.experimental?.mdxRs
+      ? {
+          loader: require.resolve('./mdx-rs-loader'),
+          options: {
+            providerImportSource: '@mdx-js/react',
+            ...pluginOptions.options,
+          },
+        }
+      : {
+          loader: require.resolve('@mdx-js/loader'),
+          options: {
+            providerImportSource: '@mdx-js/react',
+            ...pluginOptions.options,
+          },
+        }
+
     return Object.assign({}, nextConfig, {
       webpack(config, options) {
         config.module.rules.push({
           test: extension,
           use: [
-            options.defaultLoaders.babel,
-            {
-              loader: require.resolve('@mdx-js/loader'),
-              options: {
-                providerImportSource: '@mdx-js/react',
-                ...pluginOptions.options,
-              },
-            },
-          ],
+            nextConfig?.experimental?.mdxRs
+              ? undefined
+              : options.defaultLoaders.babel,
+            loader,
+          ].filter(Boolean),
         })
 
         if (typeof nextConfig.webpack === 'function') {

packages/next-mdx/mdx-rs-loader.js

@@ -0,0 +1,136 @@
+const { SourceMapGenerator } = require('source-map')
+const path = require('path')
+const { createHash } = require('crypto')
+
+const markdownExtensions = [
+  'md',
+  'markdown',
+  'mdown',
+  'mkdn',
+  'mkd',
+  'mdwn',
+  'mkdown',
+  'ron',
+]
+const mdx = ['.mdx']
+const md = markdownExtensions.map((/** @type {string} */ d) => '.' + d)
+
+const own = {}.hasOwnProperty
+
+const marker = {}
+const cache = new WeakMap()
+
+/**
+ * A webpack loader for mdx-rs. This is largely based on existing @mdx-js/loader,
+ * replaces internal compilation logic to use mdx-rs instead.
+ */
+function loader(value, bindings, callback) {
+  const defaults = this.sourceMap ? { SourceMapGenerator } : {}
+  const options = this.getOptions()
+  const config = { ...defaults, ...options }
+  const hash = getOptionsHash(options)
+  const compiler = this._compiler || marker
+
+  let map = cache.get(compiler)
+
+  if (!map) {
+    map = new Map()
+    cache.set(compiler, map)
+  }
+
+  let process = map.get(hash)
+
+  if (!process) {
+    process = createFormatAwareProcessors(bindings, config).compile
+    map.set(hash, process)
+  }
+
+  process({ value, path: this.resourcePath }).then(
+    (code) => {
+      // TODO: no sourcemap
+      callback(null, code, null)
+    },
+    (error) => {
+      const fpath = path.relative(this.context, this.resourcePath)
+      error.message = `${fpath}:${error.name}: ${error.message}`
+      callback(error)
+    }
+  )
+}
+
+function getOptionsHash(options) {
+  const hash = createHash('sha256')
+  let key
+
+  for (key in options) {
+    if (own.call(options, key)) {
+      const value = options[key]
+
+      if (value !== undefined) {
+        const valueString = JSON.stringify(value)
+        hash.update(key + valueString)
+      }
+    }
+  }
+
+  return hash.digest('hex').slice(0, 16)
+}
+
+function createFormatAwareProcessors(bindings, compileOptions = {}) {
+  const mdExtensions = compileOptions.mdExtensions || md
+  const mdxExtensions = compileOptions.mdxExtensions || mdx
+
+  let cachedMarkdown
+  let cachedMdx
+
+  return {
+    extnames:
+      compileOptions.format === 'md'
+        ? mdExtensions
+        : compileOptions.format === 'mdx'
+        ? mdxExtensions
+        : mdExtensions.concat(mdxExtensions),
+    compile,
+  }
+
+  function compile({ value, path: p }) {
+    const format =
+      compileOptions.format === 'md' || compileOptions.format === 'mdx'
+        ? compileOptions.format
+        : path.extname(p) &&
+          (compileOptions.mdExtensions || md).includes(path.extname(p))
+        ? 'md'
+        : 'mdx'
+
+    const options = {
+      parse: compileOptions.parse,
+      development: compileOptions.development,
+      providerImportSource: compileOptions.providerImportSource,
+      jsx: compileOptions.jsx,
+      jsxRuntime: compileOptions.jsxRuntime,
+      jsxImportSource: compileOptions.jsxImportSource,
+      pragma: compileOptions.pragma,
+      pragmaFrag: compileOptions.pragmaFrag,
+      pragmaImportSource: compileOptions.pragmaImportSource,
+      filepath: p,
+    }
+
+    const compileMdx = (input) => bindings.mdx.compile(input, options)
+
+    const processor =
+      format === 'md'
+        ? cachedMarkdown || (cachedMarkdown = compileMdx)
+        : cachedMdx || (cachedMdx = compileMdx)
+
+    return processor(value)
+  }
+}
+
+module.exports = function (code) {
+  const callback = this.async()
+  const { loadBindings } = require('next/dist/build/swc')
+
+  loadBindings().then((bindings) => {
+    return loader.call(this, code, bindings, callback)
+  })
+}

packages/next-mdx/package.json

@@ -10,5 +10,8 @@
   "peerDependencies": {
     "@mdx-js/loader": ">=0.15.0",
     "@mdx-js/react": "*"
+  },
+  "dependencies": {
+    "source-map": "^0.7.0"
   }
 }

packages/next-swc/crates/napi/Cargo.toml

@@ -57,6 +57,7 @@ tracing-subscriber = "0.3.9"
 tracing-chrome = "0.5.0"
 next-dev = { git = "https://github.com/vercel/turbo.git", rev = "a11422fdf6b1b3cde9072d90aab6d9eebfacb591", features = ["serializable"] }
 node-file-trace = { git = "https://github.com/vercel/turbo.git", rev = "a11422fdf6b1b3cde9072d90aab6d9eebfacb591", default-features = false, features = ["node-api"] }
+mdxjs = { version = "0.1.1", features = ["serializable"] }
 # There are few build targets we can't use native-tls which default features rely on,
 # allow to specify alternative (rustls) instead via features.
 # Note to opt in rustls default-features should be disabled

packages/next-swc/crates/napi/src/lib.rs

@@ -44,6 +44,7 @@ use swc_core::{
     common::{sync::Lazy, FilePathMapping, SourceMap},
 };
 
+pub mod mdx;
 pub mod minify;
 pub mod parse;
 pub mod transform;

packages/next-swc/crates/napi/src/mdx.rs

@@ -0,0 +1,43 @@
+use mdxjs::{compile, Options};
+use napi::bindgen_prelude::*;
+
+pub struct MdxCompileTask {
+    pub input: String,
+    pub option: Buffer,
+}
+
+impl Task for MdxCompileTask {
+    type Output = String;
+    type JsValue = String;
+
+    fn compute(&mut self) -> napi::Result<Self::Output> {
+        let options: Options = serde_json::from_slice(&self.option)?;
+        compile(&self.input, &options)
+            .map_err(|err| napi::Error::new(Status::GenericFailure, format!("{:?}", err)))
+    }
+
+    fn resolve(&mut self, _env: Env, output: Self::Output) -> napi::Result<Self::JsValue> {
+        Ok(output)
+    }
+}
+
+#[napi]
+pub fn mdx_compile(
+    value: String,
+    option: Buffer,
+    signal: Option<AbortSignal>,
+) -> napi::Result<AsyncTask<MdxCompileTask>> {
+    let task = MdxCompileTask {
+        input: value,
+        option,
+    };
+    Ok(AsyncTask::with_optional_signal(task, signal))
+}
+
+#[napi]
+pub fn mdx_compile_sync(value: String, option: Buffer) -> napi::Result<String> {
+    let option: Options = serde_json::from_slice(&option)?;
+
+    compile(value.as_str(), &option)
+        .map_err(|err| napi::Error::new(Status::GenericFailure, format!("{:?}", err)))
+}

packages/next-swc/crates/wasm/Cargo.toml

@@ -31,6 +31,7 @@ wasm-bindgen-futures = "0.4.8"
 getrandom = { version = "0.2.5", optional = true, default-features = false }
 js-sys = "0.3.59"
 serde-wasm-bindgen = "0.4.3"
+mdxjs = { version = "0.1.1", features = ["serializable"] }
 
 swc_core = { features = [
   "common_concurrent",

packages/next-swc/crates/wasm/src/lib.rs

@@ -13,6 +13,8 @@ use swc_core::{
     ecma::transforms::base::pass::noop,
 };
 
+pub mod mdx;
+
 fn convert_err(err: Error) -> JsValue {
     format!("{:?}", err).into()
 }

packages/next-swc/crates/wasm/src/mdx.rs

@@ -0,0 +1,21 @@
+use js_sys::JsString;
+use mdxjs::{compile, Options};
+use wasm_bindgen::prelude::*;
+use wasm_bindgen_futures::future_to_promise;
+
+#[wasm_bindgen(js_name = "mdxCompileSync")]
+pub fn mdx_compile_sync(value: JsString, opts: JsValue) -> Result<JsValue, JsValue> {
+    let value: String = value.into();
+    let option: Options = serde_wasm_bindgen::from_value(opts)?;
+
+    compile(value.as_str(), &option)
+        .map(|v| serde_wasm_bindgen::to_value(&v).expect("Should able to convert to JsValue"))
+        .map_err(|v| serde_wasm_bindgen::to_value(&v).expect("Should able to convert to JsValue"))
+}
+
+#[wasm_bindgen(js_name = "mdxCompile")]
+pub fn mdx_compile(value: JsString, opts: JsValue) -> js_sys::Promise {
+    // TODO: This'll be properly scheduled once wasm have standard backed thread
+    // support.
+    future_to_promise(async { mdx_compile_sync(value, opts) })
+}

packages/next/build/swc/index.js

@@ -215,6 +215,10 @@ async function loadWasm(importPath = '') {
             Log.error('Wasm binding does not support trace yet')
           },
         },
+        mdx: {
+          compile: (src, options) => bindings.mdxCompile(src, options),
+          compileSync: (src, options) => bindings.mdxCompileSync(src, options),
+        },
       }
       return wasmBindings
     } catch (e) {
@@ -352,6 +356,12 @@ function loadNative() {
         startTrace: (options = {}) =>
           bindings.runTurboTracing(toBuffer({ exact: true, ...options })),
       },
+      mdx: {
+        compile: (src, options) =>
+          bindings.mdxCompile(src, toBuffer(options ?? {})),
+        compileSync: (src, options) =>
+          bindings.mdxCompileSync(src, toBuffer(options ?? {})),
+      },
     }
     return nativeBindings
   }