Migrate server actions to single-graph-traversal

crates/next-api/src/app.rs

@@ -1121,8 +1121,21 @@ impl AppEndpoint {
             };
 
         let server_action_manifest_loader = if process_client_components {
+            let reduced_graphs = get_reduced_graphs_for_page(
+                this.app_project.project(),
+                *rsc_entry,
+                Vc::upcast(this.app_project.client_module_context()),
+            );
+            let actions = reduced_graphs.get_server_actions_for_page(
+                *rsc_entry,
+                match runtime {
+                    NextRuntime::Edge => Vc::upcast(this.app_project.edge_rsc_module_context()),
+                    NextRuntime::NodeJs => Vc::upcast(this.app_project.rsc_module_context()),
+                },
+            );
+
             let server_action_manifest = create_server_actions_manifest(
-                *ResolvedVc::upcast(app_entry.rsc_entry),
+                actions,
                 this.app_project.project().project_path(),
                 node_root,
                 app_entry.original_name.clone(),

crates/next-api/src/module_graph.rs

@@ -1,9 +1,13 @@
-use std::collections::{HashMap, HashSet};
+use std::{
+    borrow::Cow,
+    collections::{HashMap, HashSet},
+};
 
 use anyhow::Result;
 use next_core::{
     mode::NextMode,
     next_client_reference::{find_server_entries, ServerEntries},
+    next_manifests::ActionLayer,
 };
 use petgraph::{
     graph::{DiGraph, NodeIndex},
@@ -19,6 +23,7 @@ use turbopack_core::{
 use crate::{
     dynamic_imports::{map_next_dynamic, DynamicImportsHashMap},
     project::Project,
+    server_actions::{map_server_actions, to_rsc_context, AllActions, AllModuleActions},
 };
 
 #[turbo_tasks::value(cell = "new", eq = "manual", into = "new")]
@@ -221,13 +226,98 @@ impl NextDynamicGraph {
     }
 }
 
+#[turbo_tasks::value]
+pub struct ServerActionsGraph {
+    is_single_page: bool,
+    graph: ResolvedVc<SingleModuleGraph>,
+    /// (Layer, RSC or Browser module) -> list of actions
+    data: ResolvedVc<AllModuleActions>,
+}
+
+#[turbo_tasks::value_impl]
+impl ServerActionsGraph {
+    #[turbo_tasks::function]
+    pub async fn new_with_entries(
+        graph: ResolvedVc<SingleModuleGraph>,
+        is_single_page: bool,
+    ) -> Result<Vc<Self>> {
+        let mapped = map_server_actions(*graph);
+
+        // TODO shrink graph here
+
+        Ok(ServerActionsGraph {
+            is_single_page,
+            graph,
+            data: mapped.to_resolved().await?,
+        }
+        .cell())
+    }
+
+    #[turbo_tasks::function]
+    pub async fn get_server_actions_for_page(
+        &self,
+        entry: ResolvedVc<Box<dyn Module>>,
+        rsc_asset_context: Vc<Box<dyn AssetContext>>,
+    ) -> Result<Vc<AllActions>> {
+        let data = &*self.data.await?;
+        let data = if self.is_single_page {
+            // The graph contains the page (= `entry`) only, no need to filter.
+            Cow::Borrowed(data)
+        } else {
+            // The graph contains the whole app, traverse and collect all reachable imports.
+            let SingleModuleGraph { graph, entries } = &*self.graph.await?;
+
+            let mut result = HashMap::new();
+
+            let entry_node = *entries.get(&entry).unwrap();
+            let mut dfs = Dfs::new(&graph, entry_node);
+            while let Some(nx) = dfs.next(&graph) {
+                let weight = *graph.node_weight(nx).unwrap();
+                if let Some(node_data) = data.get(&weight) {
+                    result.insert(weight, *node_data);
+                }
+            }
+
+            Cow::Owned(result)
+        };
+
+        let actions = data
+            .iter()
+            .map(|(module, (layer, actions))| async move {
+                actions
+                    .await?
+                    .iter()
+                    .map(|(hash, name)| async move {
+                        Ok((
+                            hash.to_string(),
+                            (
+                                *layer,
+                                name.to_string(),
+                                if *layer == ActionLayer::Rsc {
+                                    *module
+                                } else {
+                                    to_rsc_context(**module, rsc_asset_context).await?
+                                },
+                            ),
+                        ))
+                    })
+                    .try_join()
+                    .await
+            })
+            .try_flat_join()
+            .await?;
+        Ok(Vc::cell(actions.into_iter().collect()))
+    }
+}
+
 /// The consumers of this shoudln't need to care about the exact contents since it's abstracted away
 /// by the accessor functions, but
 /// - In dev, contains information about the modules of the current page only
 /// - In prod, there is a single `ReducedGraphs` for the whole app, containing all pages
 #[turbo_tasks::value]
 pub struct ReducedGraphs {
     next_dynamic: Vec<ResolvedVc<NextDynamicGraph>>,
+    server_actions: Vec<ResolvedVc<ServerActionsGraph>>,
     // TODO add other graphs
 }
 
@@ -261,6 +351,33 @@ impl ReducedGraphs {
             Ok(DynamicImportsHashMap(result.into_iter().collect()).cell())
         }
     }
+
+    /// Returns the server actions for the given page.
+    #[turbo_tasks::function]
+    pub async fn get_server_actions_for_page(
+        &self,
+        entry: Vc<Box<dyn Module>>,
+        rsc_asset_context: Vc<Box<dyn AssetContext>>,
+    ) -> Result<Vc<AllActions>> {
+        if let [graph] = &self.server_actions[..] {
+            // Just a single graph, no need to merge results
+            Ok(graph.get_server_actions_for_page(entry, rsc_asset_context))
+        } else {
+            let result = self
+                .server_actions
+                .iter()
+                .map(|graph| async move {
+                    Ok(graph
+                        .get_server_actions_for_page(entry, rsc_asset_context)
+                        .await?
+                        .clone_value())
+                })
+                .try_flat_join()
+                .await?;
+
+            Ok(Vc::cell(result.into_iter().collect()))
+        }
+    }
 }
 
 /// Generates a [ReducedGraph] for the given project and page containing information that is either
@@ -294,5 +411,15 @@ pub async fn get_reduced_graphs_for_page(
         .try_join()
         .await?;
 
-    Ok(ReducedGraphs { next_dynamic }.cell())
+    let server_actions = graphs
+        .iter()
+        .map(|graph| ServerActionsGraph::new_with_entries(**graph, is_single_page).to_resolved())
+        .try_join()
+        .await?;
+
+    Ok(ReducedGraphs {
+        next_dynamic,
+        server_actions,
+    }
+    .cell())
 }