add different keyed app route elements

crates/next-core/js/src/entry/app-renderer.tsx

@@ -1,9 +1,21 @@
 import type { Ipc } from "@vercel/turbopack-next/internal/ipc";
 
 // Provided by the rust generate code
+type FileType =
+  | "layout"
+  | "template"
+  | "error"
+  | "loading"
+  | "not-found"
+  | "head";
 declare global {
   // an array of all layouts and the page
-  const LAYOUT_INFO: { segment: string; module: any; chunks: string[] }[];
+  const LAYOUT_INFO: ({
+    segment: string;
+    page?: { module: any; chunks: string[] };
+  } & {
+    [componentKey in FileType]?: { module: any; chunks: string[] };
+  })[];
   // array of chunks for the bootstrap script
   const BOOTSTRAP: string[];
   const IPC: Ipc<unknown, unknown>;
@@ -21,7 +33,6 @@ import "next/dist/server/node-polyfill-web-streams";
 import { RenderOpts, renderToHTMLOrFlight } from "next/dist/server/app-render";
 import { PassThrough } from "stream";
 import { ServerResponseShim } from "@vercel/turbopack-next/internal/http";
-import { structuredError } from "@vercel/turbopack-next/internal/error";
 import { ParsedUrlQuery } from "node:querystring";
 
 globalThis.__next_require__ = (data) => {
@@ -78,15 +89,10 @@ type IpcOutgoingMessage = {
 // TODO expose these types in next.js
 type ComponentModule = () => any;
 export type ComponentsType = {
-  readonly [componentKey in
-    | "layout"
-    | "template"
-    | "error"
-    | "loading"
-    | "not-found"]?: ComponentModule;
+  [componentKey in FileType]?: ComponentModule;
 } & {
-  readonly layoutOrPagePath?: string;
-  readonly page?: ComponentModule;
+  layoutOrPagePath?: string;
+  page?: ComponentModule;
 };
 type LoaderTree = [
   segment: string,
@@ -103,30 +109,27 @@ type ServerComponentsManifestModule = {
 
 async function runOperation(renderData: RenderData) {
   const pageItem = LAYOUT_INFO[LAYOUT_INFO.length - 1];
-  const pageModule = pageItem.module;
+  const pageModule = pageItem.page!.module;
   const Page = pageModule.default;
   let tree: LoaderTree = [
     "",
     {},
     { page: () => Page, layoutOrPagePath: "page.js" },
   ];
+  const layoutInfoChunks: string[][] = [];
   for (let i = LAYOUT_INFO.length - 2; i >= 0; i--) {
     const info = LAYOUT_INFO[i];
-    const mod = info.module;
-    if (mod) {
-      const Layout = mod.default;
-      tree = [
-        info.segment,
-        { children: tree },
-        { layout: () => Layout, layoutOrPagePath: `layout${i}.js` },
-      ];
-    } else {
-      tree = [
-        info.segment,
-        { children: tree },
-        { layoutOrPagePath: `layout${i}.js` },
-      ];
+    const components: ComponentsType = { layoutOrPagePath: `layout${i}.js` };
+    layoutInfoChunks[i] = [];
+    for (const key of Object.keys(info)) {
+      if (key === "segment") {
+        continue;
+      }
+      const k = key as FileType;
+      components[k] = () => info[k]!.module.default;
+      layoutInfoChunks[i].push(...info[k]!.chunks);
     }
+    tree = [info.segment, { children: tree }, components];
   }
 
   const proxyMethodsForModule = (
@@ -157,16 +160,16 @@ async function runOperation(renderData: RenderData) {
   const manifest: FlightManifest = new Proxy({} as any, proxyMethods(false));
   const serverCSSManifest: FlightCSSManifest = {};
   serverCSSManifest.__entry_css__ = {};
-  for (let i = 0; i < LAYOUT_INFO.length - 1; i++) {
-    const { chunks } = LAYOUT_INFO[i];
-    const cssChunks = (chunks || []).filter((path) => path.endsWith(".css"));
+  for (let i = 0; i < layoutInfoChunks.length; i++) {
+    const chunks = layoutInfoChunks[i];
+    const cssChunks = chunks.filter((path) => path.endsWith(".css"));
     serverCSSManifest[`layout${i}.js`] = cssChunks.map((chunk) =>
       JSON.stringify([chunk, [chunk]])
     );
   }
   serverCSSManifest.__entry_css__ = {
-    page: pageItem.chunks
-      .filter((path) => path.endsWith(".css"))
+    page: pageItem
+      .page!.chunks.filter((path) => path.endsWith(".css"))
       .map((chunk) => JSON.stringify([chunk, [chunk]])),
   };
   serverCSSManifest["page.js"] = serverCSSManifest.__entry_css__.page;
@@ -185,7 +188,8 @@ async function runOperation(renderData: RenderData) {
     dev: true,
     buildManifest: {
       polyfillFiles: [],
-      rootMainFiles: LAYOUT_INFO.flatMap(({ chunks }) => chunks || [])
+      rootMainFiles: layoutInfoChunks
+        .flat()
         .concat(BOOTSTRAP)
         .filter((path) => !path.endsWith(".css")),
       devFiles: [],

crates/next-core/src/app_render/mod.rs

@@ -1,10 +1,12 @@
+use std::collections::HashMap;
+
 use turbo_tasks_fs::FileSystemPathVc;
 
 pub mod next_layout_entry_transition;
 
 #[turbo_tasks::value(shared)]
 pub struct LayoutSegment {
-    pub file: Option<FileSystemPathVc>,
+    pub files: HashMap<String, FileSystemPathVc>,
     pub target: FileSystemPathVc,
 }
 

crates/next-core/src/app_source.rs

@@ -1,4 +1,7 @@
-use std::{collections::HashMap, io::Write};
+use std::{
+    collections::{BTreeMap, HashMap},
+    io::Write,
+};
 
 use anyhow::{anyhow, Context, Result};
 use turbo_tasks::{
@@ -284,20 +287,18 @@ async fn create_app_source_for_directory(
     let mut layouts = layouts;
     let mut sources = Vec::new();
     let mut page = None;
-    let mut layout = None;
+    let mut files = HashMap::new();
     if let DirectoryContent::Entries(entries) = &*input_dir.read_dir().await? {
         for (name, entry) in entries.iter() {
-            if let DirectoryEntry::File(file) = entry {
+            if let &DirectoryEntry::File(file) = entry {
                 if let Some((name, _)) = name.rsplit_once('.') {
                     match name {
                         "page" => {
                             page = Some(file);
                         }
-                        "layout" => {
-                            layout = Some(file);
+                        "layout" | "error" | "loading" | "template" | "not-found" | "head" => {
+                            files.insert(name.to_string(), file);
                         }
-                        "error" => { /* TODO */ }
-                        "loading" => { /* TODO */ }
                         _ => {
                             // Any other file is ignored
                         }
@@ -306,8 +307,7 @@ async fn create_app_source_for_directory(
             }
         }
 
-        let layout_js = input_dir.join("layout.js");
-        let layout_tsx = input_dir.join("layout.tsx");
+        let layout = files.get("layout");
 
         // TODO: Use let Some(page_file) = page in expression below when
         // https://rust-lang.github.io/rfcs/2497-if-let-chains.html lands
@@ -320,23 +320,23 @@ async fn create_app_source_for_directory(
             // stable does.
             let is_tsx = *page_file.extension().await? == "tsx";
 
-            if is_tsx {
-                layout.replace(&layout_tsx);
+            let layout = if is_tsx {
+                input_dir.join("layout.tsx")
             } else {
-                layout.replace(&layout_js);
-            }
+                input_dir.join("layout.js")
+            };
+            files.insert("layout".to_string(), layout);
             let content = if is_tsx {
                 include_str!("assets/layout.tsx")
             } else {
                 include_str!("assets/layout.js")
             };
 
-            let layout = layout.context("required")?;
             layout.write(FileContentVc::from(File::from(content)));
 
             AppSourceIssue {
                 severity: IssueSeverity::Warning.into(),
-                path: *page_file,
+                path: page_file,
                 message: StringVc::cell(format!(
                     "Your page {} did not have a root layout, we created {} for you.",
                     page_file.await?.path,
@@ -349,15 +349,9 @@ async fn create_app_source_for_directory(
         }
 
         let mut list = layouts.await?.clone_value();
-        list.push(
-            LayoutSegment {
-                file: layout.copied(),
-                target,
-            }
-            .cell(),
-        );
+        list.push(LayoutSegment { files, target }.cell());
         layouts = LayoutSegmentsVc::cell(list);
-        if let Some(page_path) = page.copied() {
+        if let Some(page_path) = page {
             sources.push(create_node_rendered_source(
                 server_root,
                 regular_expression_for_path(server_root, target, false),
@@ -438,14 +432,14 @@ impl NodeEntry for AppRenderer {
             .copied()
             .chain(std::iter::once(
                 LayoutSegment {
-                    file: Some(page),
+                    files: HashMap::from([("page".to_string(), page)]),
                     target: self.target,
                 }
                 .cell(),
             ))
             .try_join()
             .await?;
-        let segments: Vec<(String, Option<(String, String, String)>)> = layout_and_page
+        let segments: Vec<(String, BTreeMap<String, (String, String, String)>)> = layout_and_page
             .into_iter()
             .fold(
                 (self.server_root, Vec::new()),
@@ -455,7 +449,8 @@ impl NodeEntry for AppRenderer {
                             let target = &*segment.target.await?;
                             let segment_path =
                                 last_path.await?.get_path_to(target).unwrap_or_default();
-                            if let Some(file) = segment.file {
+                            let mut imports = BTreeMap::new();
+                            for (key, file) in segment.files.iter() {
                                 let file_str = file.to_string().await?;
                                 let identifier = magic_identifier::encode(&format!(
                                     "imported namespace {}",
@@ -466,25 +461,21 @@ impl NodeEntry for AppRenderer {
                                     file_str
                                 ));
                                 if let Some(p) = path_value.get_relative_path_to(&*file.await?) {
-                                    Ok((
-                                        stringify_str(segment_path),
-                                        Some((p, identifier, chunks_identifier)),
-                                    ))
+                                    imports.insert(
+                                        key.to_string(),
+                                        (p, identifier, chunks_identifier),
+                                    );
                                 } else {
-                                    Err(anyhow!(
+                                    return Err(anyhow!(
                                         "Unable to generate import as there
                                 is no relative path to the layout module {} from context
                                 path {}",
                                         file_str,
                                         path.to_string().await?
-                                    ))
+                                    ));
                                 }
-                            } else {
-                                Ok::<(String, Option<(String, String, String)>), _>((
-                                    stringify_str(segment_path),
-                                    None,
-                                ))
                             }
+                            Ok((stringify_str(segment_path), imports))
                         });
                         futures
                     })
@@ -498,8 +489,8 @@ impl NodeEntry for AppRenderer {
             "import IPC, { Ipc } from \"@vercel/turbopack-next/internal/ipc\";\n",
         );
 
-        for (_, import) in segments.iter() {
-            if let Some((p, identifier, chunks_identifier)) = import {
+        for (_, imports) in segments.iter() {
+            for (p, identifier, chunks_identifier) in imports.values() {
                 result += r#"("TURBOPACK { transition: next-layout-entry; chunking-type: parallel }");
 "#;
                 writeln!(
@@ -522,16 +513,16 @@ import BOOTSTRAP from {};
         }
 
         result += "const LAYOUT_INFO = [";
-        for (segment_str_lit, import) in segments.iter() {
-            if let Some((_, identifier, chunks_identifier)) = import {
+        for (segment_str_lit, imports) in segments.iter() {
+            writeln!(result, "  {{\n    segment: {segment_str_lit},")?;
+            for (key, (_, identifier, chunks_identifier)) in imports {
                 writeln!(
                     result,
-                    "  {{ segment: {segment_str_lit}, module: {identifier}, chunks: \
-                     {chunks_identifier} }},",
-                )?
-            } else {
-                writeln!(result, "  {{ segment: {segment_str_lit} }},",)?
+                    "    {key}: {{ module: {identifier}, chunks: {chunks_identifier} }},",
+                    key = stringify_str(key),
+                )?;
             }
+            result += "  },";
         }
         result += "];\n\n";
 

crates/turbo-tasks/src/manager.rs

@@ -16,7 +16,7 @@ use std::{
     time::{Duration, Instant},
 };
 
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, Context, Result};
 use event_listener::{Event, EventListener};
 use futures::FutureExt;
 use serde::{de::Visitor, Deserialize, Serialize};
@@ -355,13 +355,22 @@ impl<B: Backend> TurboTasks<B> {
                             FormatDuration(duration)
                         )
                     }
-                    let result = result.map_err(|any| match any.downcast::<String>() {
-                        Ok(owned) => Some(Cow::Owned(*owned)),
-                        Err(any) => match any.downcast::<&'static str>() {
-                            Ok(str) => Some(Cow::Borrowed(*str)),
-                            Err(_) => None,
-                        },
-                    });
+                    let result = result
+                        .map(|r| {
+                            r.with_context(|| {
+                                format!(
+                                    "while executing {}",
+                                    this.backend.get_task_description(task_id)
+                                )
+                            })
+                        })
+                        .map_err(|any| match any.downcast::<String>() {
+                            Ok(owned) => Some(Cow::Owned(*owned)),
+                            Err(any) => match any.downcast::<&'static str>() {
+                                Ok(str) => Some(Cow::Borrowed(*str)),
+                                Err(_) => None,
+                            },
+                        });
                     this.backend.task_execution_result(task_id, result, &*this);
                     this.notify_scheduled_tasks_internal();
                     let reexecute = this.backend.task_execution_completed(