feat(turbopack-ecmascript): implement acyclic SCC graph for ESM impor…

…ts (vercel/turborepo#5506)

crates/turbo-tasks/src/join_iter_ext.rs

@@ -1,4 +1,8 @@
-use std::future::{Future, IntoFuture};
+use std::{
+    future::{Future, IntoFuture},
+    pin::Pin,
+    task::Poll,
+};
 
 use anyhow::Result;
 use futures::{
@@ -108,3 +112,68 @@ where
         }
     }
 }
+
+pin_project! {
+    /// Future for the [TryFlatJoinIterExt::try_flat_join] method.
+    pub struct TryFlatJoin<F>
+    where
+        F: Future,
+    {
+        #[pin]
+        inner: JoinAll<F>,
+    }
+}
+
+impl<F, I, U> Future for TryFlatJoin<F>
+where
+    F: Future<Output = Result<I>>,
+    I: IntoIterator<IntoIter = U, Item = U::Item>,
+    U: Iterator,
+{
+    type Output = Result<Vec<U::Item>>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
+        match self.project().inner.poll_unpin(cx) {
+            Poll::Ready(res) => {
+                let mut v = Vec::new();
+                for r in res {
+                    v.extend(r?);
+                }
+
+                Poll::Ready(Ok(v))
+            }
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}
+
+pub trait TryFlatJoinIterExt<F, I, U>: Iterator
+where
+    F: Future<Output = Result<I>>,
+    I: IntoIterator<IntoIter = U, Item = U::Item>,
+    U: Iterator,
+{
+    /// Returns a future that resolves to a vector of the outputs of the futures
+    /// in the iterator, or to an error if one of the futures fail.
+    ///
+    /// It also flattens the result.
+    ///
+    /// Unlike `Futures::future::try_join_all`, this returns the Error that
+    /// occurs first in the list of futures, not the first to fail in time.
+    fn try_flat_join(self) -> TryFlatJoin<F>;
+}
+
+impl<F, IF, It, I, U> TryFlatJoinIterExt<F, I, U> for It
+where
+    F: Future<Output = Result<I>>,
+    IF: IntoFuture<Output = Result<I>, IntoFuture = F>,
+    It: Iterator<Item = IF>,
+    I: IntoIterator<IntoIter = U, Item = U::Item>,
+    U: Iterator,
+{
+    fn try_flat_join(self) -> TryFlatJoin<F> {
+        TryFlatJoin {
+            inner: join_all(self.map(|f| f.into_future())),
+        }
+    }
+}

crates/turbo-tasks/src/lib.rs

@@ -83,7 +83,7 @@ pub use id::{
 pub use invalidation::{
     DynamicEqHash, InvalidationReason, InvalidationReasonKind, InvalidationReasonSet,
 };
-pub use join_iter_ext::{JoinIterExt, TryJoinIterExt};
+pub use join_iter_ext::{JoinIterExt, TryFlatJoinIterExt, TryJoinIterExt};
 pub use manager::{
     dynamic_call, emit, get_invalidator, mark_finished, mark_stateful, run_once,
     run_once_with_reason, spawn_blocking, spawn_thread, trait_call, turbo_tasks, Invalidator,

crates/turbopack-core/src/asset.rs

@@ -25,6 +25,15 @@ impl Assets {
     }
 }
 
+#[turbo_tasks::value_impl]
+impl AssetsSet {
+    /// Creates an empty set of [Asset]s
+    #[turbo_tasks::function]
+    pub fn empty() -> Vc<AssetsSet> {
+        Vc::cell(IndexSet::new())
+    }
+}
+
 /// An asset. It also forms a graph when following [Asset::references].
 #[turbo_tasks::value_trait]
 pub trait Asset {

crates/turbopack-core/src/chunk/available_assets.rs

@@ -84,7 +84,7 @@ impl AvailableAssets {
 }
 
 #[turbo_tasks::function]
-async fn chunkable_assets_set(root: Vc<Box<dyn Module>>) -> Result<Vc<ModulesSet>> {
+pub async fn chunkable_assets_set(root: Vc<Box<dyn Module>>) -> Result<Vc<ModulesSet>> {
     let assets = AdjacencyMap::new()
         .skip_duplicates()
         .visit(once(root), |&asset: &Vc<Box<dyn Module>>| async move {

crates/turbopack-core/src/module.rs

@@ -40,6 +40,14 @@ impl Modules {
 #[turbo_tasks::value(transparent)]
 pub struct ModulesSet(IndexSet<Vc<Box<dyn Module>>>);
 
+#[turbo_tasks::value_impl]
+impl ModulesSet {
+    #[turbo_tasks::function]
+    pub fn empty() -> Vc<Self> {
+        Vc::cell(IndexSet::new())
+    }
+}
+
 /// This is a temporary function that should be removed once the [Module]
 /// trait completely replaces the [Asset] trait.
 /// It converts an [Asset] into a [Module], but either casting it or wrapping it

crates/turbopack-ecmascript/Cargo.toml

@@ -19,7 +19,7 @@ num-bigint = "0.4"
 num-traits = "0.2.15"
 once_cell = { workspace = true }
 parking_lot = { workspace = true }
-petgraph = "0.6.2"
+petgraph = { workspace = true }
 pin-project-lite = { workspace = true }
 regex = { workspace = true }
 rustc-hash = { workspace = true }

crates/turbopack-ecmascript/src/chunk/esm_scope.rs

@@ -0,0 +1,165 @@
+use std::collections::HashMap;
+
+use anyhow::{Context, Result};
+use petgraph::{algo::tarjan_scc, prelude::DiGraphMap};
+use turbo_tasks::{TryFlatJoinIterExt, Value, Vc};
+use turbopack_core::{
+    chunk::{availability_info::AvailabilityInfo, available_assets::chunkable_assets_set},
+    module::{Module, ModulesSet},
+};
+
+use crate::{
+    chunk::EcmascriptChunkPlaceable,
+    references::esm::{base::ReferencedAsset, EsmAssetReference},
+    EcmascriptModuleAssets,
+};
+
+/// A graph representing all ESM imports in a chunk group.
+#[turbo_tasks::value(serialization = "none", cell = "new", eq = "manual")]
+pub(crate) struct EsmScope {
+    scc_map: HashMap<Vc<Box<dyn EcmascriptChunkPlaceable>>, Vc<EsmScopeScc>>,
+    #[turbo_tasks(trace_ignore, debug_ignore)]
+    scc_graph: DiGraphMap<Vc<EsmScopeScc>, ()>,
+}
+
+#[turbo_tasks::value(transparent)]
+pub(crate) struct EsmScopeScc(Vec<Vc<Box<dyn EcmascriptChunkPlaceable>>>);
+
+#[turbo_tasks::value(transparent)]
+pub(crate) struct OptionEsmScopeScc(Option<Vc<EsmScopeScc>>);
+
+#[turbo_tasks::value(transparent)]
+pub(crate) struct EsmScopeSccs(Vec<Vc<EsmScopeScc>>);
+
+#[turbo_tasks::value_impl]
+impl EsmScope {
+    #[turbo_tasks::function]
+    pub(crate) async fn new(availability_info: Value<AvailabilityInfo>) -> Result<Vc<Self>> {
+        let assets = if let Some(root) = availability_info.current_availability_root() {
+            chunkable_assets_set(root)
+        } else {
+            ModulesSet::empty()
+        };
+
+        let esm_assets = get_ecmascript_module_assets(assets);
+        let import_references = collect_import_references(esm_assets).await?;
+
+        let mut graph = DiGraphMap::new();
+
+        for (parent, child) in &*import_references {
+            graph.add_edge(*parent, *child, ());
+        }
+
+        let sccs = tarjan_scc(&graph);
+
+        let mut scc_map = HashMap::new();
+        for scc in sccs {
+            let scc_vc = EsmScopeScc(scc.clone()).cell();
+
+            for placeable in scc {
+                scc_map.insert(placeable, scc_vc);
+            }
+        }
+
+        let mut scc_graph = DiGraphMap::new();
+        for (parent, child, _) in graph.all_edges() {
+            let parent_scc_vc = *scc_map
+                .get(&parent)
+                .context("unexpected missing SCC in map")?;
+            let child_scc_vc = *scc_map
+                .get(&child)
+                .context("unexpected missing SCC in map")?;
+
+            if parent_scc_vc != child_scc_vc {
+                scc_graph.add_edge(parent_scc_vc, child_scc_vc, ());
+            }
+        }
+
+        Ok(Self::cell(EsmScope { scc_map, scc_graph }))
+    }
+
+    #[turbo_tasks::function]
+    pub(crate) async fn get_scc(
+        self: Vc<Self>,
+        placeable: Vc<Box<dyn EcmascriptChunkPlaceable>>,
+    ) -> Result<Vc<OptionEsmScopeScc>> {
+        let this = self.await?;
+
+        Ok(Vc::cell(this.scc_map.get(&placeable).copied()))
+    }
+
+    #[turbo_tasks::function]
+    pub(crate) async fn get_scc_children(
+        self: Vc<Self>,
+        scc: Vc<EsmScopeScc>,
+    ) -> Result<Vc<EsmScopeSccs>> {
+        let this = self.await?;
+
+        let children = this.scc_graph.neighbors(scc).collect();
+
+        Ok(Vc::cell(children))
+    }
+}
+
+#[turbo_tasks::function]
+async fn get_ecmascript_module_assets(
+    modules: Vc<ModulesSet>,
+) -> Result<Vc<EcmascriptModuleAssets>> {
+    let esm_assets = modules
+        .await?
+        .iter()
+        .copied()
+        .map(|r| async move { anyhow::Ok(Vc::try_resolve_downcast_type(r).await?) })
+        .try_flat_join()
+        .await?;
+
+    Ok(Vc::cell(esm_assets))
+}
+
+// for clippy
+type PlaceableVc = Vc<Box<dyn EcmascriptChunkPlaceable>>;
+
+/// A directional reference between 2 [EcmascriptChunkPlaceable]s.
+#[turbo_tasks::value(transparent)]
+struct ImportReferences(Vec<(PlaceableVc, PlaceableVc)>);
+
+#[turbo_tasks::function]
+async fn collect_import_references(
+    esm_assets: Vc<EcmascriptModuleAssets>,
+) -> Result<Vc<ImportReferences>> {
+    let import_references = esm_assets
+        .await?
+        .iter()
+        .copied()
+        .map(|a| async move {
+            let placeable = Vc::upcast::<Box<dyn EcmascriptChunkPlaceable>>(a)
+                .resolve()
+                .await?;
+
+            a.references()
+                .await?
+                .iter()
+                .copied()
+                .map(|r| async move {
+                    let Some(r) = Vc::try_resolve_downcast_type::<EsmAssetReference>(r).await?
+                    else {
+                        return Ok(None);
+                    };
+
+                    let ReferencedAsset::Some(child_placeable) = &*r.get_referenced_asset().await?
+                    else {
+                        return Ok(None);
+                    };
+
+                    let child_placeable = child_placeable.resolve().await?;
+
+                    anyhow::Ok(Some((placeable, child_placeable)))
+                })
+                .try_flat_join()
+                .await
+        })
+        .try_flat_join()
+        .await?;
+
+    Ok(Vc::cell(import_references))
+}

crates/turbopack-ecmascript/src/chunk/mod.rs

@@ -1,6 +1,7 @@
 pub(crate) mod content;
 pub(crate) mod context;
 pub(crate) mod data;
+pub(crate) mod esm_scope;
 pub(crate) mod item;
 pub(crate) mod placeable;
 

crates/turbopack-ecmascript/src/lib.rs

@@ -198,6 +198,10 @@ pub struct EcmascriptModuleAsset {
 #[turbo_tasks::value(transparent)]
 pub struct OptionEcmascriptModuleAsset(Option<Vc<EcmascriptModuleAsset>>);
 
+/// A list of [EcmascriptModuleAsset]s
+#[turbo_tasks::value(transparent)]
+pub struct EcmascriptModuleAssets(Vec<Vc<EcmascriptModuleAsset>>);
+
 impl EcmascriptModuleAsset {
     pub fn builder(
         source: Vc<Box<dyn Source>>,

crates/turbopack-ecmascript/src/references/esm/base.rs

@@ -119,7 +119,7 @@ impl EsmAssetReference {
 #[turbo_tasks::value_impl]
 impl EsmAssetReference {
     #[turbo_tasks::function]
-    pub(super) async fn get_referenced_asset(self: Vc<Self>) -> Result<Vc<ReferencedAsset>> {
+    pub(crate) async fn get_referenced_asset(self: Vc<Self>) -> Result<Vc<ReferencedAsset>> {
         let this = self.await?;
 
         Ok(ReferencedAsset::from_resolve_result(