[turbopack] move edge entry wrapper to build template (#86699)

Move the edge-wrapper virtual module to templates

Author: huozhi

crates/next-core/src/next_edge/entry.rs

@@ -1,62 +1,39 @@
 use anyhow::Result;
-use indoc::formatdoc;
 use turbo_rcstr::{RcStr, rcstr};
 use turbo_tasks::{ResolvedVc, Vc, fxindexmap};
-use turbo_tasks_fs::{File, FileContent, FileSystemPath};
-use turbopack_core::{
-    asset::AssetContent, context::AssetContext, module::Module, reference_type::ReferenceType,
-    virtual_source::VirtualSource,
-};
-use turbopack_ecmascript::utils::StringifyJs;
+use turbo_tasks_fs::FileSystemPath;
+use turbopack_core::{context::AssetContext, module::Module, reference_type::ReferenceType};
+
+use crate::util::load_next_js_template_no_imports;
 
 #[turbo_tasks::function]
-pub fn wrap_edge_entry(
+pub async fn wrap_edge_entry(
     asset_context: Vc<Box<dyn AssetContext>>,
     project_root: FileSystemPath,
     entry: ResolvedVc<Box<dyn Module>>,
     pathname: RcStr,
 ) -> Result<Vc<Box<dyn Module>>> {
-    // The wrapped module could be an async module, we handle that with the proxy
-    // here. The comma expression makes sure we don't call the function with the
-    // module as the "this" arg.
-    // Turn exports into functions that are also a thenable. This way you can await the whole object
-    // or  exports (e.g. for Components) or call them directly as though they are async functions
-    // (e.g. edge functions/middleware, this is what the Edge Runtime does).
-    // Catch promise to prevent UnhandledPromiseRejectionWarning, this will be propagated through
-    // the awaited export(s) anyway.
-    let source = formatdoc!(
-        r#"
-            self._ENTRIES ||= {{}};
-            const modProm = import('MODULE');
-            modProm.catch(() => {{}});
-            self._ENTRIES[{}] = new Proxy(modProm, {{
-                get(modProm, name) {{
-                    if (name === "then") {{
-                        return (res, rej) => modProm.then(res, rej);
-                    }}
-                    let result = (...args) => modProm.then((mod) => (0, mod[name])(...args));
-                    result.then = (res, rej) => modProm.then((mod) => mod[name]).then(res, rej);
-                    return result;
-                }},
-            }});
-        "#,
-        StringifyJs(&format_args!("middleware_{pathname}"))
-    );
-    let file = File::from(source);
-
-    // TODO(alexkirsz) Figure out how to name this virtual asset.
-    let virtual_source = VirtualSource::new(
-        project_root.join("edge-wrapper.js")?,
-        AssetContent::file(FileContent::Content(file).cell()),
-    );
+    // The actual wrapper lives in the Next.js templates directory as `edge-wrapper.js`.
+    // We use the template expansion helper so this code is kept in sync with other
+    // Next.js runtime templates. This particular template does not have any imports
+    // of its own, so we use the variant that allows templates without relative
+    // imports to be rewritten.
+    let template_source = load_next_js_template_no_imports(
+        "edge-wrapper.js",
+        project_root,
+        &[("VAR_ENTRY_NAME", &format!("middleware_{pathname}"))],
+        &[],
+        &[],
+    )
+    .await?;
 
     let inner_assets = fxindexmap! {
         rcstr!("MODULE") => entry
     };
 
     Ok(asset_context
         .process(
-            Vc::upcast(virtual_source),
+            template_source,
             ReferenceType::Internal(ResolvedVc::cell(inner_assets)),
         )
         .module())

crates/next-core/src/util.rs

@@ -1,7 +1,7 @@
 use std::{fmt::Display, str::FromStr};
 
 use anyhow::{Result, anyhow, bail};
-use next_taskless::expand_next_js_template;
+use next_taskless::{expand_next_js_template, expand_next_js_template_no_imports};
 use serde::{Deserialize, Serialize, de::DeserializeOwned};
 use turbo_rcstr::{RcStr, rcstr};
 use turbo_tasks::{FxIndexMap, NonLocalValue, TaskInput, Vc, trace::TraceRawVcs};
@@ -268,6 +268,41 @@ pub async fn load_next_js_template(
     Ok(Vc::upcast(source))
 }
 
+/// Loads a next.js template but does **not** require that any relative imports are present
+/// or rewritten. This is intended for small internal templates that do not have their own
+/// imports but still use template variables/injections.
+pub async fn load_next_js_template_no_imports(
+    template_path: &str,
+    project_path: FileSystemPath,
+    replacements: &[(&str, &str)],
+    injections: &[(&str, &str)],
+    imports: &[(&str, Option<&str>)],
+) -> Result<Vc<Box<dyn Source>>> {
+    let template_path = virtual_next_js_template_path(project_path.clone(), template_path).await?;
+
+    let content = file_content_rope(template_path.read()).await?;
+    let content = content.to_str()?;
+
+    let package_root = get_next_package(project_path).await?;
+
+    let content = expand_next_js_template_no_imports(
+        &content,
+        &template_path.path,
+        &package_root.path,
+        replacements.iter().copied(),
+        injections.iter().copied(),
+        imports.iter().copied(),
+    )?;
+
+    let file = File::from(content);
+    let source = VirtualSource::new(
+        template_path,
+        AssetContent::file(FileContent::Content(file).cell()),
+    );
+
+    Ok(Vc::upcast(source))
+}
+
 #[turbo_tasks::function]
 pub async fn file_content_rope(content: Vc<FileContent>) -> Result<Vc<Rope>> {
     let content = &*content.await?;

crates/next-taskless/src/lib.rs

@@ -25,6 +25,48 @@ pub fn expand_next_js_template<'a>(
     replacements: impl IntoIterator<Item = (&'a str, &'a str)>,
     injections: impl IntoIterator<Item = (&'a str, &'a str)>,
     imports: impl IntoIterator<Item = (&'a str, Option<&'a str>)>,
+) -> Result<String> {
+    expand_next_js_template_inner(
+        content,
+        template_path,
+        next_package_dir_path,
+        replacements,
+        injections,
+        imports,
+        true,
+    )
+}
+
+/// Same as [`expand_next_js_template`], but does not enforce that at least one relative
+/// import is present and rewritten. This is useful for very small templates that only
+/// use template variables/injections and have no imports of their own.
+pub fn expand_next_js_template_no_imports<'a>(
+    content: &str,
+    template_path: &str,
+    next_package_dir_path: &str,
+    replacements: impl IntoIterator<Item = (&'a str, &'a str)>,
+    injections: impl IntoIterator<Item = (&'a str, &'a str)>,
+    imports: impl IntoIterator<Item = (&'a str, Option<&'a str>)>,
+) -> Result<String> {
+    expand_next_js_template_inner(
+        content,
+        template_path,
+        next_package_dir_path,
+        replacements,
+        injections,
+        imports,
+        false,
+    )
+}
+
+fn expand_next_js_template_inner<'a>(
+    content: &str,
+    template_path: &str,
+    next_package_dir_path: &str,
+    replacements: impl IntoIterator<Item = (&'a str, &'a str)>,
+    injections: impl IntoIterator<Item = (&'a str, &'a str)>,
+    imports: impl IntoIterator<Item = (&'a str, Option<&'a str>)>,
+    require_import_replacement: bool,
 ) -> Result<String> {
     let template_parent_path = normalize_path(get_parent_path(template_path))
         .context("failed to normalize template path")?;
@@ -92,10 +134,11 @@ pub fn expand_next_js_template<'a>(
     })
     .context("replacing imports failed")?;
 
-    // Verify that at least one import was replaced. It's the case today where every template file
-    // has at least one import to update, so this ensures that we don't accidentally remove the
-    // import replacement code or use the wrong template file.
-    if count == 0 {
+    // Verify that at least one import was replaced when required. It's the case today where every
+    // template file (except a few small internal helpers) has at least one import to update, so
+    // this ensures that we don't accidentally remove the import replacement code or use the wrong
+    // template file.
+    if require_import_replacement && count == 0 {
         bail!("Invariant: Expected to replace at least one import")
     }
 

packages/next/src/build/templates/edge-wrapper.js

@@ -0,0 +1,23 @@
+// The wrapped module could be an async module, we handle that with the proxy
+// here. The comma expression makes sure we don't call the function with the
+// module as the "this" arg.
+// Turn exports into functions that are also a thenable. This way you can await the whole object
+// or  exports (e.g. for Components) or call them directly as though they are async functions
+// (e.g. edge functions/middleware, this is what the Edge Runtime does).
+// Catch promise to prevent UnhandledPromiseRejectionWarning, this will be propagated through
+// the awaited export(s) anyway.
+self._ENTRIES ||= {}
+const modProm = import('MODULE')
+modProm.catch(() => {})
+self._ENTRIES['VAR_ENTRY_NAME'] = new Proxy(modProm, {
+  get(innerModProm, name) {
+    if (name === 'then') {
+      return (res, rej) => innerModProm.then(res, rej)
+    }
+    let result = (...args) =>
+      innerModProm.then((mod) => (0, mod[name])(...args))
+    result.then = (res, rej) =>
+      innerModProm.then((mod) => mod[name]).then(res, rej)
+    return result
+  },
+})