Skip to content

Commit

Permalink
refactor(turbopack): Use graph for server action workaround (#72812)
Browse files Browse the repository at this point in the history 
### What?

Previously, we had a workaround while consuming the graph and converting them to `Module`s.

### Why?

Everything should be stored in the dependency graph because we are going to optimize it using graph node traversal.

### How?
  • Loading branch information
@kdy1 @wyattjoh
kdy1 authored and wyattjoh committed Nov 28, 2024
1 parent 11fe507 commit 0c61f77
Showing 1 changed file with 162 additions and 61 deletions.
223 changes: 162 additions & 61 deletions turbopack/crates/turbopack-ecmascript/src/tree_shake/graph.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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(),
})));
}
}

Expand Down Expand Up @@ -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,
};
Expand Down Expand Up @@ -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__";
Expand Down

0 comments on commit 0c61f77

Please sign in to comment.