refactor(turbopack): Use graph for server action workaround

turbopack/crates/turbopack-ecmascript/src/tree_shake/graph.rs

@@ -2,8 +2,11 @@ use std::{fmt, hash::Hash};
 
 use petgraph::{
     algo::{condensation, has_path_connecting},
+    graph::NodeIndex,
     graphmap::GraphMap,
     prelude::DiGraphMap,
+    visit::EdgeRef,
+    Direction, Graph,
 };
 use rustc_hash::{FxHashMap, FxHashSet};
 use swc_core::{
@@ -257,7 +260,7 @@ impl DepGraph {
     ///
     /// Note: ESM imports are immutable, but we do not handle it.
     pub(super) fn split_module(
-        &self,
+        &mut self,
         directives: &[ModuleItem],
         data: &FxHashMap<ItemId, ItemData>,
     ) -> SplitModuleResult {
@@ -381,13 +384,10 @@ impl DepGraph {
                 }
             }
 
-            let mut use_export_instead_of_declarator = false;
-
             for item in group {
                 match item {
                     ItemId::Group(ItemIdGroupKind::Export(..)) => {
                         if let Some(export) = &data[item].export {
-                            use_export_instead_of_declarator = true;
                             outputs.insert(Key::Export(export.as_str().into()), ix as u32);
 
                             let s = ExportSpecifier::Named(ExportNamedSpecifier {
@@ -421,68 +421,50 @@ impl DepGraph {
                 }
             }
 
-            // Workaround for implcit export issue of server actions.
-            //
-            // Inline server actions require the generated `$$RSC_SERVER_0` to be **exported**.
-            //
-            // But tree shaking works by removing unused code, and the **export** of $$RSC_SERVER_0
-            // is cleary not used from the external module as it does not exist at all
-            // in the user code.
-            //
-            // So we need to add an import for $$RSC_SERVER_0 to the module, so that the export is
-            // preserved.
-            if use_export_instead_of_declarator {
-                for (other_ix, other_group) in groups.graph_ix.iter().enumerate() {
-                    if other_ix == ix {
-                        continue;
-                    }
-
-                    let deps = part_deps.entry(ix as u32).or_default();
-
-                    for other_item in other_group {
-                        if let ItemId::Group(ItemIdGroupKind::Export(export, _)) = other_item {
-                            if !export.0.as_str().starts_with("$$RSC_SERVER_") {
-                                continue;
-                            }
+            for dep in groups
+                .idx_graph
+                .neighbors_directed(ix as u32, Direction::Outgoing)
+            {
+                if dep == ix as u32 {
+                    continue;
+                }
 
-                            let Some(&declarator) = declarator.get(export) else {
-                                continue;
-                            };
+                let dep_item_ids = groups.graph_ix.get_index(dep as usize).unwrap();
 
-                            if declarator == ix as u32 {
-                                continue;
-                            }
+                for dep_item_id in dep_item_ids {
+                    let ItemId::Group(ItemIdGroupKind::Export(var, export)) = dep_item_id else {
+                        continue;
+                    };
 
-                            if !has_path_connecting(&groups.idx_graph, ix as u32, declarator, None)
-                            {
-                                continue;
-                            }
+                    if !export.starts_with("$$RSC_SERVER_") {
+                        continue;
+                    }
 
-                            let s = ImportSpecifier::Named(ImportNamedSpecifier {
-                                span: DUMMY_SP,
-                                local: export.clone().into(),
-                                imported: None,
-                                is_type_only: false,
-                            });
+                    required_vars.swap_remove(var);
 
-                            required_vars.swap_remove(export);
+                    let dep_part_id = PartId::Export(export.as_str().into());
+                    let specifiers = vec![ImportSpecifier::Named(ImportNamedSpecifier {
+                        span: DUMMY_SP,
+                        local: var.clone().into(),
+                        imported: None,
+                        is_type_only: false,
+                    })];
 
-                            deps.push(PartId::Export(export.0.as_str().into()));
+                    part_deps
+                        .entry(ix as u32)
+                        .or_default()
+                        .push(dep_part_id.clone());
 
-                            chunk.body.push(ModuleItem::ModuleDecl(ModuleDecl::Import(
-                                ImportDecl {
-                                    span: DUMMY_SP,
-                                    specifiers: vec![s],
-                                    src: Box::new(TURBOPACK_PART_IMPORT_SOURCE.into()),
-                                    type_only: false,
-                                    with: Some(Box::new(create_turbopack_part_id_assert(
-                                        PartId::Export(export.0.as_str().into()),
-                                    ))),
-                                    phase: Default::default(),
-                                },
-                            )));
-                        }
-                    }
+                    chunk
+                        .body
+                        .push(ModuleItem::ModuleDecl(ModuleDecl::Import(ImportDecl {
+                            span: DUMMY_SP,
+                            specifiers,
+                            src: Box::new(TURBOPACK_PART_IMPORT_SOURCE.into()),
+                            type_only: false,
+                            with: Some(Box::new(create_turbopack_part_id_assert(dep_part_id))),
+                            phase: Default::default(),
+                        })));
                 }
             }
 
@@ -713,12 +695,15 @@ impl DepGraph {
     /// Note that [ModuleItem] and [Module] are represented as [ItemId] for
     /// performance.
     pub(super) fn finalize(
-        &self,
+        &mut self,
         data: &FxHashMap<ItemId, ItemData>,
     ) -> InternedGraph<Vec<ItemId>> {
-        let graph = self.g.idx_graph.clone().into_graph::<u32>();
+        let mut graph = self.g.idx_graph.clone().into_graph::<u32>();
+
+        self.workaround_server_action(&mut graph, data);
 
         let mut condensed = condensation(graph, true);
+
         let optimizer = GraphOptimizer {
             graph_ix: &self.g.graph_ix,
         };
@@ -1458,6 +1443,122 @@ impl DepGraph {
 
         has_path_connecting(&self.g.idx_graph, from, to, None)
     }
+
+    /// Workaround for implcit export issue of server actions.
+    ///
+    /// Inline server actions require the generated `$$RSC_SERVER_0` to be **exported**.
+    ///
+    /// But tree shaking works by removing unused code, and the **export** of $$RSC_SERVER_0
+    /// is cleary not used from the external module as it does not exist at all
+    /// in the user code.
+    ///
+    /// So we need to add an import for $$RSC_SERVER_0 to the module, so that the export is
+    /// preserved.
+    fn workaround_server_action(
+        &mut self,
+        g: &mut Graph<u32, Dependency>,
+        data: &FxHashMap<ItemId, ItemData>,
+    ) {
+        fn collect_deps(
+            g: &Graph<u32, Dependency>,
+            done: &mut FxHashSet<NodeIndex>,
+            node: NodeIndex,
+        ) -> Vec<NodeIndex> {
+            let direct_deps = g
+                .edges_directed(node, Direction::Outgoing)
+                .map(|e| e.target())
+                .collect::<Vec<_>>();
+
+            if direct_deps.iter().all(|dep| done.contains(dep)) {
+                return direct_deps;
+            }
+
+            direct_deps
+                .into_iter()
+                .flat_map(|dep| {
+                    let mut v = if !done.insert(dep) {
+                        vec![]
+                    } else {
+                        collect_deps(g, done, dep)
+                    };
+
+                    v.push(dep);
+                    v
+                })
+                .collect()
+        }
+
+        let mut server_action_decls = FxHashMap::default();
+        let mut server_action_exports = FxHashMap::default();
+
+        for node in g.node_indices() {
+            let Some(ix) = g.node_weight(node) else {
+                continue;
+            };
+
+            let item_id = self.g.graph_ix.get_index(*ix as _).unwrap();
+
+            if let ItemId::Group(ItemIdGroupKind::Export(v, name)) = item_id {
+                if name.starts_with("$$RSC_SERVER_") {
+                    server_action_exports.insert(v.0.clone(), node);
+                }
+            }
+
+            let item_data = &data[item_id];
+
+            for v in item_data.var_decls.iter() {
+                if v.0.starts_with("$$RSC_SERVER_") {
+                    server_action_decls.insert(node, v.0.clone());
+                }
+            }
+        }
+
+        if server_action_decls.is_empty() || server_action_exports.is_empty() {
+            return;
+        }
+
+        let mut queue = vec![];
+
+        for node in g.node_indices() {
+            let Some(ix) = g.node_weight(node) else {
+                continue;
+            };
+
+            let is_export_node = {
+                let item_id = self.g.graph_ix.get_index(*ix as _).unwrap();
+                matches!(item_id, ItemId::Group(ItemIdGroupKind::Export(..)))
+            };
+
+            if !is_export_node {
+                continue;
+            }
+
+            // If an export uses $$RSC_SERVER_0, depend on "export $$RSC_SERVER_0"
+
+            let mut done = FxHashSet::default();
+            let dependencies = collect_deps(g, &mut done, node);
+
+            for &dependency in dependencies.iter() {
+                if dependency == node {
+                    continue;
+                }
+
+                let Some(action_item_id) = server_action_decls.get(&dependency) else {
+                    continue;
+                };
+
+                let Some(action_export_node) = server_action_exports.get(action_item_id) else {
+                    continue;
+                };
+
+                queue.push((node, *action_export_node));
+            }
+        }
+
+        for (export_node, dep) in queue {
+            g.add_edge(export_node, dep, Dependency::Strong);
+        }
+    }
 }
 
 const ASSERT_CHUNK_KEY: &str = "__turbopack_part__";